feat(awex): add separation AdamW delta weight transfer - #1604
Open
pyq623 wants to merge 5 commits into
Open
Conversation
pyq623
requested review from
TaoZex,
fishcrap,
garrett4wade,
nuzant,
rchardx and
sitabulaixizawaluduo
as code owners
August 14, 2026 05:38
pyq623
requested review from
CormickKneey,
HwVanICI and
PrometheusComing
as code owners
August 14, 2026 06:31
| beta1 = 0.0 | ||
| beta2 = 0.0 | ||
| return ( | ||
| float(param_group.get("_areal_last_step_lr", param_group["lr"])), |
Collaborator
There was a problem hiding this comment.
The _areal_last_step_lr is not seen where it is passed in, and it always uses param_group["lr"]. However, the update of param_group["lr"] needs to be before the current calculation, and is there a possibility of a silent calculation error here?
| force_dense = True | ||
| else: | ||
| skipped_step_unchanged += 1 | ||
| elif step_delta != 1: |
Collaborator
There was a problem hiding this comment.
I understand that if ppo_n_minibatches is greater than 1, step_delta here is also greater than 1, which forces the use of full sync. However, this information is not described in the documentation. If it is the case, it is recommended to add validation and warning logs in cli_args.py, and describe it in examples/dte/README.md. If not, please correct me.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR adds an opt-in, DTE-backed AdamW delta weight transfer path for AReaL separation mode.
The first synchronization and periodic anchors use the existing full-weight AWEX path. Safe, contiguous intermediate versions use sparse AdamW-
derived deltas.
The feature is disabled by default and does not change existing full-weight behavior.
Dependency
This is a stacked PR and depends on:
AReaL-DTE PR #1 adds the external integer-index payload construction and
DeltaTracker.mark_delta_committed()lifecycle required by thisintegration.
AWEX v0.8.0 and the Gloo weight-update sidecar are already available on the target AReaL
mainbranch and are not duplicated here.Configuration
The supported configuration is intentionally limited to the validated separation AdamW path:
When
actor.dte.enabled=false, the existing AWEX full-weight path is unchanged.Unsupported topology, transfer mode, or delta method combinations are rejected explicitly.
Changes
Configuration and worker environment
This PR adds:
actor.dteconfiguration;AdamW change detection
The training side reconstructs pre-step parameters from distributed AdamW state using:
exp_avg;exp_avg_sq;The reconstructed tensors pass through the same AReaL Megatron-to-HF conversion path used by full-weight synchronization.
The detector then produces compact
int32/int64flat change indices. Missing or unsafe optimizer state, version discontinuities, unsupported precision-aware optimizer state, and other ambiguous conditions fall back to full synchronization.Streaming reconstruction avoids retaining a complete second HF payload while constructing masks.
Separation transfer lifecycle
The Megatron and SGLang adapters now support:
A failed transfer does not advance the delta version, anchor counter, or detector watermark.
Scope
This PR intentionally includes only the separation AdamW delta path.
It does not add or modify:
The previously extracted, unconnected zero-delta probe fast path was also removed from this PR. Zero-gradient steps continue through the normal AdamW inversion and mask calculation because momentum and weight decay may still change model-visible weights.
Tests
Added focused coverage for:
Validation against AReaL-DTE PR #1 and AWEX v0.8.0:
Additional checks:
Runtime validation boundary
The implementation was extracted from the Qwen3-30B-A3B separation integration used for a completed 100-step AdamW delta run.
However, the exact combination of this public AReaL PR and public AReaL-DTE PR #1 commits has not yet been rerun through a fresh cluster smoke or 100-step trial. The focused tests above validate the API and lifecycle integration but should not be interpreted as a completed cluster validation of the public PR pair.