Skip to content

feat: add TrainerRank#731

Closed
bradhilton wants to merge 115 commits into
mainfrom
feat/trainer-rank-gdn-tree
Closed

feat: add TrainerRank#731
bradhilton wants to merge 115 commits into
mainfrom
feat/trainer-rank-gdn-tree

Conversation

@bradhilton

@bradhilton bradhilton commented Jun 20, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • add art.megatron.trainer_rank plus a minimal dev/trainer_rank.py torchrun demo
  • add shared-prefix packing/tree helpers and unify GDN execution around the generic tree path
  • add TrainerRank request-head support for target logprobs, multi-target labels, top-k, logits, and hidden states
  • add topology/perf/parity dev harnesses and unit/integration coverage

Validation

  • uv run ruff check src/art/megatron/context_parallel/builder.py src/art/megatron/shared_prefix_state.py tests/unit/test_shared_prefix_attention_builder.py dev/trainer_rank_perf.py
  • uv run pytest tests/unit/test_shared_prefix_packing.py tests/unit/test_shared_prefix_tree.py tests/unit/test_shared_prefix_grad_parity.py tests/unit/test_trainer_rank_validation.py tests/unit/test_shared_prefix_attention_builder.py (34 passed, 8 skipped locally; Megatron-only attention builder passes on H200)
  • H200: shared-prefix attention builder 7 passed
  • H200: GDN CP packed correctness 7 passed, 2 skipped
  • H200: real GDN/native FLA CP 4 passed, 2 skipped
  • H200: Qwen35 full-model CP1 packed vs flattened 1 passed
  • H200: TrainerRank topology matrix 120/120 passed across DP/TP/CP <= 4 and depths 0..3
  • H200 35B/A3B CP=4 EP=4 perf guards: Austin 198k, depth-3 random, no-sharing 90k, mixed hidden/logits/top-k outputs

@bradhilton bradhilton changed the title feat: add TrainerRank and generic tree GDN feat: add TrainerRank Jun 22, 2026
@bradhilton bradhilton force-pushed the feat/trainer-rank-gdn-tree branch from 2cecc5a to 94afb0f Compare June 22, 2026 19:01
bradhilton added 20 commits June 22, 2026 16:13
@bradhilton
fix: guard dynamic slot optimizer reductions
327ced4
@bradhilton
Merge remote-tracking branch 'origin/main' into feat/trainer-rank-gdn…
02e7103 
…-tree

# Conflicts:
#	scripts/build-gpu-image.sh
@bradhilton
test: add megatron review readiness harness
9dcc9f9
@bradhilton
test: add Austin-focused TrainerRank validation
d78eab7
@bradhilton
test: harden Austin validation harness
991aaff
@bradhilton
fix: prune false exact CP mask blocks
896d344
@bradhilton
fix: refine full exact CP mask blocks
fb8b188
@bradhilton
perf: avoid unnecessary exact CP mask refinement
0c8a16b
@bradhilton
fix: keep exact CP refinement for deep trees
1ac8edf
@bradhilton
test: make CP mask parity slice aware
19cb41a
@bradhilton
fix: broadcast slice-aware mask parity
f464939
@bradhilton
perf: vectorize depth-one CP mask refinement
f640167
@bradhilton
perf: keep depth-one CP mask refinement vectorized
3cf1fe5
@bradhilton
fix: handle empty CP correctness shards
63ca045
@bradhilton
fix: collect perf metadata on all ranks
5fb5924
@bradhilton
bench: use CP stage masks for flex timing
f79f520
@bradhilton
fix: benchmark production sparse flex path
2526641
@bradhilton
fix: mirror production flex stage padding in review bench
1e95437
@bradhilton
fix: default flex review bench to qwen head dim
3578fa4
@bradhilton
perf: skip exact refinement for homogeneous mask blocks
76ded79
bradhilton added 18 commits June 25, 2026 05:48
@bradhilton
refactor: unify trainer rank head projection
9b5c5da
@bradhilton
refactor: simplify trainer rank planning
e282ddd
@bradhilton
refactor: simplify block mask validation
58b5e5d
@bradhilton
refactor: remove alternate cp planner strategies
a67fd86
@bradhilton
refactor: trim trainer rank gdn surfaces
15c01b6
@bradhilton
refactor: shrink shared prefix parser metadata
ca369df
@bradhilton
refactor: use shared prefix pack in trainer rank
3429e47
@bradhilton
refactor: trim trainer rank packing helpers
5ccb6f9
@bradhilton
refactor: share trainer rank grad sync helpers
3893e9b
@bradhilton
refactor: unify gdn tree bucket planning
b42b76f
@bradhilton
refactor: trim adaptive planner cache keys
3519caf
@bradhilton
refactor: inline stale trainer rank helpers
3b14f35
@bradhilton
test: allow tensor fields in gdn layout fixture
bbdd9d5
@bradhilton
fix: type triton topk launch constants
d1a5283
@bradhilton
test: handle fused expert fc1 lora in oracle
865e945
@bradhilton
test: canonicalize padded forward trace rows
8a827d1
@bradhilton
fix: run cp tree child buckets recurrently
0a4bb28
@bradhilton
Revert "fix: run cp tree child buckets recurrently"
16e51b8 
This reverts commit 0a4bb28.
@bradhilton bradhilton force-pushed the feat/trainer-rank-gdn-tree branch 5 times, most recently from ed32a0c to 6443d99 Compare June 25, 2026 16:52
@bradhilton bradhilton force-pushed the feat/trainer-rank-gdn-tree branch from 6443d99 to fa09126 Compare June 25, 2026 16:58
@bradhilton

bradhilton commented Jun 25, 2026

Copy link
Copy Markdown
Collaborator Author

Superseded by the split draft PRs: #739 for the Austin-facing Megatron/CP/GDN core changes, and #740 for the dependent art.trainer_rank API layer.

@bradhilton bradhilton closed this Jun 25, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@bradhilton