Repository navigation
refactor: unify AWEX converters and adapters across v1 and v2 - #1684
Conversation
There was a problem hiding this comment.
AwexFSDPAdapter now imports and calls batch_send_recv_by_peer, so this mock targets a symbol that no longer exists in the FSDP module. The FSDP test case raises AttributeError before reaching the barrier assertion. Please patch batch_send_recv_by_peer instead, and verify that the transfer is invoked before the Gloo completion barrier.
There was a problem hiding this comment.
Fixed in a17882b. The parameterized FSDP/Megatron/SGLang test now patches batch_send_recv_by_peer, checks its arguments for both use_group settings, and asserts that the transfer returns before the Gloo completion barrier. The SGLang copy test now asserts transfer -> copy -> synchronize -> barrier. I also updated the two DTE module-loader stubs to expose the new helper.
Validation: full pre-commit, Ruff, and Python compilation passed. Unit tests were not run locally.
b377204 to
4d69307
Compare
GitHub Actions
Failure
GitHub Actions
Failure
sitabulaixizawaluduo
added
safe-to-test
GitHub Actions
Inactive
GitHub Actions
Inactive
| model_config, _ = PretrainedConfig.get_config_dict(actor.path) | ||
| if model_config.get("model_type") == "qwen2" or "Qwen2ForCausalLM" in ( | ||
| model_config.get("architectures") or [] | ||
| ): | ||
| raise ValueError( | ||
| "Qwen2/Qwen2.5 text models do not support AWEX weight updates. " | ||
| "For RL, set actor._version=v1, rollout._version=v1, and " | ||
| "actor.weight_update_mode=xccl or disk. Controller v2 is not " | ||
| "supported for these models' RL weight updates." | ||
| ) |
There was a problem hiding this comment.
This guard also rejects Qwen2/Qwen2.5 with v2 + use_lora=True, but the existing trainer routes that configuration through disk-based adapter updates, bypassing AWEX conversion. Is there a separate incompatibility that requires rejecting this path? If not, could we scope the guard to configurations that actually select AWEX?
| for peer in sorted(ops_by_peer): | ||
| peer_ops = ops_by_peer[peer] | ||
| batch_send_recv( | ||
| send_ops=peer_ops if send_ops else [], | ||
| recv_ops=peer_ops if recv_ops else [], | ||
| blocking=True, | ||
| use_group=True, | ||
| ) |
There was a problem hiding this comment.
Is peer serialization required for the shared adapter/native converter integration to work? It appears independent of the refactor and changes the default transfer scheduling. Unless there is a demonstrated correctness or resource issue that makes it necessary, could we preserve the existing scheduling here and move this change to a separate PR with before/after measurements?
Share Megatron and SGLang weight-update adapters so both versions use native AWEX conversion and sharding metadata. Key changes: - Consolidate v1 colocate and v2 separation adapter implementations - Preserve legacy entry points and shared Megatron weight residency - Forward inference converter context through the weight-update gateway - Propagate parameter layouts between training and inference workers - Retain full-weight and delta-weight separation transfer paths
Defer peer-serialized transport to a separate PR while keeping the shared adapter and native converter changes. Key changes: - Restore blocking AWEX batch_send_recv on training and inference. - Update transfer mocks while retaining completion-order assertions.
Disk updates bypass AWEX native conversion, so the model support guard must allow them for both controller versions.
7f09d5a to
599718c
Compare
|
This pull request has been automatically marked as stale because it has not had recent activity within the last 14 days. Please add a comment or push new commits to keep it active. Thank you for your contribution! |
Carry the new CPU-staged optimizer residency and IPC export lifetime through the shared AWEX adapter while preserving the PR module layout.
sitabulaixizawaluduo
added
safe-to-test
Controller v2 selects disk updates for LoRA and AWEX otherwise, regardless of weight_update_mode. Match the guard to that dispatch so Qwen2 LoRA runs are allowed and non-LoRA AWEX runs remain rejected.
The Megatron AWEX fallback lacks Qwen2 conversion mappings, while FSDP publishes HF-layout weights without that converter. Avoid rejecting the FSDP path before its own metadata validation.
2 participants
Description
Unify the Megatron and SGLang AWEX integrations so v1 colocated weight updates and v2 separated weight updates share adapter implementations and native AWEX model conversion.
Shared adapters
areal/engine/awex/.areal.v2in the shared AWEX package.MegatronWeightResidencyfor weight, optimizer, and gradient memory management.Native conversion and metadata
Type of Change