Skip to content

feat(awex): add separation AdamW delta weight transfer - #1604

Open
pyq623 wants to merge 5 commits into
areal-project:mainfrom
pyq623:feat/dte-separation-adamw
Open

feat(awex): add separation AdamW delta weight transfer#1604
pyq623 wants to merge 5 commits into
areal-project:mainfrom
pyq623:feat/dte-separation-adamw

Conversation

@pyq623

@pyq623 pyq623 commented Aug 14, 2026

Copy link
Copy Markdown
Collaborator

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 this
integration.

AWEX v0.8.0 and the Gloo weight-update sidecar are already available on the target AReaL main branch and are not duplicated here.

Configuration

The supported configuration is intentionally limited to the validated separation AdamW path:

actor:
  dte:
    enabled: true
    transfer: delta
    delta_method: adamw
    anchor_interval: 20

rollout:
  scheduling_strategy:
    type: separation

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:

  • an explicit actor.dte configuration;
  • separation-only topology validation;
  • propagation of the delta, anchor, and streaming-reconstruction settings to actor and rollout workers;
  • generated English and Chinese CLI reference updates.
  • Added a runnable Qwen3-30B-A3B GSM8K separation example.
  • Example configuration and DTE environment propagation were validated in the Python 3.12 training image.

AdamW change detection

The training side reconstructs pre-step parameters from distributed AdamW state using:

  • current parameters;
  • exp_avg;
  • exp_avg_sq;
  • optimizer step and AdamW hyperparameters.

The reconstructed tensors pass through the same AReaL Megatron-to-HF conversion path used by full-weight synchronization.

The detector then produces compact int32/int64 flat 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:

  • initial and periodic full anchors;
  • rank-consistent full/delta decisions;
  • sparse payload construction from the AWEX transfer plan;
  • dtype-grouped, two-round sparse P2P through AReaL-DTE;
  • in-place application to live inference parameters;
  • temporary Megatron grad-buffer release and guaranteed restoration;
  • CUDA completion followed by a Gloo completion barrier;
  • tracker and detector state advancement only after receiver apply succeeds.

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:

  • colocate/shared-GPU delta transfer;
  • snapshot or dirty-bit detectors;
  • post-apply digest or weight fingerprint diagnostics;
  • deterministic rollout scheduling or request tracing;
  • recover/checkpoint or stale-process cleanup;
  • experiment runners, Slurm configuration, or W&B settings;
  • FP32 LM-head or logprob numerical alignment;
  • AWEX native model converters;
  • the AReaL-DTE remap or sparse P2P protocol implementation.

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:

  • default-disabled behavior;
  • supported and rejected topology/method combinations;
  • DTE environment propagation;
  • AdamW inversion round trips;
  • DTE tracker construction;
  • full/delta control-group initialization;
  • rank-consistent fallback;
  • receiver completion barriers;
  • tracker commit ordering;
  • failure paths not advancing tracker state;
  • tensor-parallel metadata preservation;
  • streaming reconstruction collective ordering.

Validation against AReaL-DTE PR #1 and AWEX v0.8.0:

tests/test_dte_topology_gating.py
tests/test_awex_delta_common.py
tests/test_awex_separation_delta.py

19 passed

Additional checks:

scoped pre-commit hooks: passed
Ruff lint/format: passed
git diff --check: passed
target Python py_compile: passed
English/Chinese CLI reference regeneration: no diff

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.

beta1 = 0.0
beta2 = 0.0
return (
float(param_group.get("_areal_last_step_lr", param_group["lr"])),

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@pyq623
pyq623 requested a review from geshi001 as a code owner August 14, 2026 08:49

@sitabulaixizawaluduo sitabulaixizawaluduo left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

sitabulaixizawaluduo

This comment was marked as outdated.

@sitabulaixizawaluduo sitabulaixizawaluduo added the safe-to-test Ready to run unit-tests in a PR. label Aug 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

safe-to-test Ready to run unit-tests in a PR.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants