Repository navigation

feat(draft): support packed context-parallel co-training - #3757

Open

seonjinn wants to merge 398 commits into

NVIDIA-NeMo:mainfrom

seonjinn:seonjinn/draft-cotrain-cp-v1-linear

Open

feat(draft): support packed context-parallel co-training#3757
seonjinn wants to merge 398 commits into
NVIDIA-NeMo:mainfrom
seonjinn:seonjinn/draft-cotrain-cp-v1-linear

This pull request is big! We're only showing the most recent 250 commits

Commits on Aug 21, 2026
- fix(draft): synchronize TP input shapes
  seonjinn
  committed
  0f097a2
- Merge NVIDIA-NeMo/RL main into PR9
  seonjinn
  committed
  4ce4c30
- Merge NVIDIA-NeMo/RL main into PR9
  seonjinn
  committed
  3dfef44
- style(vllm): normalize speculator runtime imports
  seonjinn
  committed
  9721015
- fix(vllm): fail closed on FP8 KV postprocess errors
  seonjinn
  committed
  ef32181
- feat(refit): represent train context parallelism
  seonjinn
  committed
  3921a78
- fix(refit): validate CP refit mesh setup
  seonjinn
  committed
  b05e245
- fix(refit): validate PP before communicator setup
  seonjinn
  committed
  52e50b4
- feat(draft): bucket refit payloads by CP lane
  seonjinn
  committed
  c8bbc46
- fix(draft): synchronize CP lane refit preflight
  seonjinn
  committed
  4539a29
- feat(draft): bind refit owners to MCore CP lanes
  seonjinn
  committed
  39ee71e
- fix(draft): forward metadata through quant refit
  seonjinn
  committed
  692d3ea
- fix(draft): satisfy refit metadata typing
  seonjinn
  committed
  bc3d4e5
- style(refit): format CP validation test
  seonjinn
  committed
  fb7a1c7
- test(draft): cover real-quant metadata transport
  seonjinn
  committed
  c467540
- test(refit): cover cross-group draft failure
  seonjinn
  committed
  74ab3d2
- fix(refit): fail together before draft transport
  seonjinn
  committed
  6de4b13
- fix(refit): preflight draft metadata before meta context
  seonjinn
  committed
  6bbf215
- fix(refit): narrow synchronized error chaining
  seonjinn
  committed
  3e36e97
- feat(draft): support packed context-parallel co-training
  seonjinn
  committed
  56494c3
- fix(dspark): require target-owned vocabulary
  seonjinn
  committed
  7d73fce
- test(draft): compose TP2 with packed CP groups
  seonjinn
  committed
  c26ff0a
- chore(draft): include stable IDs in type checks
  seonjinn
  committed
  438134e
- test(dspark): patch isolated provider fixture
  seonjinn
  committed
  1ae1a9c
- chore(draft): add missing copyright headers
  seonjinn
  committed
  0198f9e
- docs(draft): document optimizer update probes
  seonjinn
  committed
  a393e03
- test(draft): cover target TP and optional draft config
  seonjinn
  committed
  915a2db
- fix(draft): enforce supported topology and preserve no-draft training
  seonjinn
  committed
  7f51cb4
- docs(draft): add block-draft guide and nightly
  seonjinn
  committed
  d0f67cc
- test(draft): require CP-invariant split gradients
  seonjinn
  committed
  1888640
- test(dspark): require revision-aware checkpoint loading
  seonjinn
  committed
  30b76f2
- test(draft): require CP-invariant leader metrics
  seonjinn
  committed
  ed3dac7
- fix(draft): normalize CP gradients and metrics
  seonjinn
  committed
  e166171
- fix(dspark): separate checkpoint revision from model ID
  seonjinn
  committed
  4e17585
- fix(dflash): preserve CP autograd for empty queries
  seonjinn
  committed
  0195822
- test(dspark): allowlist public revision literal
  seonjinn
  committed
  0dc5f9a
- fix(dflash): accept inherited auxiliary layer field
  seonjinn
  committed
  55e677b
- test(dspark): add packed CP nightly coverage
  seonjinn
  committed
  4be3c96
- fix(draft): allow capable providers with packed policies
  seonjinn
  committed
  27cdb2d
- test(draft): require post-update refit coverage
  seonjinn
  committed
  fbe73bf
- test(draft): mark completed post-update refits
  seonjinn
  committed
  24956bc
- test(draft): expose resumed refit marker ordering
  seonjinn
  committed
  d832b1a
- fix(draft): track pending post-update refits
  seonjinn
  committed
  44b13d3
- test(dflash): align CP attention oracles
  seonjinn
  committed
  218ec64
- fix(dflash): accept singleton CP groups
  seonjinn
  committed
  3856435
- test(dflash): require stable CUDA graph recipe
  seonjinn
  committed
  8a90d63
- fix(dflash): use stable CUDA graph compilation
  seonjinn
  committed
  a39b04f
- ci(dspark): allowlist official artifact hashes
  seonjinn
  committed
  b4e6c7b
- fix(dflash): pass unpacked benchmark layout
  seonjinn
  committed
  8871c2f
- test(dflash): populate expanded attention plans
  seonjinn
  committed
  2b6295f
- test(dflash): pass unpacked block mask layout
  seonjinn
  committed
  ba7ecbc
- test(draft): assert CP denominator group by call intent
  seonjinn
  committed
  2471561
- test(draft): compare CUDA metrics as scalars
  seonjinn
  committed
  ee222de
Commits on Aug 22, 2026
- Fix TP checkpoint group handling and shard tagging in DSpark heads
  seonjinn
  committed
  54a3d9b
- fix(test): forward GPU topology to Ray launcher
  seonjinn
  committed
  829a1d6
- Remove redundant TP metadata collective and document DFlash loss contracts
  seonjinn
  committed
  7700a07
- fix(draft): route packed CP training through split API
  seonjinn
  committed
  a5ef756
- fix(draft): align split routing with provider capabilities
  seonjinn
  committed
  7b80930
- Document artifact pin and TP ceiling, test error branches, rename PP-named tests
  seonjinn
  committed
  abc74d7
- fix(dflash): reconstruct target-SP mask embeddings
  seonjinn
  committed
  446ea90
- fix(draft): isolate transient outputs from packed policy loss
  seonjinn
  committed
  ddae591
- fix(draft): resolve context parallel size at finish
  seonjinn
  committed
  88643a6
- test(draft): assert context parallel gradient correction
  seonjinn
  committed
  c49fd46
- Add chunk parametrization, raise tests, rename shadowed var, document invariant
  seonjinn
  committed
  42cefa8
- fix(dflash): satisfy import ordering
  seonjinn
  committed
  d43a5b6
- fix(draft): share TP-safe DSpark Markov embeddings
  seonjinn
  committed
  3e6ecae
- test(draft): assert streaming loss normalization shape
  seonjinn
  committed
  8887579
- fix(grpo): aggregate CUDA draft metrics safely
  seonjinn
  committed
  0728112
- fix(metrics): share CUDA-safe step reduction
  seonjinn
  committed
  3922322
- Validate caller-supplied bin_ids range in streaming soft CE
  seonjinn
  committed
  895e0c2
- test(dspark): cover target-SP masks and graph config
  seonjinn
  committed
  f272799
- test(dspark): assert reconstructed mask block shape
  seonjinn
  committed
  9735386
- fix(dspark): reconstruct target-SP mask blocks
  seonjinn
  committed
  815b783
- Guard nccl_reshard bulk receive with fail-closed refit lifecycle
  seonjinn
  committed
  8a03385
- Test wire-decode failure surface, document vestigial control hook
  seonjinn
  committed
  f85be2e
- Merge PR9 dependency into PR10
  seonjinn
  committed
  59c875a
- Merge refreshed PR10 dependency into PR11
  seonjinn
  committed
  be3b9d1
- test(refit): match wire-status error contract
  seonjinn
  committed
  be40c54
- Merge PR10 refit test correction into PR11
  seonjinn
  committed
  4649436
- Restore TP shape agreement, gate CP inheritance, tolerate zero-window probes
  seonjinn
  committed
  1574946
- test(refit): initialize NCCL lifecycle fixture
  seonjinn
  committed
  da1bd0e
- Merge PR10 NCCL fixture correction into PR11
  seonjinn
  committed
  2c97a5b
- test(refit): cover non-owner payload allocation failure
  seonjinn
  committed
  ea57fed
- test(draft): prove CP finish normalization parity
  seonjinn
  committed
  d4129f6
- test(draft): report CP normalization mutation
  seonjinn
  committed
  47922aa
- Merge refreshed PR10 payload coverage into PR11
  seonjinn
  committed
  443e724
- Run reshard preflight inside the fail-closed refit guard
  seonjinn
  committed
  b4f7744
- Gate embedding TP reduction on consumer sharding, not head sharding
  seonjinn
  committed
  1deb4be
- Match real-quant finalize to the finalize_draft lifecycle signature
  seonjinn
  committed
  c31cead
- Pass finalize_draft in real-quant lifecycle test
  seonjinn
  committed
  4a7f8ca
- Accept coverage kwarg in real-quant weight loading
  seonjinn
  committed
  52e2553
- Merge current PR9 into PR10
  seonjinn
  committed
  8f39756
- Merge refreshed PR10 into PR11
  seonjinn
  committed
  a021b04
- fix(draft): preserve typed disabled refit config
  seonjinn
  committed
  df9daf6
Commits on Aug 23, 2026
- test(draft): cover asymmetric qwen moe attention widths
  seonjinn
  committed
  c8f1b6c
- test(draft): use locked python for focused runner
  seonjinn
  committed
  c47fd50
- test(draft): add qwen moe body behavior probe
  seonjinn
  committed
  2ba2da2
- fix(draft): allow asymmetric qwen attention projections
  seonjinn
  committed
  273dfc1
- fix(draft): preserve asymmetric qwen tp export shapes
  seonjinn
  committed
  fadae4e
- test(draft): include asymmetric tp export regression
  seonjinn
  committed
  d0c4f11
- test(vllm): accept coverage in real-quant IPC load
  seonjinn
  committed
  14239a8
- test(draft): select qwen mcore regressions
  seonjinn
  committed
  87c7572
- chore(draft): add headers to qwen test tools
  seonjinn
  committed
  a063fba
- test(draft): isolate qwen runner arguments
  seonjinn
  committed
  1fe5fde
- test(vllm): match unprepared coverage fixture
  seonjinn
  committed
  3c3fc10
- Merge current PR9 CI fix into PR10
  seonjinn
  committed
  3c9218b
- Merge current PR10 into PR11
  seonjinn
  committed
  edc838a
- Reject near-miss typo keys in Eagle3DraftConfig
  seonjinn
  committed
  af952e8
- Validate caller-supplied bin_ids range in streaming soft CE
  seonjinn
  committed
  9c792dc
- Add split parity regression test and shared draft_loss metric key
  seonjinn
  committed
  fd12c9c
- Fix parity test tolerance and stub for draft loss metric key
  seonjinn
  committed
  ec3a939
- Document and pin zero decay for draft norm and bias params
  seonjinn
  committed
  246b391
- Skip draft optimizer test collection without megatron bridge
  seonjinn
  committed
  866de64
- Fix TP checkpoint group handling and shard tagging in DSpark heads
  seonjinn
  committed
  8442dc9
- fix(draft): share TP-safe DSpark Markov embeddings
  seonjinn
  committed
  9f5c59c
- Gate embedding TP reduction on consumer sharding, not head sharding
  seonjinn
  committed
  f63dee4
- Remove redundant TP metadata collective and document DFlash loss contracts
  seonjinn
  committed
  ed16b85
- Add chunk parametrization, raise tests, rename shadowed var, document invariant
  seonjinn
  committed
  f0df99e
- Format restacked draft config regression
  seonjinn
  committed
  c4c0f32
- fix(refit): frame packed collective failures
  seonjinn
  committed
  eee45f4
- fix(modelopt): finalize cotrained draft refits
  seonjinn
  committed
  fa2afa0
- fix(refit): drain packed collective failures
  seonjinn
  committed
  bc853fe
- fix(loss): validate projected soft CE bin ids
  seonjinn
  committed
  e2da091
- fix(loss): avoid duplicate projected bin validation
  seonjinn
  committed
  03cc0e2
- test(vllm): add lifecycle test copyright header
  seonjinn
  committed
  4ee518b
Commits on Aug 24, 2026
- Merge remote-tracking branch 'nvidia-current/main' into fix/pr3715-hotpath-20260824
  seonjinn
  committed
  f185535
- test(draft): cover single DFlash metadata agreement
  seonjinn
  committed
  2d4ed91
Commits on Aug 27, 2026
- Merge origin/main into sna/megatron-sft-prepacked-inputs
  seonjinn
  committed
  accd30e
- Fix submodule pointers to match main after merge
  seonjinn
  committed
  8dfcf14
- Fix cp_normalize in direct-packed loss and lazy-import megatron in test
  seonjinn
  committed
  bc265c8
- fix(sft): finalize direct-packed config and docs
  seonjinn
  committed
  d0cff7e
- docs(sft): clarify context-parallel packing requirement
  seonjinn
  committed
  77bd36c
Commits on Sep 2, 2026
- Merge remote-tracking branch 'origin/main' into review/pr-3703-main-merge
  seonjinn
  committed
  064fcf9
Commits on Sep 3, 2026
- Merge remote-tracking branch 'origin/main' into review/pr-3703-main-merge
  seonjinn
  committed
  a141bfa
Commits on Sep 4, 2026
- Merge remote-tracking branch 'origin/main' into review/pr-3703-main-merge
  seonjinn
  committed
  cbda213
- feat(draft): expose draft soft-CE token_chunk_size as config
  seonjinn
  committed
  3e4cc82
Commits on Sep 5, 2026
- Fix import ordering in loss wrapper
  seonjinn
  committed
  f633cab
Commits on Sep 7, 2026
- Add draft.token_chunk_size to grpo reference config
  seonjinn
  committed
  d9068ad
Commits on Sep 10, 2026
Commits on Sep 11, 2026
- fix(sft): harden prepacked input validation
  seonjinn
  committed
  8ca4869
- fix(refit): align weight finalizer contract
  seonjinn
  committed
  80625d5
- fix(draft): address co-training review feedback
  seonjinn
  committed
  6017c37
Commits on Sep 12, 2026
- Merge branch 'main' into seonjinn/pr9-v3-latest-main
  seonjinn
  committed
  711cb08
- Merge branch 'main' into seonjinn/draft-cotrain-streaming-loss-20260818
  seonjinn
  committed
  122af10
- Merge branch 'main' into seonjinn/draft-cotrain-dspark-core
  seonjinn
  committed
  18c8f46
- Merge main into sna/megatron-sft-prepacked-inputs
  seonjinn
  committed
  b84eba1
- Merge streamed soft CE dependency updates
  seonjinn
  committed
  864e872
- Merge commit '864e8720eef9a56a3f3072e18ee3842d23f1e677' into HEAD
  seonjinn
  committed
  cf40c00
- Merge commit '122af108cea19cef0781cd9cf3a8c163ed34e56e' into HEAD
  seonjinn
  committed
  4dc91df
- Merge commit '18c8f468cae199290c015c5fde9f4f9fffb133f7' into HEAD
  seonjinn
  committed
  c6e1df1
- Merge commit '711cb08e6276bbb6460188966fb67ad6f641d75a' into review/pr3750-review-fixes-20260911
  seonjinn
  committed
  b86a79b
- Merge context-parallel refit dependency into PR #3757
  seonjinn
  committed
  cfc4de3
- Merge prepacked SFT dependency into PR #3757
  seonjinn
  committed
  626a143
- style: apply lint fixes
  seonjinn
  committed
  acffbba
- Merge branch 'main' into PR #3723
  seonjinn
  committed
  040ddcc
- Merge branch 'main' into PR #3726
  seonjinn
  committed
  77084ae
- Merge branch 'main' into PR #3744
  seonjinn
  committed
  37ba21a
- Merge branch 'main' into PR #3750
  seonjinn
  committed
  4ed8a37
- Merge branch 'main' into PR #3757
  seonjinn
  committed
  3ff97fd
Commits on Sep 15, 2026
Commits on Sep 16, 2026
- Merge remote-tracking branch 'origin/main' into review/pr-3723-remote-main-merge
  seonjinn
  committed
  b6bbca2
- Merge remote-tracking branch 'origin/main' into review/pr-3726-remote-main-merge
  seonjinn
  committed
  e9726dd
Commits on Sep 18, 2026
Commits on Sep 19, 2026
- Merge branch 'main' into review/pr-3723-remote-main-merge
  seonjinn
  committed
  e606829
- Merge branch 'main' into review/pr-3726-remote-main-merge
  seonjinn
  committed
  430d777
Commits on Sep 21, 2026
- chore(draft): merge latest main for PR #3723
  seonjinn
  committed
  a6f15d5
- chore(draft): merge latest main for PR #3726
  seonjinn
  committed
  868223f
Commits on Sep 22, 2026
- chore(draft): merge latest main for PR #3723
  seonjinn
  committed
  3477890
- chore(draft): merge latest main for PR #3726
  seonjinn
  committed
  0fc5b29
- Merge latest main after #3706 into PR #3726
  seonjinn
  committed
  9b0561b
- Merge latest main after #3706 into PR #3723
  seonjinn
  committed
  3186bd1
- test: collect draft worker coverage and exercise input guards
  seonjinn
  committed
  9641ffb
Commits on Sep 23, 2026
- test: record coverage after pytest-cov finalizes collection
  seonjinn
  committed
  76929b6
Commits on Sep 30, 2026
Commits on Oct 7, 2026
- Merge latest main into DFlash body
  seonjinn
  committed
  9f5aad7
- Merge latest main into DSpark objectives
  seonjinn
  committed
  0d1abdb
- docs: explain DFlash body contract and limits
  seonjinn
  committed
  816aae5
- docs: describe DSpark objectives and normalization
  seonjinn
  committed
  a8efe44
- docs: correct DSpark token chunk description
  seonjinn
  committed
  2e7667e
- Merge latest main into PR #3744
  seonjinn
  committed
  12c2799
- Merge latest main into PR #3750 and clarify CP draft refit
  seonjinn
  committed
  e88f972
- docs(draft): clarify packed DFlash and DSpark setup
  seonjinn
  committed
  6d6c30b
Commits on Oct 8, 2026
- Merge NVIDIA-NeMo/RL main into DFlash body PR
  seonjinn
  committed
  6c6f18c
- fix(draft): use supported all-gather for TP metadata
  seonjinn
  committed
  3f10b44
Commits on Oct 13, 2026
- Merge latest main into DFlash tensor-parallel body
  seonjinn
  committed
  2f3d2e0
- docs(draft): clarify DFlash weight ownership
  seonjinn
  committed
  e7becda
- docs(draft): explain DFlash body usage and computation
  seonjinn
  committed
  9ef0240
- docs(draft): explain DFlash body at a conceptual level
  seonjinn
  committed
  c8b3c60
- docs(draft): keep DFlash overview conceptual
  seonjinn
  committed
  227e965
Commits on Oct 15, 2026
- fix(draft): address DFlash body review threads
  seonjinn
  committed
  8176428
Commits on Oct 18, 2026

Commits

Commits on Aug 21, 2026

Commits on Aug 22, 2026

Commits on Aug 23, 2026

Commits on Aug 24, 2026

Commits on Aug 27, 2026

Commits on Sep 2, 2026

Commits on Sep 3, 2026

Commits on Sep 4, 2026

Commits on Sep 5, 2026