feat(awex): add separation AdamW delta weight transfer

[pyq623]

Summary

This PR adds an opt-in, AReaL-DTE-backed AdamW delta weight transfer path for AWEX separation mode.

The initial synchronization and periodic anchors use the existing full-weight path. Safe, contiguous intermediate versions use sparse deltas derived from distributed AdamW optimizer state.

The feature is disabled by default and does not change existing full-weight behavior.

PR migration

This PR supersedes #1604.

The implementation has been rebased onto the latest upstream main and moved from a fork branch to an areal-project/AReaL repository branch so the complete protected CI workflow can run.

Previous review context and approval:

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

The cumulative feature scope remains the same as #1604, with additional protocol-safety fixes described below.

Dependency

The required AReaL-DTE support has been merged:

• feat(core): support external sparse delta construction and commit tracking AReaL-DTE#1

It provides external sparse-delta construction, index remapping, and durable DeltaTracker commit tracking used by this integration.

Supported configuration

The supported topology is intentionally restricted to the path validated by this PR:

actor:
  backend: megatron
  _version: v2
  weight_update_mode: awex
  weight_update_transfer: delta
  weight_update_delta_method: adamw
  weight_update_anchor_interval: 20
  use_lora: false
  ppo_n_minibatches: 1

The combined inference and training world size must be a positive power of two, as required by the recursive sparse transport scheduler in AWEX 0.8.0.

Unsupported backends, versions, topologies, LoRA mode, minibatch counts, and world sizes fail fast before worker/process-group initialization.

Changes

AdamW change detection

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

• current parameters;
• exp_avg and exp_avg_sq;
• optimizer step;
• AdamW hyperparameters and the effective step learning rate.

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

Missing or ambiguous optimizer state, unsupported precision-aware state, skipped or discontinuous versions, and other unsafe conditions fall back to a full synchronization.

Separation transfer lifecycle

The Megatron and SGLang adapters now support:

• initial and periodic full-weight anchors;
• rank-consistent full/delta decisions;
• sparse payload construction from the AWEX transfer plan;
• globally canonical mixed-dtype transfer rounds;
• empty local participation in dtype rounds when required for collective symmetry;
• in-place application to live inference parameters;
• CUDA completion followed by a Gloo completion barrier;
• tracker and detector advancement only after receiver apply succeeds;
• temporary Megatron grad-buffer release with guaranteed restoration.

A failed transfer does not advance the delta version, anchor counter, or detector watermark.

Protocol safety

This version additionally:

• validates the exact supported Megatron-v2/SGLang-v2 AWEX topology before exporting DTE environment variables;
• rejects non-power-of-two combined world sizes before process-group creation;
• builds and caches a rank-consistent canonical dtype union;
• requires every rank to enter every dtype round, including ranks with an empty local plan;
• preserves dense AWEX behavior when DTE is disabled.

Configuration and documentation

This PR includes:

• delta-transfer configuration integrated into TrainEngineConfig;
• worker environment propagation;
• generated English and Chinese CLI reference updates;
• a runnable Qwen3-30B-A3B GSM8K separation example;
• validation and documentation for ppo_n_minibatches=1.

Scope

This PR intentionally covers only separation-mode AdamW delta transfer.

It does not add or modify:

• colocate/shared-GPU delta transfer;
• snapshot or dirty-bit detectors;
• arbitrary-world-size AWEX recursive scheduling;
• post-apply weight fingerprints;
• deterministic rollout scheduling or request tracing;
• checkpoint/recovery behavior;
• experiment runners, Slurm configuration, or W&B settings;
• FP32 LM-head or logprob numerical alignment;
• AWEX native model converters.

Zero-gradient steps continue through normal AdamW inversion and mask calculation because momentum and weight decay may still modify model-visible weights.

Validation

The final branch was rebased onto:

areal-project/main@3a0b37222e10100755581cfa3d123f0f2009697f

Final post-rebase validation:

Focused DTE/AWEX, topology, NCCL-group, CLI and GAE tests:
  103 passed, 1 warning

Qwen3-30B flat TrainEngineConfig parse and DTE environment propagation:
  passed

Scoped pre-commit checks:
  passed

Ruff lint and formatting:
  passed

py_compile:
  passed

git diff --check:
  passed

English/Chinese CLI documentation regeneration:
  no diff

The implementation lineage was exercised in a completed 100-step Qwen3-30B-A3B separation AdamW delta run.

[sitabulaixizawaluduo]

Can the option be described using enable_delta_weight_update(bool), as the current name does not sound understandable of its function?

[pyq623]

Good suggestion. I replaced weight_update_transfer: full|delta with enable_delta_weight_update: bool = False in TrainEngineConfig, and updated the validation, examples, generated CLI docs, and regression tests accordingly. The runtime delta-transfer protocol is unchanged. Addressed in 56145b7.

[sitabulaixizawaluduo]

Can an annotation be added to describe the functions of these environment variables?

[pyq623]

Added inline comments describing the roles of all four environment variables: separation AWEX path selection, sparse delta transfer, periodic full-sync anchors, and streaming pre-step reconstruction. Their names, values, and runtime behavior are unchanged. Addressed in 56145b7.