feat: implement megatron weight update with awex#1239
Conversation
ran-yan-hk
requested review from
garrett4wade,
nuzant and
rchardx
as code owners
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the AwexMegatronAdapter to support weight updates for Megatron-based training engines, enabling integration with the Awex weight update system for both Data Parallelism (DP) and Tensor Parallelism (TP) configurations. The changes include the adapter implementation, factory integration, and new end-to-end tests. Review feedback highlights critical memory optimization opportunities, specifically recommending that parameter gathering and iteration be filtered by required names to prevent Out-Of-Memory (OOM) errors on large models. Additionally, it was noted that the adapter's docstring should be updated to correctly reflect its support for TP.
ran-yan-hk
changed the title
ran-yan-hk
force-pushed
the
yr/megatron-awex
branch
from
4308d0f to
8254e59
Compare
| gathered = all_gather_param( | ||
| mcore_name, | ||
| param, | ||
| fp8_direct_convert=False, | ||
| quantization_config=None, | ||
| duplicated_param_names=self._engine._duplicated_param_names, | ||
| ) |
There was a problem hiding this comment.
Just FYI this all-gather should be removed in the future.
There was a problem hiding this comment.
Sure. This experimental version is runnable but performance issues will be fixed soon.
ran-yan-hk
force-pushed
the
yr/megatron-awex
branch
from
8254e59 to
71346fc
Compare
ran-yan-hk
force-pushed
the
yr/megatron-awex
branch
from
71346fc to
bfbacec
Compare
2 participants
Description
Following PR#1214, add support for Megatron engine for weight update with awex.
This implements a full connect → update_weights → disconnect lifecycle orchestrated
by an HTTP gateway. Megatron and SGLang adapters handle parameter shard mapping. For pure DP, Megatron is integrated in the same way as FSDP; however, Megatron TP's performance may not full optimized (but runnable) due to its built-in all-gather operations.
Current scope: Megatron engine (pure DP or pure TP ore hybrid DP and TP) + SGLang (pure DP) only. PP/EP sharding
is not yet supported.
Related Issue
NA
Type of Change
Checklist
pre-commit run --all-files)./docs/build_all.sh)main/review-prcommand/create-prBreaking Change Details (if applicable):
NA
Additional Context
Key files:
areal/experimental/training_service/worker/awex.py: Implement auto detection on training backend.areal/experimental/weight_update/awex/megatron_adapter.py: Implement Megatron adapter to be compatible with awex.tests/experimental/weight_update/test_nccl_integration.py: Integrate Megatron engine test for pure DP, pure TP, and hybrid DP + TP.