Skip to content

feat: add v2 DTE colocated weight transfer - #1587

Open
dingzhiqiang wants to merge 1 commit into
mainfrom
chucai.dzq/feat-dte-v2-clean
Open

feat: add v2 DTE colocated weight transfer#1587
dingzhiqiang wants to merge 1 commit into
mainfrom
chucai.dzq/feat-dte-v2-clean

Conversation

@dingzhiqiang

Copy link
Copy Markdown
Collaborator

Summary

  • add v2 AWEX colocated full and DTE delta weight transfer
  • support snapshot and AdamW inversion delta detection
  • preserve the receiver delta base across rollout memory transitions
  • add DTE configuration, dependency locking, compact examples, and focused tests

Validation

  • 95 passed, 1 skipped
  • Ruff lint and format checks passed
  • pre-commit passed except uv-lock, which could not download a PyTorch aarch64 wheel due to a network timeout; committed lockfiles already include the pinned AReaL-DTE revision

Risk Areas

  • Megatron AdamW state inversion and MoE DP-group mapping
  • SGLang tagged weight/KV-cache offload lifecycle
  • colocated metadata merge and rank/device mapping

Integrate full and delta weight updates into the v2 AWEX colocated path while preserving the receiver base across rollout offload.

Key changes:
- add snapshot and AdamW delta detection with DTE encoding
- propagate actor DTE configuration to train and rollout workers
- add tagged memory handling and SGLang weight CPU backup support
- add focused metadata, optimizer, configuration, and round-trip tests

def __init__(self, engine: MegatronEngine):
self._engine = engine
setattr(engine, "_awex_weight_update_adapter", self)

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 method set here is not seen to be read or used.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants