feat(megatron-bridge): SFT-masked data support in distillation - #2113
feat(megatron-bridge): SFT-masked data support in distillation#2113yueshen2016 wants to merge 3 commits into
Conversation
The distillation example only consumes pretraining-style data (GPTDataset over pre-tokenized
blends, NullTokenizer), so the loss is computed over every token. For distilling an
instruction-tuned model it is usually preferable to train on prompt/response pairs and mask
the loss to the response, matching how the model was fine-tuned.
Adds --sft and --sft_dataset_root, which switch the data path to Bridge's
FinetuningDatasetConfig (NeMo-style GPTSFTDataset) reading training.jsonl / validation.jsonl of
{"input": <prompt>, "output": <response>} records.
Details:
* prompt_template="{input}{output}" tokenizes input+output verbatim (adjacent placeholders,
no separator), label_key="output" with answer_only_loss=True masks the loss to the response
(answer_start_idx == len(context_ids)), and truncation_field="input" truncates the context
when a pair exceeds seq_length.
* SFT reads raw text, so it uses the model's real HuggingFace tokenizer; the pretraining path
consumes pre-tokenized data and keeps NullTokenizer.
* The response-only loss mask requires per-token loss reduction to combine correctly across
context-parallel ranks, so calculate_per_token_loss is enabled and average_in_collective is
disabled under --sft. Both are untouched on the pretraining path.
Opt-in: without --sft the existing mock/blend data path is unchanged.
Signed-off-by: James Shen <yueshen@nvidia.com>
📝 WalkthroughWalkthroughThe distillation example adds SFT options, validates the dataset root and input combinations, loads prompt/completion JSONL data with response-only masking, uses HuggingFace tokenization, calculates per-token loss, and applies SFT-specific distributed gradient reduction. ChangesSFT distillation flow
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant FinetuningDatasetConfig
participant HuggingFaceTokenizer
participant ModelProviders
participant DistributedGradientReduction
CLI->>FinetuningDatasetConfig: configure SFT dataset
FinetuningDatasetConfig->>HuggingFaceTokenizer: tokenize prompt/completion records
HuggingFaceTokenizer->>ModelProviders: provide tokenized inputs with per-token loss
ModelProviders->>DistributedGradientReduction: reduce SFT gradients without pre-averaging
Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/megatron_bridge/distill.py`:
- Around line 275-279: Update the existing pretraining data-input validation
before the SFT-specific check so it does not require --data_paths or
--use_mock_data when args.sft is enabled. Preserve the current requirement for
non-SFT runs, and keep the --sft_dataset_root validation in the SFT argument
flow unchanged.
- Around line 403-417: Before constructing FinetuningDatasetConfig in the --sft
flow, validate args.sft_dataset_root, required split files, each JSONL record’s
schema, and its size against the approved limits; reject invalid input before
the dataset builder parses it. Do not rely on seq_length truncation to constrain
parsing, and only pass validated data into the existing dataset construction
path.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 14bae0c7-1e6d-41f6-9d65-0a89b9ce393b
📒 Files selected for processing (1)
examples/megatron_bridge/distill.py
|
/claude review |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2113 +/- ##
=======================================
Coverage 78.60% 78.60%
=======================================
Files 522 522
Lines 60167 60167
=======================================
Hits 47294 47294
Misses 12873 12873
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| if args.sft and not args.sft_dataset_root: | ||
| raise ValueError( | ||
| "--sft requires --sft_dataset_root (a directory with training.jsonl / validation.jsonl)." | ||
| ) |
There was a problem hiding this comment.
[CRITICAL Algorithm] The pre-existing data-source check at line 265-266 runs before this one and does not know about --sft, so the invocation documented in the PR description fails immediately:
python distill.py --sft --sft_dataset_root /path/to/data ...
ValueError: Must provide either --data_paths or set --use_mock_data.
--use_mock_data defaults to False and --data_paths to None, so with --sft alone line 266 raises and the SFT branch at line 396 is never reached. The only way to run the feature today is to also pass --data_paths <anything> or --use_mock_data — whose values are then silently ignored, because if args.sft: wins the dataset branch. That makes the headline feature unreachable as documented, and reachable only via a misleading incantation.
Fix: teach the existing check about the new source, and reject the ignored combination so a stale --data_paths in a launch script doesn't look like it's in use:
# Sanity checks
if not args.sft and not args.use_mock_data and not args.data_paths:
raise ValueError("Must provide one of --data_paths, --use_mock_data, or --sft.")
if args.sft and (args.data_paths or args.use_mock_data):
raise ValueError("--sft is mutually exclusive with --data_paths / --use_mock_data.")(CodeRabbit raised the same line; repeating because it blocks the feature end-to-end.)
There was a problem hiding this comment.
Accepted, both parts — 080f4de and 8c19f38.
Verified this was a genuine regression rather than a theoretical one: the working branch these commits were ported from carried the exemption, and the port onto main dropped it, so the PR as first pushed could not run its own documented invocation.
if not args.sft and not args.use_mock_data and not args.data_paths:
raise ValueError("Must provide either --data_paths or set --use_mock_data.")
if args.sft and (args.data_paths or args.use_mock_data):
raise ValueError(
"--sft is mutually exclusive with --data_paths / --use_mock_data: the SFT branch wins "
"the dataset selection, so those inputs would be silently ignored."
)Took the mutual-exclusion check too — that is the more valuable half, since it turns the misleading incantation into an error instead of leaving it as a silent no-op.
| dataset_kwargs={ | ||
| "prompt_template": "{input}{output}", | ||
| "label_key": "output", | ||
| "truncation_field": "input", | ||
| "answer_only_loss": True, | ||
| "add_bos": False, | ||
| "add_eos": True, | ||
| }, |
There was a problem hiding this comment.
[IMPORTANT Compatibility] add_bos: False is hardcoded, and prompt_template="{input}{output}" applies no chat template. Together these mean every training sequence starts directly at the first token of input — no BOS, no role markers.
For a model family whose tokenizer/chat template always prepends BOS (Llama <|begin_of_text|>, Gemma <bos>, Mistral <s>), that is a train/inference skew: distillation runs on sequences the model never sees at serving time, and nothing warns about it. The PR body's stated intent is the opposite — "matching how the model was fine-tuned" — and this silently doesn't for those models. It happens to be correct for the Nemotron-Nano-3 run in the Testing section, which is why the run looked clean.
The escape hatch (bake the full chat-formatted prompt, including BOS and role markers, into the "input" field) is real but undocumented — neither the --sft_dataset_root help text nor the README says the text is fed verbatim, so the natural reading of {"input": <prompt>, "output": <response>} is plain instruction text.
Two options, either is fine:
- Derive it from the tokenizer instead of hardcoding, so BOS-requiring models get BOS:
"add_bos": AutoTokenizer.from_pretrained( args.student_hf_path, trust_remote_code=args.trust_remote_code ).bos_token is not None,
- Keep
add_bos=False(verbatim is a defensible contract) but state the requirement where users will read it — in the--sft_dataset_roothelp string and the README:inputmust contain the fully templated prompt, including any BOS and role/turn markers the model expects; no chat template or BOS is added.
There was a problem hiding this comment.
Accepted — took option 2, in 080f4de.
Deriving add_bos from the tokenizer is the more automatic fix, but it would silently change tokenization for the runs already validated against this path, and it only covers BOS while leaving the role/turn-marker half of the skew unaddressed. "The fields are tokenized verbatim" is the contract I actually want; it was just undocumented.
So the requirement is now stated in all three places a user could look:
- the
--sft_dataset_roothelp text, - a comment next to the
dataset_kwargsthat implement it, - the Data Preparation section of
examples/megatron_bridge/README.md.
Each says the same thing: both fields are tokenized verbatim, no chat template is applied and no BOS is prepended, so bake in any role/turn markers and BOS the model expects.
| ), | ||
| tokenizer=TokenizerConfig( | ||
| tokenizer_type="NullTokenizer", vocab_size=distill_provider.vocab_size | ||
| tokenizer=( | ||
| # SFT reads raw text, so it needs the model's real tokenizer; the pretraining path | ||
| # consumes pre-tokenized data and keeps NullTokenizer. | ||
| TokenizerConfig( | ||
| tokenizer_type="HuggingFaceTokenizer", | ||
| tokenizer_model=args.student_hf_path, | ||
| hf_tokenizer_kwargs={"trust_remote_code": args.trust_remote_code}, | ||
| ) |
There was a problem hiding this comment.
[IMPORTANT Compatibility] The SFT path tokenizes with --student_hf_path's tokenizer, but the KD loss is computed against teacher logits over the same token ids. If the teacher and student come from different families with different vocabularies (which the existing --teacher_hf_path / --student_hf_path interface permits, and which the README's "distill a 4B student from an 8B teacher" framing invites), the teacher receives the student's ids and its logits are meaningless — silently producing a garbage KD target rather than an error.
The pretraining path was structurally immune to this: NullTokenizer + pre-tokenized --data_paths meant the user chose one tokenization for both. Picking the student's tokenizer here quietly introduces the coupling.
Worth a guard next to the other --sft sanity checks, since both configs are already loaded for the VLM detection:
if args.sft and student_provider.vocab_size != teacher_provider.vocab_size:
raise ValueError(
"--sft tokenizes with the student's tokenizer; student and teacher must share a "
f"vocabulary (got {student_provider.vocab_size} vs {teacher_provider.vocab_size})."
)A warn_rank_0 would be acceptable too — the point is that this failure mode is currently invisible. Also worth stating in the help text / README that --sft requires teacher and student to share a tokenizer.
There was a problem hiding this comment.
Accepted — added the guard in 8c19f38, as a hard error rather than a warning.
The failure is silent and total (the teacher scores ids it never saw, so the KD target is noise), and the check cannot produce a false positive: both providers are built with the same make_vocab_size_divisible_by and TP here, so their padded vocab_size values are equal iff the base vocabularies are.
if args.sft and student_provider.vocab_size != teacher_provider.vocab_size:
raise ValueError(
"--sft tokenizes with the student's tokenizer, so student and teacher must share a "
f"vocabulary (got {student_provider.vocab_size} vs {teacher_provider.vocab_size})."
)Your framing of why this is new — the pretraining path was structurally immune because NullTokenizer plus pre-tokenized data means one tokenization feeds both — is now the comment above it.
| dataset_config = FinetuningDatasetConfig( | ||
| seq_length=args.seq_length, | ||
| dataset_root=args.sft_dataset_root, | ||
| seed=args.seed, | ||
| dataloader_type="batch", | ||
| do_validation=True, | ||
| do_test=False, |
There was a problem hiding this comment.
[SUGGESTION] do_validation=True is hardcoded, but --eval_iters accepts 0 (_nonnegative_int) and the rest of the script honors that as "no validation". Under --sft --eval_iters 0, FinetuningDatasetConfig will still build the validation split and therefore still require validation.jsonl to exist — so a user with training data only has to fabricate a dummy validation file to run a config the script otherwise supports.
Deriving it keeps the two knobs consistent:
| dataset_config = FinetuningDatasetConfig( | |
| seq_length=args.seq_length, | |
| dataset_root=args.sft_dataset_root, | |
| seed=args.seed, | |
| dataloader_type="batch", | |
| do_validation=True, | |
| do_test=False, | |
| dataset_config = FinetuningDatasetConfig( | |
| seq_length=args.seq_length, | |
| dataset_root=args.sft_dataset_root, | |
| seed=args.seed, | |
| dataloader_type="batch", | |
| do_validation=args.eval_iters > 0, | |
| do_test=False, |
There was a problem hiding this comment.
Accepted — applied verbatim in 8c19f38.
do_validation=args.eval_iters > 0. You are right that the rest of the script already treats --eval_iters 0 as "no validation", so requiring a dummy validation.jsonl to run a supported configuration was an inconsistency, not a constraint.
| if args.sft: | ||
| # The SFT loss mask covers only the response tokens, so the reduction must be | ||
| # per-token for it to combine correctly across context-parallel ranks. | ||
| provider.calculate_per_token_loss = True |
There was a problem hiding this comment.
[SUGGESTION] _build_model_provider is called for both student and teacher, so calculate_per_token_loss = True lands on both providers. That's harmless (the teacher's LM loss is zeroed out in adjust_distillation_model_for_mcore, and the flag has to agree with the average_in_collective=not args.sft setting on the single shared DistributedDataParallelConfig anyway), but it reads as if it were a student-only knob.
Since this and line 450 are two halves of one decision that must stay in sync, a one-line note here pointing at the DDP setting — or moving the comment to mention both — would keep a future edit from flipping one without the other.
There was a problem hiding this comment.
Accepted — 8c19f38.
Kept the flag where it is (it does have to agree with the single shared DistributedDataParallelConfig, so scoping it to the student would be wrong) and made the coupling explicit instead: the comment now says it lands on both providers, why that is harmless, and that it must stay in sync with average_in_collective=not args.sft below.
There was a problem hiding this comment.
Claude review — feat(megatron-bridge): SFT-masked data support in distillation
Scope: full review (trigger comment was bare /claude review). 1 file changed (+62/−4): examples/megatron_bridge/distill.py, reviewed in full, plus surrounding context in _distillation_provider.py, modelopt/torch/distill/plugins/megatron.py, and examples/megatron_bridge/README.md.
Findings
CRITICAL: 1 · IMPORTANT: 2 · SUGGESTION: 3
| # | Severity | Location | Issue |
|---|---|---|---|
| 1 | CRITICAL Algorithm | distill.py:275-278 |
The pre-existing --data_paths/--use_mock_data check runs first and doesn't know about --sft, so the documented invocation raises before the SFT branch is reached |
| 2 | IMPORTANT Compatibility | distill.py:410-417 |
Hardcoded add_bos=False with no chat template — train/inference skew for BOS-requiring model families; the verbatim-text contract is undocumented |
| 3 | IMPORTANT Compatibility | distill.py:462-470 |
SFT tokenizes with the student's tokenizer while KD targets come from the teacher — mismatched vocabularies silently produce garbage KD targets instead of erroring |
| 4 | SUGGESTION | distill.py:403-409 |
do_validation=True hardcoded ignores --eval_iters 0, forcing a dummy validation.jsonl |
| 5 | SUGGESTION | distill.py:303-306 |
calculate_per_token_loss and average_in_collective are two halves of one decision, ~145 lines apart, with no cross-reference |
| 6 | SUGGESTION | README.md |
The Distillation section documents only pre-tokenized --data_paths and mock data; --sft has no README entry and no note on the expected JSONL schema |
Most impactful
Finding 1 blocks the feature. With --sft --sft_dataset_root ... and nothing else, distill.py:266 raises Must provide either --data_paths or set --use_mock_data. before the new dataset branch at line 396 is ever evaluated. --use_mock_data defaults to False and --data_paths to None, so the only way to reach the SFT path today is to also pass one of them — after which its value is silently discarded, since if args.sft: wins the branch. The usage block in the PR description does not run as written. This is worth confirming with an actual launch before merge; the Testing section's Nemotron-Nano-3 run presumably passed --data_paths too, which would mask it.
Findings 2 and 3 share a shape: both are silent-wrong-result paths rather than crashes, and both are invisible in the one model configuration that was tested. The Nemotron-Nano-3 run has a shared teacher/student tokenizer and no mandatory BOS, so neither would have shown up in the reported loss curve. A vocab-size assertion and a documented prompt-format contract are cheap relative to debugging a distillation run that trains cleanly but produces a subtly wrong model.
The parts I checked and found correct:
- Loss reduction.
calculate_per_token_loss=Truepaired withaverage_in_collective=Falseis the right combination for a response-only mask under CP>1 — a pre-averaged collective would weight ranks equally regardless of how many unmasked tokens each holds. Consistent with the# zero when one CP rank has only context tokensguard already inLogitsAndIntermediatesLossBalancer.forward. - Opt-in isolation. Every behavioral change is gated on
args.sft; the mock and blend paths,NullTokenizer, andaverage_in_collective=Trueare byte-for-byte unchanged without the flag. Nomodelopt_stateschema, mode registration, or public API surface is touched, so there is no checkpoint-compat or restore-fidelity exposure. - Dataset kwargs.
prompt_template="{input}{output}"+label_key="output"+answer_only_loss=Truedoes mask to the response (answer_start_idx == len(context_ids)), andtruncation_field="input"correctly truncates the context rather than the labels.
Risk
Medium. The blast radius is one example script and fully opt-in, so no existing user workflow can regress — but the new path does not currently run as documented, and two of its silent-failure modes are outside what the reported test exercised. Finding 1 is a small, self-contained fix; findings 2 and 3 are a guard plus a docs sentence each.
Not blocking, but noted: the PR checklist marks no tests and no docs. tests/examples/megatron_bridge/test_distill.py exists and a --sft smoke case over a two-record JSONL fixture would have caught finding 1 outright.
… document verbatim SFT format --sft supplies its own data via --sft_dataset_root, but the pretraining sanity check still demanded --data_paths or --use_mock_data, so a valid "--sft --sft_dataset_root <dir>" invocation raised before reaching the SFT branch. Exempt SFT from that check. Also state the SFT record contract where users will read it (--sft_dataset_root help, the dataset_kwargs comment, and the README): add_bos=False plus a placeholder-only prompt_template means "input"/"output" are tokenized verbatim -- no chat template, no BOS, no role markers -- so models that expect those need them baked into the fields. Addresses CodeRabbit and claude[bot] review comments. Signed-off-by: James Shen <yueshen@nvidia.com>
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/megatron_bridge/README.md`:
- Around line 135-138: Update the validation.jsonl tokenization description near
the “Both fields are tokenized” text to clarify that inputs and outputs are
tokenized verbatim before the dataset automatically appends an EOS token via
add_eos=True. Preserve the existing guidance about manually including chat
markers or BOS tokens.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 997a1191-cc96-4369-ac17-019ddc8ccd91
📒 Files selected for processing (2)
examples/megatron_bridge/README.mdexamples/megatron_bridge/distill.py
🚧 Files skipped from review as they are similar to previous changes (1)
- examples/megatron_bridge/distill.py
| `validation.jsonl` of `{"input": <prompt>, "output": <response>}` records, which are tokenized with | ||
| the model's own HuggingFace tokenizer. Both fields are tokenized **verbatim** — no chat template is | ||
| applied and no BOS token is prepended — so if your model expects role/turn markers or a BOS token, | ||
| include them in the `"input"` field yourself. |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Document the automatic EOS token.
Line 136 says both fields are tokenized verbatim. However, examples/megatron_bridge/distill.py:398-423 sets add_eos=True, so the dataset adds an EOS token automatically. Clarify that the fields are tokenized verbatim before this automatic EOS addition.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/megatron_bridge/README.md` around lines 135 - 138, Update the
validation.jsonl tokenization description near the “Both fields are tokenized”
text to clarify that inputs and outputs are tokenized verbatim before the
dataset automatically appends an EOS token via add_eos=True. Preserve the
existing guidance about manually including chat markers or BOS tokens.
- Reject `--sft` combined with `--data_paths` / `--use_mock_data`. The SFT branch wins the dataset selection, so those inputs were silently ignored -- a stale `--data_paths` in a launch script looked like it was in use. - Fail loudly when `--sft` is used with a teacher and student that do not share a vocabulary. SFT tokenizes raw text with the student's tokenizer and the KD target comes from the teacher's logits over those same ids, so a cross-family pair produced a garbage target rather than an error. The pretraining path was structurally immune (NullTokenizer + pre-tokenized data means one tokenization feeds both). - Derive `do_validation` from `--eval_iters` instead of hardcoding True, so a training-only `dataset_root` no longer has to carry a dummy `validation.jsonl` just to satisfy the dataset builder. - Cross-reference `calculate_per_token_loss` and `average_in_collective`, which are two halves of one decision that must stay in sync. Signed-off-by: James Shen <yueshen@nvidia.com>
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
examples/megatron_bridge/distill.py (1)
141-145: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick winMake
validation.jsonlconditional in the SFT documentation.When
--eval_iters 0,do_validationis false, so onlytraining.jsonlis required. Update the CLI help, error text, and README to state thatvalidation.jsonlis required only when--eval_iters > 0.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/megatron_bridge/distill.py` around lines 141 - 145, Update the SFT dataset documentation near the CLI help for the training/validation JSONL files, plus the corresponding validation/error text and README, to state that training.jsonl is always required while validation.jsonl is required only when --eval_iters > 0. Keep the existing record format and tokenization guidance unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/megatron_bridge/distill.py`:
- Around line 341-350: Update the SFT validation near the existing
student_provider.vocab_size check to compare the actual tokenizer mappings used
by args.student_hf_path and the teacher, including token IDs, special-token IDs,
and added-token IDs. Reject distillation with a clear ValueError when mappings
are not equivalent, while preserving training only when both tokenizers
interpret every batch ID identically.
---
Outside diff comments:
In `@examples/megatron_bridge/distill.py`:
- Around line 141-145: Update the SFT dataset documentation near the CLI help
for the training/validation JSONL files, plus the corresponding validation/error
text and README, to state that training.jsonl is always required while
validation.jsonl is required only when --eval_iters > 0. Keep the existing
record format and tokenization guidance unchanged.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 93033df1-1caa-48a1-9b2f-ad8ba8609c92
📒 Files selected for processing (1)
examples/megatron_bridge/distill.py
| if args.sft and student_provider.vocab_size != teacher_provider.vocab_size: | ||
| # The pretraining path is structurally immune to this: NullTokenizer plus pre-tokenized | ||
| # --data_paths means one tokenization feeds both models. SFT tokenizes raw text with the | ||
| # student's tokenizer, so a teacher from another family would score ids it never saw and | ||
| # silently produce a garbage KD target instead of an error. | ||
| raise ValueError( | ||
| "--sft tokenizes with the student's tokenizer, so student and teacher must share a " | ||
| f"vocabulary (got {student_provider.vocab_size} vs {teacher_provider.vocab_size})." | ||
| ) | ||
|
|
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 8 --type py \
'convert_to_distillation_provider|HuggingFaceTokenizer|input_ids|student_provider|teacher_provider' \
examples/megatron_bridge testsRepository: NVIDIA/Model-Optimizer
Length of output: 50379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- distill.py: argument and provider setup ---'
sed -n '80,180p' examples/megatron_bridge/distill.py
sed -n '280,410p' examples/megatron_bridge/distill.py
sed -n '450,515p' examples/megatron_bridge/distill.py
printf '%s\n' '--- distillation provider implementation ---'
sed -n '1,125p' examples/megatron_bridge/_distillation_provider.py
printf '%s\n' '--- focused tokenizer and distillation references ---'
rg -n -C 5 --type py \
'TokenizerConfig|tokenizer_model|teacher_provider|student_provider|DistillationProvider|kd_loss|teacher' \
examples/megatron_bridge \
-g 'distill.py' \
-g '_distillation_provider.py'Repository: NVIDIA/Model-Optimizer
Length of output: 39297
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
distill = Path("examples/megatron_bridge/distill.py").read_text()
provider = Path("examples/megatron_bridge/_distillation_provider.py").read_text()
checks = {
"sft_uses_student_tokenizer": (
'tokenizer_model=args.student_hf_path' in distill
and '"HuggingFaceTokenizer"' in distill
),
"guard_checks_only_vocab_sizes": (
'if args.sft and student_provider.vocab_size != teacher_provider.vocab_size:' in distill
and 'student_provider.vocab_size' in distill
and 'teacher_provider.vocab_size' in distill
),
"teacher_model_enters_kd_wrapper": (
'teacher = unwrap_model(' in provider
and '"teacher_model": teacher' in provider
and 'mtd.convert(student, mode=[("kd_loss", modelopt_cfg)])' in provider
),
"no_teacher_tokenizer_is_configured": (
'tokenizer_model=args.teacher_hf_path' not in distill
),
}
for name, passed in checks.items():
print(f"{name}: {'PASS' if passed else 'FAIL'}")
if not all(checks.values()):
raise SystemExit(1)
PY
printf '%s\n' '--- distillation README tokenizer/model constraints ---'
rg -n -C 4 --type md \
'distill|student_hf_path|teacher_hf_path|tokenizer|vocab' \
examples/megatron_bridge README.mdRepository: NVIDIA/Model-Optimizer
Length of output: 50380
Validate tokenizer ID compatibility for SFT distillation.
SFT uses args.student_hf_path as the only tokenizer, and the KD wrapper supplies the same batch IDs to the teacher. Equal vocab_size values do not guarantee matching token, special-token, or added-token IDs. Require equivalent tokenizer mappings before training.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/megatron_bridge/distill.py` around lines 341 - 350, Update the SFT
validation near the existing student_provider.vocab_size check to compare the
actual tokenizer mappings used by args.student_hf_path and the teacher,
including token IDs, special-token IDs, and added-token IDs. Reject distillation
with a clear ValueError when mappings are not equivalent, while preserving
training only when both tokenizers interpret every batch ID identically.
1 participant
What does this PR do ?
Type of change: New feature
Overview: Adds SFT-masked data support to the Megatron-Bridge distillation example, so a
model can be distilled on prompt/response pairs with the loss masked to the response.
Today
examples/megatron_bridge/distill.pyonly consumes pretraining-style data —GPTDatasetover pre-tokenized blends with
NullTokenizer— so the loss is computed over every token. Whendistilling an instruction-tuned model it is usually preferable to train on prompt/response pairs
and mask the loss to the response, matching how the model was fine-tuned.
Usage
where
/path/to/dataholdstraining.jsonl/validation.jsonlof records:{"input": "<prompt>", "output": "<response>"}How it works
Switches the data path to Bridge's
FinetuningDatasetConfig(NeMo-styleGPTSFTDataset):prompt_template="{input}{output}"tokenizes input+output verbatim — adjacent placeholders,no separator — so the text is fed exactly as provided
label_key="output"withanswer_only_loss=Truemasks the loss to the response(
answer_start_idx == len(context_ids))truncation_field="input"truncates the context when a pair exceedsseq_lengthTwo supporting changes, both scoped to
--sft:pretraining path consumes pre-tokenized data and keeps
NullTokenizer.context-parallel ranks, so
calculate_per_token_lossis enabled andaverage_in_collectiveis disabled. Both are untouched on the pretraining path.
Testing
Used for quantization-aware distillation of Nemotron-Nano-3 (W4A16 NVFP4) at
seq_length=32768with CP>1: 200 iterations, logits-distillation loss
3.37e-2 -> 1.91e-2monotonically, routerseq_load_balancing_losssteady, and the resulting checkpoint exports and serves correctly.Opt-in: without
--sftthe existing mock/blend data path is unchanged.Before your PR is "Ready for review"
--sft.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Documentation