Skip to content

[codex] Add DMuon param policy and isolated process groups - #12

Merged
StarrickLiu merged 1 commit into
mainfrom
feat/lpt-ablation-wallx-eval
Jun 24, 2026
Merged

StarrickLiu merged 1 commit into
mainfrom
feat/lpt-ablation-wallx-eval

Conversation

@StarrickLiu

Copy link
Copy Markdown
Collaborator

Summary

This PR adds a structured DMuon parameter policy and isolated process group support for Pure DMuon training.

The main changes are:

  • Add param_policy support to configure per-parameter optimizer route and dtype policy.
  • Support mixed precision routing such as fp32 action-head compute with bf16 backbone compute.
  • Add process_group_policy="isolated" so DMuon can own dedicated communication process groups instead of reusing trainer process groups.
  • Keep DMUON_ISOLATED_PG_BARRIER disabled by default to preserve cross-step overlap.
  • Update checkpointing, diagnostics, docs, and tests to cover the new policy semantics.

Motivation

Pure DMuon integrations need more precise control than the previous global compute_dtype and route hint API could express. In WallX, action-head
modules may require fp32 forward compute, while the rest of the model should continue using bf16. The previous API could not represent this cleanly.

Separately, DMuon communication should not depend on trainer-owned process groups. Isolating DMuon process groups makes integration safer across
different training frameworks and avoids subtle process-group ownership or ordering issues.

Implementation Notes

  • Introduced dmuon.policy with structured DMuonParamPolicy.
  • Extended dedicate_params() to accept param_policy and param_policy_fn.
  • Preserved legacy route_hint_fn behavior, but disallow mixing it with structured policy.
  • Added isolated DP/HSDP/TP process group creation.
  • Added optional diagnostic barrier via DMUON_ISOLATED_PG_BARRIER=1; default remains off.
  • Updated checkpoint logic to save master parameter storage where appropriate.
  • Extended diagnostics to report route and dtype policy decisions.

Validation

  • python -m compileall -q dmuon tests
  • python -m pytest -q tests/unit/test_param_policy.py tests/unit/test_param_policy_state_cast.py tests/unit/test_diagnostics.py
  • torchrun --standalone --nproc_per_node=2 tests/distributed/test_process_group_policy.py
  • Validated with a 12-node WallX Pure DMuon training job that successfully entered training steps.

@StarrickLiu
StarrickLiu merged commit d1fa2ee into main Jun 24, 2026
2 checks passed
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.

1 participant