feat(mopd): support multiple teacher checkpoints in full-vocabulary o… - #4045
Merged
Merged
Conversation
RayenTian
force-pushed
the
ruit/full_vocab_v2_multi_teacher
branch
from
September 8, 2026 14:57
7384eba to
5dc5548
Compare
RayenTian
force-pushed
the
ruit/mopd_full_vocabulary_v2
branch
2 times, most recently
from
September 8, 2026 21:09
b4c4f6e to
0823c83
Compare
RayenTian
force-pushed
the
ruit/full_vocab_v2_multi_teacher
branch
from
September 8, 2026 21:38
5dc5548 to
6305f72
Compare
RayenTian
force-pushed
the
ruit/mopd_full_vocabulary_v2
branch
from
September 9, 2026 19:41
0823c83 to
8603a46
Compare
RayenTian
force-pushed
the
ruit/full_vocab_v2_multi_teacher
branch
from
September 9, 2026 21:18
b8bafa3 to
15839dc
Compare
RayenTian
force-pushed
the
ruit/full_vocab_v2_multi_teacher
branch
from
September 10, 2026 18:07
15839dc to
63c2754
Compare
github-actions
Bot
added
the
Documentation
RayenTian
added
the
CI:Lfast
Contributor
Author
|
/ok to test 63c2754 |
RayenTian
marked this pull request as ready for review
September 11, 2026 17:54
Squashes four commits from this branch's iteration on top of #3978's opd_full feature: - Multiple teacher checkpoints on the hidden_states path: each physical teacher gets a stable index (OPD_FULL_TEACHER_INDEX_FIELD), the student loads one LM-head shard per teacher, and a microbatch's rows are grouped by teacher_index and projected through the matching shard. - Student pipeline parallelism on the hidden_states path: earlier PP stages hold no teacher LM-head shard (only the last stage runs the loss), but still join the whole-world dist_checkpointing.load collective with an empty sharded request so the collective stays balanced. - Review-finding fixes: repaired 16 test-double failures the teacher_index plumbing left behind, raised instead of silently falling through to the single-teacher projection when a routing column is missing, fixed pipeline-rank-asymmetric eviction tracking, and caught heterogeneous teacher hidden sizes/dtypes at load time instead of an opaque jagged- collate error later. - Rebuilt the multi-teacher routing after a silent rebase-onto-main conflict: main reordered reconstruct_opd_full_teacher_logits into validate -> narrow -> project (so the LM-head matmul only ever sees this rank's CP window), but this branch's multi-teacher commits were written against the older project-then-narrow order and the rebase merged both textually into an incoherent function. Kept main's order and slotted the routing checks into the validate phase. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: ruit <ruit@nvidia.com>
…_size
chunk_size on the opd_full divergence kernels
(ChunkedDistributedReverseKLToFixedLogits / CrossEntropyToFixedLogits /
Entropy) only bounded the transient fp32 working set inside each kernel's
internal forward/backward loop over sequence chunks -- the single
autograd.Function node wrapping that loop still called
ctx.save_for_backward() on the whole, un-chunked [B, S, V_local] student
and teacher logits. So backward's peak retained memory scaled with the
full local vocabulary width regardless of chunk_size, which is what an
OOM during opd_full backward traced back to.
Restructures the three kernels into thin dispatchers over a new
_ChunkedStudentTeacherChunk autograd.Function: one node per sequence
chunk, each saving only its own chunk (materialized via .clone(), not
.contiguous() -- a middle-dim slice of an already-contiguous tensor can
itself report is_contiguous() == True, so .contiguous() would be a no-op
that still aliases the full parent's storage). torch.cat's own backward
routes the upstream gradient back to each chunk node. External
ClassName.apply(...) call sites are unchanged.
Verified forward/backward numerics against a from-scratch reference
(reverse_kl, cross_entropy, entropy; chunk_size in {1, 2, 3, 100}) plus
self-distillation (KL ~= 0), teacher-no-gradient, misaligned-shape
rejection, and bf16 round-trip -- all within rtol=1e-4. Confirmed
structurally that the autograd graph now holds one independent
_ChunkedStudentTeacherChunkBackward node per chunk, each fed by its own
CloneBackward0, rather than one node holding the whole tensor.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Signed-off-by: ruit <ruit@nvidia.com>
… padding/truncation materialize()'s pad_to_seqlen step forced teacher_full_hidden_states / teacher_full_logits to pad to the global cross-DP max sequence length even though they never feed the model forward (only the per-sequence loss, which reads each row's own real length) -- at 98304 tokens x hidden_size wide across a whole DP-rank's batch this produced 100+ GiB single allocations and was the direct cause of opd_full OOMs on long-context full-vocabulary runs. Exempts OPD_FULL_FIELDS from that padding in codec.py::materialize, and adds the matching exemption to the two downstream checks that assumed every tensor field shares one padded seqlen: megatron/data.py's get_and_validate_seqlen (hit live once the payload was allowed to keep its own shorter length) and batched_data_dict.py's truncate_tensors (the dynamic_batching microbatch path's analog, not yet exercised but same gap). Signed-off-by: ruit <ruit@nvidia.com>
RayenTian
force-pushed
the
ruit/full_vocab_v2_multi_teacher
branch
from
September 11, 2026 23:31
63c2754 to
549953d
Compare
Contributor
Author
|
/ok to test 549953d |
…ption test_materialize_pads_a_3d_payload_on_the_seq_dim_only used OPD_FULL_LOGITS_FIELD as its example 3-D column, so it broke when materialize() started exempting OPD_FULL_FIELDS from pad_to_seqlen (see 549953d). Swaps it to a synthetic field name to keep covering the generic 3-D pad-spec construction it actually tests, and adds a dedicated regression test asserting the new exemption itself. Signed-off-by: ruit <ruit@nvidia.com>
Contributor
Author
|
/ok to test 7db9f0d |
yfw
reviewed
Sep 15, 2026
5 tasks
The hang the non-last pipeline stages guard against comes from Megatron-Bridge's read_train_state broadcast during path resolution, not from dist_checkpointing.load, which is rank-local under validate_access_integrity=False. Reword the six spots that claimed otherwise. Also make the opd_full transport keys required rather than .get-with- fallback in TQPolicy, and drop the teacher_index default on load_opd_full_teacher_lm_head: both silently degrade to "teacher 0" on a caller bug. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: bg51717 <biguo@nvidia.com>
…allbacks Three silent-corruption paths in full-vocabulary MOPD: - teacher_index was numbered from the sorted alias list, but which alias represents a checkpoint shared by several agents is first-seen-wins over teacher_model_by_agent_name. The index outlives a run inside a data-plane checkpoint, so reordering the YAML repointed already-tagged rows at another teacher's LM head. Number by the deduplicated checkpoint instead, record the per-index checkpoint paths in the data-plane checkpoint metadata, and reject a restore that would renumber them. Stored as a list ordered by index rather than a map, since the metadata round trips through JSON. - reconstruct_opd_full_teacher_logits branched on how many LM heads were loaded, so a single-teacher run ignored the routing column entirely and projected a row tagged for an unloaded teacher through the only head it had. Route whenever the column is present. - truncate_tensors exempted the teacher payload from dynamic batching's narrow, which made every microbatch shorter than the fetch raise in the loss. Clamp to min(truncated_len, its own width) instead: materialize leaves the payload at its natural width, which can already be shorter. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: bg51717 <biguo@nvidia.com>
…ycle Every raise and lifecycle branch added for multi-teacher opd_full was invisible to the suite: deleting the guard, the row-count check, or the evict bookkeeping left it green. - teacher worker: a routing column configured without an index is rejected before the forward runs. - reconstruct/prepare: several heads without a routing column, and an index column whose length does not match the payload. - TQPolicy: prepare_step, prepare_val_partition and train_from_meta all request the teacher-index column, not just train_microbatches_from_meta. - megatron worker: evict reloads every teacher once per release and resets the flag, ranks off the last pipeline stage re-enter the reload from their recorded paths, offload parks every head, a second teacher with another hidden size is rejected at load, a shardless rank still records its path, and owner resolution returns None only off the last stage. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: bg51717 <biguo@nvidia.com>
This reverts commit a6e20cc. Splitting the backward into per-chunk autograd nodes did not bound the retained memory: the per-chunk clones sum to exactly one full copy each of the student and teacher logits, so the retained set never varied with chunk_size. Each chunk's SliceBackward0 then materializes another full [B, S, V_local] zero tensor that the engine accumulates into a second full buffer, and the forward loop briefly holds the originals alongside their clones. Measured peak (B=2, S=1024, V_local=8192, chunk 128): fp32 232 -> 296 (+28%), bf16 unchanged. defer_fp32_logits defaults to False, so the fp32 row is the default path. Backward also got slower in both dtypes. The two designs agree numerically and the kernel tests pass either way. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: bg51717 <biguo@nvidia.com>
torch.empty_like(..., dtype=torch.float32) held a full fp32 [B, S, V_local] buffer that autograd immediately cast back down to the leaf's dtype. Allocate it in the logits' dtype instead and move both multiplies ahead of the copy_, so the arithmetic still happens in fp32 and bf16 avoids an extra rounding. Measured peak (B=2, S=1024, V_local=8192): bf16 168 -> 144, fp32 unchanged, no forward cost. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: bg51717 <biguo@nvidia.com>
Signed-off-by: bg51717 <biguo@nvidia.com> # Conflicts: # nemo_rl/algorithms/single_controller_utils/setup.py # nemo_rl/data_plane/codec.py # nemo_rl/models/megatron/data.py # nemo_rl/models/policy/tq_policy.py
Member
|
/ok to test 5db3ec9 |
Defaulting it to 0 tags every teacher's payload rows as teacher 0, and both lookups that key on the attribute -- setup.py's per-checkpoint RPC dict and opd.py's checkpoint-by-index map -- are dict comprehensions, so the collision collapses them to a single entry instead of raising. The student would then project every row through one teacher's LM head at legal shapes. Also fix the attribute's comment: create_teacher_worker_groups has ordered the index by checkpoint since teacher_configs_by_index landed, not by alias. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: bg51717 <biguo@nvidia.com>
Buffered rows in a data-plane checkpoint carry a teacher_index each and nothing re-tags them on restore, so _maybe_restore_native_data_plane_checkpoint refuses a checkpoint whose opd_full teacher list differs from the one this run would number. That guard shipped without a test -- no test in the repo referenced opd_full_teacher_checkpoints at all. Covers renumbering, a dropped teacher, both directions of the untagged/tagged transition, and the matching case that must still restore. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: bg51717 <biguo@nvidia.com>
Member
|
/ok to test 105b4f6 |
…estore fixture This test builds SingleControllerActor via object.__new__, bypassing __init__, so it never gets the _teacher_coordinator attribute __init__ always sets (to a coordinator or None). The data-plane checkpoint-save path added by this PR reads self._teacher_coordinator unconditionally, so this fixture's controller raised AttributeError instead of running. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: bg51717 <biguo@nvidia.com>
Member
|
/ok to test c6fdfc3 |
yfw
approved these changes
Sep 21, 2026
4 tasks
lauradang
added a commit
to lauradang/RL
that referenced
this pull request
Sep 22, 2026
Bring main a9c3178 (25 commits since the previous merge base 702cd79) into the branch on top of the #5 submodule bump (00621c2). Conflicts: - nemo_rl/algorithms/single_controller_utils/setup.py and tests/unit/single_controller/test_setup.py: keep main's sharded Gym stacks (NVIDIA-NeMo#3373), multi-teacher MOPD (NVIDIA-NeMo#4045), critic optimizer rebuild (NVIDIA-NeMo#4213) and Mooncake storage checkpoints (NVIDIA-NeMo#3898) next to the Megatron token-capture setup guards and backend wiring; the vLLM-only capture guard is replaced by the backend-aware one. - tests/unit/models/generation/test_vllm_token_capture_hosting.py: keep main's stripped staged token arrays (NVIDIA-NeMo#4206) on top of the shared vLLM chain-cache prefix resolution. - pyproject.toml (semantic conflict, auto-merged by git): main NVIDIA-NeMo#3566 removed `exclude-dependencies = ["nvidia-cutlass-dsl-libs-base"]` in favour of per-version dependency-metadata overrides for CUTLASS DSL 4.5.x, while this branch's mcore extra pinned nvidia-cutlass-dsl[cu13]==4.8.0.dev0 (68b63da). Without the exclusion that pin drags in the pre-release nvidia-cutlass-dsl-libs-base==4.8.0.dev0, which uv refuses for a transitive dependency, so `uv lock` no longer resolves. Pin 4.6.2 instead: it is the stable release flashinfer-python 0.6.18.post1 actually requires (>=4.6.2a0), and it is the pin NVIDIA-NeMo#4139 uses for the same Megatron-Bridge bump on main. - uv.lock: regenerated from main's copy with uv 0.11.28 (the docker/Dockerfile pin), so it carries main's new git pins (Model-Optimizer, sglang, Lens, TransferQueue, nvidia-resiliency-ext, torch_memory_saver) together with the Gym and Megatron-Bridge main pins and the flashinfer index spelling from #5. `uv lock --check` passes with uv 0.11.28 and 0.12.5. Submodule pins: Automodel follows main (72daceffa); Gym (9fc05c0ff) and Megatron-Bridge (fcea83a85) stay at the #5 pins, which descend from main's. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Laura Dang <laurad@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
…pd_full
Drops the "exactly one unique teacher checkpoint" restriction on on_policy_distillation.full's hidden_states path. Each physical teacher gets a stable index; teacher worker groups tag every payload row they write with it (new OPD_FULL_TEACHER_INDEX_FIELD column); the student loads one LM-head shard per teacher (dict keyed by index, same offload/evict lifecycle as before) and groups a microbatch's rows by teacher_index to project each group through its own shard rather than requiring a single teacher for the whole run.
What does this PR do ?
Add a one line overview of what this PR aims to accomplish.
Issues
List issues that this PR closes (syntax):
Usage
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information