Skip to content

fix: complete TrainerRank model runtime integration#763

Merged
bradhilton merged 10 commits into
mainfrom
fix/trainer-rank-pr758-followup
Jul 17, 2026
Merged

fix: complete TrainerRank model runtime integration#763
bradhilton merged 10 commits into
mainfrom
fix/trainer-rank-pr758-followup

Conversation

@bradhilton

@bradhilton bradhilton commented Jul 16, 2026

Copy link
Copy Markdown
Collaborator

Summary

Completes the TrainerRank integration with the model-runtime behavior added by #758.

  • Canonicalize model-specific internal LoRA padding when loading dynamic slots and zero padded slot gradients before optimizer updates, without mutating unrelated slots that may retain backward graphs.
  • Pass packed positions, sliding-window variants, model state, and GDN planner configuration through the CP=1 prefix-tree path.
  • Configure HybridEP capacity and the exact per-rank model-row extent for CP=1 and CP>1 packed forwards.
  • Reject simultaneous TrainerRank data parallelism and expert parallelism before HybridEP collectives. The supported/recommended MoE topology is DP=1 with CP=EP=world_size.

Validation

  • uv run --no-sync prek run --all-files
  • uv run --no-sync python dev/trainer_rank_fast_check.py -q (75 passed, 4 skipped)
  • 2xH200 dependency gate: all 8 Megatron integration tests, including packed attention, depth>1 GDN parity/trainability, TP LoRA gradients, dynamic-slot recompute/stepping, and TP head backward
  • 2xH200 CP=2 matrix: all 16 output combinations at depths 0 and 4 with chunk sizes 17 and 8192; mean_abs_pct=0.00155049, head-backward mean_abs_pct=0.0
  • 4xH200 unsupported-topology regression: the formerly hanging CP=1/EP=4/DP=4 unequal-tail case raises the documented error on all four ranks before model collectives
  • 4xH200 supported-topology regression: Qwen/Qwen3.5-35B-A3B, CP=4/EP=4/DP=1, unequal dynamic slots, retained graphs, backward, and optimizer step
  • GPT-OSS CP=1 dynamic-slot forward/backward/optimizer smoke
  • Open-ended adversarial reviews from Claude and a separate Codex subagent found no remaining blocker

No files under src/art/megatron/** are changed.

Base automatically changed from austin/e2e_pipeline_rl_on_pr_739_unify_shared_prefix to main July 16, 2026 20:07
@bradhilton
bradhilton force-pushed the fix/trainer-rank-pr758-followup branch from c3be588 to 6e4254d Compare July 16, 2026 20:35
@bradhilton
bradhilton force-pushed the fix/trainer-rank-pr758-followup branch from edc2a3c to 9bfe4ef Compare July 16, 2026 22:14
@bradhilton
bradhilton marked this pull request as ready for review July 16, 2026 23:24
@bradhilton
bradhilton merged commit 44081f9 into main Jul 17, 2026
11 checks passed
@bradhilton
bradhilton deleted the fix/trainer-rank-pr758-followup branch July 17, 2026 00:42
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