From 108a1f5820c3f0d13df1710c2247c1c4a039b4f7 Mon Sep 17 00:00:00 2001 From: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com> Date: Fri, 2 Oct 2026 16:35:58 -0400 Subject: [PATCH 1/2] fix(distill): flatten MFTLoss labels along with the logits MFTLoss.forward flattens both logit tensors to (batch * positions, vocab) and documents that it only assumes the class dimension is last, so leading dimensions are meant to be allowed. The labels were passed through untouched, and _prepare_corrected_distributions then rejects anything that is not 1D: ValueError: Logits must be a 2D tensor and labels must be a 1D tensor. So the shapes a language model actually produces -- (batch, seq_len, vocab) logits against (batch, seq_len) labels -- could not be used, which is the setting Minifinetuning is for. A caller had to flatten the labels themselves, and nothing said so. Flatten them with the logits. The existing test drives MFTLoss through a vision model, whose logits are already 2D and labels 1D, which is why this went unseen. Signed-off-by: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com> --- CHANGELOG.rst | 1 + modelopt/torch/distill/losses.py | 6 ++++-- tests/unit/torch/distill/test_distill.py | 17 +++++++++++++++++ 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4f621df5dea..f8434620609 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -94,6 +94,7 @@ Changelog - Fix ``megatron_generate`` dropping the VLM vision inputs (``pixel_values`` / ``image_grid_thw`` / ``image_sizes``) after the first generated token when KV-cache decoding is off, including the automatic fallback under sequence parallelism, which made generation silently ignore the image. No other ModelOpt feature is affected. - Fix two issues in the vLLM offline hidden-state dump (``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py``) that only surface on large runs. **Resume:** the filter that skips conversations whose ``.pt`` already exists now runs with ``load_from_cache_file=False``. It depends on on-disk state, which is not part of the fingerprint ``datasets`` computes from the function and the dataset, so with a persistent HF cache reused across a resumed or requeued run the cached "keep everything" result from an earlier run was replayed and the dump re-generated and overwrote conversations it had already finished (observed: tens of thousands of ``.pt`` rewritten while the output count stayed flat). **Staging:** generation is now chunked (``--save-chunk-size``, default 256), so each chunk is saved and its staged hidden states freed before the next chunk is generated. Previously the whole dataset was generated before anything was saved, which kept every conversation staged in the connector's ``shared_storage_path`` (``/dev/shm``, i.e. RAM, by default) at once and exhausted it partway through large dumps. Chunking also makes the dump incrementally durable, so an interrupted run keeps its finished conversations and resumes from them. The save path now also frees each conversation's staged hidden states in a ``finally``, so a conversation skipped mid-loop (e.g. a short ``loss_mask``) can no longer leak its staging file, and conversation ids are validated as plain filenames before being used to build output paths. - Hybrid (e.g. Nemotron-H) checkpoints saved by the ``examples/megatron_bridge`` scripts now record their layer spec in ``run_config.yaml`` in a form that reloads, so they can be converted to HuggingFace; a checkpoint saved by an earlier release still needs its ``model.hybrid_stack_spec`` block replaced by hand. +- Fix ``MFTLoss`` rejecting sequence-shaped logits. Labels are now flattened alongside the logits, so ``(batch, seq_len)`` labels work against ``(batch, seq_len, vocab)`` logits without reshaping them first. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/distill/losses.py b/modelopt/torch/distill/losses.py index 832437b7e59..d094ac9e899 100644 --- a/modelopt/torch/distill/losses.py +++ b/modelopt/torch/distill/losses.py @@ -103,7 +103,9 @@ def forward( Args: logits_s: Student's logits, treated as prediction. logits_t: Teacher's logits, treated as training target. - labels: Labels for the ground truth, used to prepare the corrected teacher distributions. + labels: Labels for the ground truth, used to prepare the corrected teacher + distributions. Flattened alongside the logits, so it carries one label per + position, e.g. ``(batch, seq_len)`` against ``(batch, seq_len, vocab)`` logits. .. note:: @@ -115,7 +117,7 @@ def forward( target_logits: torch.Tensor = logits_t / self._temperature # (B, ..., C) target_logits = target_logits.view(-1, target_logits.size(-1)) # (new B, C) soft_targets = self._prepare_corrected_distributions( - target_logits, labels, self._threshold, apply_threshold_to_all=True + target_logits, labels.reshape(-1), self._threshold, apply_threshold_to_all=True ) kd_loss = F.kl_div( diff --git a/tests/unit/torch/distill/test_distill.py b/tests/unit/torch/distill/test_distill.py index 69dec86b7f5..14f7108a4b6 100644 --- a/tests/unit/torch/distill/test_distill.py +++ b/tests/unit/torch/distill/test_distill.py @@ -127,6 +127,23 @@ def test_distillation_model_mft(): assert isinstance(loss, torch.Tensor) and loss.numel() == 1 +def test_mft_loss_accepts_sequence_shaped_logits(): + """MFTLoss flattens the logits it is given, so the labels have to follow them.""" + torch.manual_seed(0) + batch, seq_len, vocab = 2, 8, 50 + logits_s = torch.randn(batch, seq_len, vocab) + logits_t = torch.randn(batch, seq_len, vocab) + labels = torch.randint(0, vocab, (batch, seq_len)) + + loss = mtd.MFTLoss()(logits_s, logits_t, labels) + + # One label per position, so flattening first must not change the result. + flattened = mtd.MFTLoss()( + logits_s.reshape(-1, vocab), logits_t.reshape(-1, vocab), labels.reshape(-1) + ) + assert torch.allclose(loss, flattened) + + def test_distillation_mode_default_config(): student = tiny_mobilenet() with pytest.raises(AssertionError): From bef29e8b9aff7f280d649eb932f3f4f7621536d3 Mon Sep 17 00:00:00 2001 From: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com> Date: Sat, 3 Oct 2026 12:08:10 -0400 Subject: [PATCH 2/2] Drop the changelog entry per review AAnoosheh noted this fix isn't significant enough to warrant a changelog entry. Signed-off-by: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com> --- CHANGELOG.rst | 1 - 1 file changed, 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index f8434620609..4f621df5dea 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -94,7 +94,6 @@ Changelog - Fix ``megatron_generate`` dropping the VLM vision inputs (``pixel_values`` / ``image_grid_thw`` / ``image_sizes``) after the first generated token when KV-cache decoding is off, including the automatic fallback under sequence parallelism, which made generation silently ignore the image. No other ModelOpt feature is affected. - Fix two issues in the vLLM offline hidden-state dump (``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py``) that only surface on large runs. **Resume:** the filter that skips conversations whose ``.pt`` already exists now runs with ``load_from_cache_file=False``. It depends on on-disk state, which is not part of the fingerprint ``datasets`` computes from the function and the dataset, so with a persistent HF cache reused across a resumed or requeued run the cached "keep everything" result from an earlier run was replayed and the dump re-generated and overwrote conversations it had already finished (observed: tens of thousands of ``.pt`` rewritten while the output count stayed flat). **Staging:** generation is now chunked (``--save-chunk-size``, default 256), so each chunk is saved and its staged hidden states freed before the next chunk is generated. Previously the whole dataset was generated before anything was saved, which kept every conversation staged in the connector's ``shared_storage_path`` (``/dev/shm``, i.e. RAM, by default) at once and exhausted it partway through large dumps. Chunking also makes the dump incrementally durable, so an interrupted run keeps its finished conversations and resumes from them. The save path now also frees each conversation's staged hidden states in a ``finally``, so a conversation skipped mid-loop (e.g. a short ``loss_mask``) can no longer leak its staging file, and conversation ids are validated as plain filenames before being used to build output paths. - Hybrid (e.g. Nemotron-H) checkpoints saved by the ``examples/megatron_bridge`` scripts now record their layer spec in ``run_config.yaml`` in a form that reloads, so they can be converted to HuggingFace; a checkpoint saved by an earlier release still needs its ``model.hybrid_stack_spec`` block replaced by hand. -- Fix ``MFTLoss`` rejecting sequence-shaped logits. Labels are now flattened alongside the logits, so ``(batch, seq_len)`` labels work against ``(batch, seq_len, vocab)`` logits without reshaping them first. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^