feat(megatron-bridge): SFT-masked data support in distillation

[yueshen2016]

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.py only consumes pretraining-style data — GPTDataset
over pre-tokenized blends with NullTokenizer — so the loss is computed over every token. When
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.

Usage

python examples/megatron_bridge/distill.py \
  --teacher_hf_path <teacher> --student_hf_path <student> \
  --sft --sft_dataset_root /path/to/data \
  ...

where /path/to/data holds training.jsonl / validation.jsonl of records:

{"input": "<prompt>", "output": "<response>"}

How it works

Switches the data path to Bridge's FinetuningDatasetConfig (NeMo-style GPTSFTDataset):

• prompt_template="{input}{output}" tokenizes input+output verbatim — adjacent placeholders,
  no separator — so the text is fed exactly as provided
• label_key="output" with answer_only_loss=True masks the loss to the response
  (answer_start_idx == len(context_ids))
• truncation_field="input" truncates the context when a pair exceeds seq_length

Two supporting changes, both scoped to --sft:

• Tokenizer. SFT reads raw text, so it uses the model's real HuggingFace tokenizer. The
  pretraining path consumes pre-tokenized data and keeps NullTokenizer.
• Loss reduction. A response-only mask requires per-token loss to combine correctly across
  context-parallel ranks, so calculate_per_token_loss is enabled and average_in_collective
  is disabled. Both are untouched on the pretraining path.

Testing

Used for quantization-aware distillation of Nemotron-Nano-3 (W4A16 NVFP4) at seq_length=32768
with CP>1: 200 iterations, logits-distillation loss 3.37e-2 -> 1.91e-2 monotonically, router
seq_load_balancing_loss steady, and the resulting checkpoint exports and serves correctly.

Opt-in: without --sft the existing mock/blend data path is unchanged.

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes — purely additive and opt-in behind --sft.
• Did you write any new necessary tests?: No
• Did you add or update any necessary documentation?: No
• Did you update Changelog?: No

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Added supervised fine-tuning (SFT) support for distillation workflows.
  • Added configuration and validation for SFT dataset locations and supported inputs.
  • Added raw prompt/completion JSONL datasets with response-only loss masking.
  • Added truncation, end-of-sequence handling, and student-tokenizer support without automatic chat templates or BOS tokens.
  • Added matching student and teacher vocabulary validation.
  • Preserved existing mock and GPT dataset modes for non-SFT runs.

• Documentation

  • Documented required filenames, record format, tokenizer behavior, and completion-only loss masking.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

The distillation example adds SFT options, validates 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.

Changes

SFT distillation flow

Layer / File(s)	Summary
SFT options and validation examples/megatron_bridge/distill.py, examples/megatron_bridge/README.md	The CLI adds --sft and --sft_dataset_root. SFT requires the dataset root, rejects data paths and mock data, checks BOS handling, and documents the required JSONL files and record fields.
SFT data and loss configuration examples/megatron_bridge/distill.py	SFT uses FinetuningDatasetConfig for prompt/completion records, response-only masking, truncation, EOS insertion, tokenizer-vocabulary validation, per-token loss, and HuggingFace tokenization. Non-SFT paths retain their existing dataset and tokenizer behavior.
SFT distributed gradient reduction examples/megatron_bridge/distill.py	SFT disables collective loss pre-averaging during distributed gradient reduction. Other modes retain loss averaging.

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 and completion records
  HuggingFaceTokenizer->>ModelProviders: provide tokenized inputs with per-token loss
  ModelProviders->>DistributedGradientReduction: reduce SFT gradients without pre-averaging

Loading

Suggested reviewers: chenhanyu

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: adding SFT-masked data support to Megatron-Bridge distillation.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	The PR adds no torch.load, allow_pickle, eval/exec, # nosec, or dependency changes; trust_remote_code is passed from a store_true flag that defaults to False.

✨ Finishing Touches 💡 1

📝 Generate docstrings 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/mbridge-sft-distillation

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

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.

👉 Steps to fix this

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

📥 Commits

Reviewing files that changed from the base of the PR and between 43e9d15 and 8abcdb4.

📒 Files selected for processing (1)

• examples/megatron_bridge/distill.py

[yueshen2016]

/claude review

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.74%. Comparing base (b96841d) to head (8fe9543).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2113      +/-   ##
==========================================
- Coverage   78.76%   78.74%   -0.03%     
==========================================
  Files         522      522              
  Lines       60461    60461              
==========================================
- Hits        47621    47607      -14     
- Misses      12840    12854      +14     

Flag	Coverage Δ
examples-megatron_bridge	25.74% <ø> (-0.07%)	⬇️
unit	55.30% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[claude]

[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:

1. 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,

2. Keep add_bos=False (verbatim is a defensible contract) but state the requirement where users will read it — in the --sft_dataset_root help string and the README: input must contain the fully templated prompt, including any BOS and role/turn markers the model expects; no chat template or BOS is added.

[yueshen2016]

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_root help text,
• a comment next to the dataset_kwargs that 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.

[yueshen2016]

Following up on this thread: I took option 2 (document the contract), and after digging further I think option 1 as written would have been actively wrong for this model.

bos_token is not None tests whether a BOS exists in the vocab, not whether the tokenizer uses one. Nemotron-3.5's config is:

bos_token     : <s>        # exists
add_bos_token : False      # but is never prepended

and its chat_template.jinja never emits one either — the first thing it produces is <|im_start|>system. So option 1 would have set add_bos=True and trained the model with a BOS it never sees at serving, creating exactly the skew the suggestion was meant to prevent.

Even the corrected predicate (tokenizer.add_bos_token) conflicts with the contract, because a fully-templated input may already contain a BOS and would then get a second one. The two designs are mutually exclusive: either the caller owns the whole prompt (verbatim, what this PR does) or the framework does (apply_chat_template).

What your comment did expose is a real gap the documentation alone does not close: a model with add_bos_token=True whose data lacks a BOS is silently skewed. That case is now detected (6160a69):

def _warn_if_bos_missing(tokenizer, dataset_root: str) -> None:
    bos = getattr(tokenizer, "bos_token", None)
    if not bos or not getattr(tokenizer, "add_bos_token", False):
        return
    ...
    if not first_input.startswith(bos):
        warn_rank_0(f"This tokenizer prepends {bos!r} at inference, but the first record ... ")

Verified across five cases: silent for add_bos_token=False, silent when no BOS token exists, silent when the caller already baked the BOS in, fires when the tokenizer adds one and the data lacks it, and silent on a missing/malformed file (left to the dataset builder to report).

For data that is not pre-templated at all, the right answer is Bridge's GPTSFTChatDataset with use_hf_tokenizer_chat_template=True, which handles BOS, role markers and multi-turn together. Tracking that as a follow-up rather than expanding this PR.

[claude]

[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.

[yueshen2016]

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.

[claude]

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=True paired with average_in_collective=False is 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 tokens guard already in LogitsAndIntermediatesLossBalancer.forward.
• Opt-in isolation. Every behavioral change is gated on args.sft; the mock and blend paths, NullTokenizer, and average_in_collective=True are byte-for-byte unchanged without the flag. No modelopt_state schema, 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=True does mask to the response (answer_start_idx == len(context_ids)), and truncation_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.

[coderabbitai]

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.

👉 Steps to fix this

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

📥 Commits

Reviewing files that changed from the base of the PR and between 8abcdb4 and 080f4de.

📒 Files selected for processing (2)

• examples/megatron_bridge/README.md
• examples/megatron_bridge/distill.py

🚧 Files skipped from review as they are similar to previous changes (1)

• examples/megatron_bridge/distill.py

[coderabbitai]

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.

👉 Steps to fix this

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 win

Make validation.jsonl conditional in the SFT documentation.

When --eval_iters 0, do_validation is false, so only training.jsonl is required. Update the CLI help, error text, and README to state that validation.jsonl is 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

📥 Commits

Reviewing files that changed from the base of the PR and between 080f4de and 8c19f38.

📒 Files selected for processing (1)

• examples/megatron_bridge/distill.py

[yueshen2016]

/claude review

[coderabbitai]

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.

👉 Steps to fix this

Actionable comments posted: 2

🧹 Nitpick comments (1)

examples/megatron_bridge/distill.py (1)

329-329: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Keep AutoTokenizer in the module import block.

Transformers is already imported for the non-SFT AutoConfig call. Move AutoTokenizer to the module imports and remove this local import.

As per coding guidelines: “Keep imports at the top of Python source and test files; use local imports only for justified circular dependencies, optional dependencies, or unusually heavy imports, with a brief explanatory comment.”

🤖 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` at line 329, Move AutoTokenizer into the
module-level transformers import block and remove the local import near the
non-SFT distillation flow. Keep the existing AutoConfig import and behavior
unchanged; no local import is needed here.

Source: Coding guidelines

🤖 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`:
- Line 289: Update the argument logging around print_args(args) to avoid
emitting raw path-bearing values, including sft_dataset_root. Log a redacted
namespace containing only safe fields, or omit all path fields while preserving
non-sensitive argument logging.
- Around line 142-147: Update the CLI help text for the training data directory
argument in the argument parser to state that validation.jsonl is required only
when validation is enabled, such as when --eval_iters is greater than zero;
retain the existing descriptions of training.jsonl and record formatting.

---

Nitpick comments:
In `@examples/megatron_bridge/distill.py`:
- Line 329: Move AutoTokenizer into the module-level transformers import block
and remove the local import near the non-SFT distillation flow. Keep the
existing AutoConfig import and behavior unchanged; no local import is needed
here.

🪄 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: 96ad575c-df4d-43c5-87ea-4a431fdee3ad

📥 Commits

Reviewing files that changed from the base of the PR and between 48e9547 and 6160a69.

📒 Files selected for processing (1)

• examples/megatron_bridge/distill.py

[claude]

Claude review — feat(megatron-bridge): SFT-masked data support in distillation

Scope: full review (trigger was a bare /claude review). 2 files changed (+139/−5): examples/megatron_bridge/distill.py reviewed in full, README.md skimmed for the docs contract, plus surrounding context in _distillation_provider.py and modelopt/torch/distill/plugins/megatron.py.

This is a re-review — the six findings from the previous pass all landed:

Prior finding	Status
1 · CRITICAL — sanity check raised before the SFT branch	Fixed (distill.py:269, not args.sft and ...), plus a new mutual-exclusion guard at 283 so --data_paths/--use_mock_data cannot be silently discarded
2 · IMPORTANT — undocumented BOS / verbatim-text contract	Addressed — _warn_if_bos_missing, arg help, README paragraph. One residual issue, see below
3 · IMPORTANT — student/teacher tokenizer mismatch	Fixed well — get_vocab() mapping comparison at 335, not just size, with a message that distinguishes the two cases. Runs before the providers are built, so a mismatch costs seconds
4 · SUGGESTION — do_validation ignored --eval_iters 0	Fixed (477)
5 · SUGGESTION — split loss-reduction decision	Fixed — both halves now cross-reference each other (362-369, 519)
6 · SUGGESTION — no README entry	Fixed

Findings this round

CRITICAL: 0 · IMPORTANT: 1 · SUGGESTION: 1

#	Severity	Location	Issue
1	IMPORTANT Compatibility	distill.py:479-486	The "verbatim, no BOS" contract is assumed rather than enforced: add_bos=False suppresses the dataset prepend, not the tokenizer wrapper's add_special_tokens=True
2	SUGGESTION	distill.py:301-303	getattr(tokenizer, "add_bos_token", False) no-ops for tokenizers that prepend BOS via the post-processor — a behavioural probe is more robust

Most impactful

Finding 1 is the residue of prior finding 2, one layer down. The docs now make an unconditional promise ("no BOS token is prepended… include them in the input field yourself"), but the only thing backing it is add_bos=False in dataset_kwargs — and that flag governs the explicit prepend inside GPTSFTDataset._process_example, not what the Bridge HuggingFaceTokenizer wrapper does when it calls self._tokenizer(text).input_ids. If the wrapper tokenizes with specials on, then (a) a Llama-3/Mistral/Gemma user who follows the README gets a doubled BOS — with _warn_if_bos_missing staying quiet, because their input does start with the BOS string — and (b) because _process_example tokenizes each prompt_template segment independently, every record gets a second BOS injected at the response boundary, exactly where answer_only_loss begins scoring.

I could not verify the wrapper from this checkout (megatron-bridge is not installed in the review env), so I have flagged it as verify then enforce rather than as a confirmed defect. The check is a two-line print(tok("hello").input_ids); the fix, if needed, is add_bos_token=False in the hf_tokenizer_kwargs you already thread through. Either way the outcome belongs in that comment, since the next reader will otherwise have to re-derive it. Neither variant would have appeared in the Nemotron-Nano-3 run if that tokenizer does not auto-add BOS — the same blind spot the earlier findings sat in.

What I checked and found correct

• Loss reduction. calculate_per_token_loss=True + average_in_collective=False remains the right pairing for a response-only mask under CP>1: a pre-averaged collective weights ranks equally regardless of how many unmasked tokens each holds. Consistent with the # zero when one CP rank has only context tokens guard at plugins/megatron.py:507, and the flag survives convert_to_distillation_provider (which mutates the student provider in place). Setting it on the teacher too is genuinely harmless — _compute_teacher_lm_loss returns zeros_like(labels).
• KD loss is masked, not just the LM loss. LogitsKLLoss.post_forward reshapes [s, b] → [b, s] "for upcoming loss masking", so the response-only mask reaches the distillation term as well. With skip_lm_loss=True (the default) that is the only live term, so this is what makes --sft mean anything.
• Vocab guard. Comparing get_vocab() mappings rather than sizes is the correct strictness — same-size/different-mapping is exactly the case that produces a cleanly-training, subtly-wrong model.
• Opt-in isolation. Every behavioural change is gated on args.sft. NullTokenizer, average_in_collective=True, the mock/blend branches, and calculate_per_token_loss are byte-for-byte unchanged without the flag. No modelopt_state schema, mode registration, or public API is touched — no checkpoint-compat or restore-fidelity exposure.
• Arg validation. --sft without --sft_dataset_root errors; --sft with either pretraining data source errors rather than silently discarding it; --validate_only still forces eval_iters > 0, which now also implies do_validation=True, so the validation file is required exactly when it is used.

Risk

Low-medium, down from medium. The feature now runs as documented and the two silent-wrong-result paths from last round are guarded. What remains is one docs/behaviour question whose worst case is a train/inference BOS skew for model families outside the one that was tested — a config-level concern, not a code-structure one, in a fully opt-in example script that cannot regress any existing workflow.

Still worth noting (raised last round, unchanged): the checklist marks no tests, and tests/examples/megatron_bridge/test_distill.py exists. A --sft smoke case over a two-record JSONL fixture would have caught last round's CRITICAL outright, and would pin the arg-validation matrix that now has four branches.

[yueshen2016]

/claude review

[claude]

[SUGGESTION] The "verbatim" contract stated here (and in --sft_dataset_root help and the README) is slightly stronger than what GPTSFTDataset actually does: _process_example applies .strip(' ') to each prompt_template field value before tokenizing, so leading/trailing ASCII spaces of "input" and "output" are dropped. (Other whitespace such as \n survives, which is why most templates are unaffected.)

Why it matters: each template string is tokenized separately, so the first token of "output" is a BPE token boundary the loss starts scoring at. {"input": "... [/INST] ", "output": "..."} (Llama-2/Mistral style, trailing space significant) and {"input": "Q: ...\nA:", "output": " 42"} both lose that space, so the model is trained on "42" where the served format produces "▁42". That is the same class of train/inference skew the BOS warning above was added to catch, and it is invisible in the loss curve.

Suggest softening the claim in all three places to name the exception — e.g. "tokenized as written, except that leading/trailing spaces on each field are stripped; put any significant separator inside the text (a newline, not a space)". Worth a quick check of _process_example in your container's megatron.bridge.data.datasets.sft to confirm the .strip(' ') carried over from NeMo.

[yueshen2016]

Accepted — 3b9eb9d. Confirmed in the source: prompt_template_values.append(example[c].strip(" ")) (gpt_sft.py:513). Spaces go, newlines survive.

I had actually noticed the newline surviving when tracing this path earlier, but did not follow it to the consequence you draw: with each field tokenized separately, a stripped trailing space changes the first token of "output" from ▁42 to 42, and the loss starts scoring at that boundary. Silent, and invisible in the loss curve.

Softened in all three places, with the concrete recommendation:

tokenized as written, except that leading and trailing spaces on each field are stripped ... express any significant separator as a newline rather than a trailing space.

[claude]

[SUGGESTION] Two things worth verifying about this kwarg, because the entire verbatim guarantee rests on it and nothing here fails loudly if it doesn't land:

1. Is it consumed? hf_tokenizer_kwargs may be forwarded to AutoTokenizer.from_pretrained, which accepts (and silently retains) unknown kwargs rather than raising. If Bridge's HuggingFaceTokenizer doesn't intercept include_special_tokens as its own __init__ parameter, this becomes a no-op and a BOS-adding tokenizer does inject a BOS at the answer boundary — the exact bug the comment says is being prevented, still present but now believed fixed.
2. The stated default. The comment asserts the default is True. In the NeMo AutoTokenizer this class is ported from, include_special_tokens defaults to False and text_to_ids goes through text_to_tokens/tokens_to_ids (no special tokens added). If Bridge kept that default, the comment's rationale is inverted — harmless behaviorally, but it's now the primary documentation of a non-obvious constraint, so a reader can't trust it.

Given the failure mode is silent, an assertion beats a comment. You already load student_tokenizer at the top of main(); a two-line check there would pin the invariant regardless of what the Bridge layer does:

# The SFT dataset tokenizes "{input}" and "{output}" separately, so a BOS injected by
# text_to_ids would land at the answer boundary where answer_only_loss starts scoring.
if student_tokenizer("x").input_ids[:1] == [student_tokenizer.bos_token_id]:
    print_rank_0("Tokenizer adds BOS; relying on include_special_tokens=False to suppress it.")

At minimum, correct the stated default once you've checked it.

[yueshen2016]

Not changing this one — both concerns are already disproven by measurement, which I should have put in the comment rather than leaving as an assertion.

I probed the actual call chain in the container: _tokenize -> DefaultTokenizerText.tokenize -> megatron.core.tokenizers.text.libraries.huggingface_tokenizer.HuggingFaceTokenizer.text_to_ids, which is

if self.include_special_tokens:
    return self.tokenizer(text).input_ids
tokens = self.text_to_tokens(text)
return self.tokens_to_ids(tokens)

and then measured both of your questions directly:

default include_special_tokens = True
hf_tokenizer_kwargs include_special_tokens=False -> False
hf_tokenizer_kwargs include_special_tokens=True  -> True

So (1) the kwarg is consumed — passing it changes the value, it is not silently retained by from_pretrained; and (2) the default in this Bridge version is True, so the comment's rationale stands. The NeMo AutoTokenizer default of False you are reasoning from did not carry over to megatron.core.tokenizers.

Your underlying point is fair though — the guarantee rested on a comment for a silent failure mode. It now rests on a measurement, and the two probes are reproducible from the container if a future Bridge version changes the default.

[claude]

[SUGGESTION] Full dict equality is stricter than the invariant that actually matters, and there's no escape hatch when it over-triggers.

What KD needs is that every id the dataset can emit means the same thing to both models. A teacher whose get_vocab() is a strict superset of the student's satisfies that — e.g. a teacher fine-tune that added a few reserved/special tokens, or a student tokenizer directory that dropped unused added_tokens. get_vocab() includes added_tokens_encoder, so a single extra token on the teacher side fails this check and the run is unlaunchable, with an error message ("must share a tokenizer") that points at the wrong thing.

Consider testing agreement on the shared tokens plus student ⊆ teacher, and reserving the hard error for genuine mapping conflicts:

conflicts = {t for t, i in student_vocab.items() if teacher_vocab.get(t, i) != i}
missing = student_vocab.keys() - teacher_vocab.keys()
if conflicts or missing:
    raise ValueError(...)  # ids genuinely disagree
if len(teacher_vocab) != len(student_vocab):
    warn_rank_0(
        f"Teacher vocabulary has {len(teacher_vocab) - len(student_vocab)} extra token(s) not in "
        "the student's; ids agree on all shared tokens, so distillation is well-defined."
    )

Also minor: AutoTokenizer.from_pretrained runs twice here (plus a third time inside Bridge for the training tokenizer), on every rank. Not hot-path, but discarding the teacher tokenizer immediately after get_vocab() makes the cost easy to miss.

[yueshen2016]

Accepted — 3b9eb9d. You are right that full equality is stricter than the invariant. Only student ids ever reach the teacher (the dataset tokenizes with the student's tokenizer), so what matters is that the teacher maps every student token to the same id.

conflicts = {t for t, i in student_vocab.items() if teacher_vocab.get(t, i) != i}
missing = student_vocab.keys() - teacher_vocab.keys()
if conflicts or missing:
    raise ValueError(...)
if len(teacher_vocab) > len(student_vocab):
    warn_rank_0(...)

Checked against five cases before pushing: identical -> OK; teacher superset -> warn; id conflict at equal size -> raise; token absent from teacher -> raise; different family at equal size -> raise. So the Llama-2/Mistral case that motivated the guard still fails, and a teacher with extra reserved tokens no longer does.

On the double from_pretrained: agreed it is not free, but it is once per run at startup against a local path, and holding the teacher tokenizer alive only to avoid a second load would trade clarity for little. Left as is.

[claude]

[SUGGESTION] Function-local import without one of the sanctioned reasons. CONTRIBUTING.md ("Keep imports at the top of the file") allows in-function imports only for circular imports, optional-dependency guards, or an unusually heavy import with explicit justification — and transformers is already imported unconditionally at module top (from transformers import AutoConfig, line 53), so none of the three applies and there is no cost to deferring.

Suggested change

from transformers import AutoTokenizer

…and extend line 53 to from transformers import AutoConfig, AutoTokenizer.

[yueshen2016]

Accepted — 3b9eb9d. from transformers import AutoConfig is already at module scope, so none of the three sanctioned reasons applied. Extended to AutoConfig, AutoTokenizer and removed the deferred import.

[claude]

Claude review — feat(megatron-bridge): SFT-masked data support in distillation

Scope: full review (trigger comment was bare /claude review). 2 files changed (+159/−5); both reviewed in full — examples/megatron_bridge/distill.py and examples/megatron_bridge/README.md — plus surrounding context in _distillation_provider.py and modelopt/torch/distill/plugins/megatron.py.

Prior round: all six findings addressed

Prior finding	Status
1 CRITICAL — pretraining data check fired before the SFT branch	✅ not args.sft guard + explicit mutual-exclusion error for --data_paths/--use_mock_data
2 IMPORTANT — undocumented BOS / verbatim-text contract	✅ documented in help + README, plus a _warn_if_bos_missing probe
3 IMPORTANT — student/teacher tokenizer mismatch silently corrupts KD targets	✅ vocab-mapping comparison before providers are built
4 SUGGESTION — do_validation ignored --eval_iters 0	✅ do_validation=args.eval_iters > 0
5 SUGGESTION — calculate_per_token_loss / average_in_collective split with no cross-reference	✅ both comments now point at each other
6 SUGGESTION — no README entry for --sft	✅ added

Findings 2 and 3 were fixed above the bar I asked for — the vocab check compares mappings rather than sizes and names the Llama-2/Mistral collision, and the BOS probe reads the tokenizer's actual output instead of trusting add_bos_token, with its own bounded-scope caveat documented.

This round

CRITICAL: 0 · IMPORTANT: 1 · SUGGESTION: 4

#	Severity	Location	Issue
1	IMPORTANT Algorithm	distill.py:494	truncation_method unset → default 'right' cuts the tail of the prompt, destroying the prompt→response seam answer_only_loss scores from
2	SUGGESTION	distill.py:474-477	"Verbatim" overstates it: GPTSFTDataset._process_example applies .strip(' ') per field, so significant leading/trailing spaces are dropped
3	SUGGESTION	distill.py:552-557	include_special_tokens=False may be silently swallowed by from_pretrained, and the stated default of True is False in the NeMo class this is ported from
4	SUGGESTION	distill.py:347-356	Full vocab-dict equality rejects a legitimate teacher-superset vocabulary with no escape hatch
5	SUGGESTION	distill.py:341	Function-local from transformers import AutoTokenizer — transformers is already imported at line 53

Most impactful

Finding 1 is the one to act on. truncation_field="input" correctly protects the labels, but which end of the context gets cut is a separate decision that this config inherits rather than makes. In the NeMo GPTSFTDataset lineage, truncation_method defaults to 'right', so an over-length pair loses the tail of the prompt — the actual question, and any role/turn markers the caller baked in (which, with no chat template applied, is the only thing marking where the response begins). The example is then trained to start a response at a boundary that never occurs at serving time. It is silent, loss keeps falling, and it affects only records exceeding seq_length — at seq_length=32768 the validation run likely truncated nothing, so the reported curve cannot rule it out. One kwarg: "truncation_method": "left".

Findings 2 and 3 share a shape with the previous round's 2 and 3: the verbatim contract is now asserted in three places (code comment, --sft_dataset_root help, README) but still enforced only by a kwarg whose consumption is not observable from this file, and the contract has one real exception (.strip(' ')) that none of the three mention. Both are cheap to close, and worth closing precisely because the guarantee is now load-bearing documentation others will rely on.

Verified correct

• Per-token loss plumbing. provider.calculate_per_token_loss = True is set inside _build_model_provider, i.e. before convert_to_distillation_provider, so it reaches distill_provider (rebased from the student provider) and thus the TransformerConfig the pipeline schedules read — and DistillationProvider.__setattr__'s teacher-mirroring makes the "lands on both providers" note in the comment literally true. Paired with average_in_collective=False this is the right combination for a response-only mask under CP>1, and it lines up with the # zero when one CP rank has only context tokens guard already in LogitsAndIntermediatesLossBalancer.forward (megatron.py:507) — the losses entering that balancer are already mask-reduced scalars, so a CP rank holding only context contributes zero rather than a NaN scale factor.
• Dataset kwargs. prompt_template="{input}{output}" splits to exactly ['{input}', '{output}'] under _separate_template's ( *?{.+?}) regex (adjacent placeholders, no injected separator); label_key="output" is last, satisfying the label-position assertion; answer_only_loss=True masks from answer_start_idx == len(context_ids); and add_eos=True goes through tokenizer.eos_id, so it is unaffected by include_special_tokens=False and the EOS lands inside the scored span — the model still learns to stop.
• Opt-in isolation. Every behavioral change is gated on args.sft. The mock and blend branches, NullTokenizer, and average_in_collective=True are byte-for-byte unchanged without the flag. No modelopt_state schema, mode registration, or public API is touched, so there is no checkpoint-compat or restore-fidelity exposure.
• Argument validation. --sft requires --sft_dataset_root; conflicting data sources now error instead of being silently discarded; --validate_only / --eval_iters stay consistent with do_validation. warn_rank_0 → dist.is_master() is safe before distributed init.

Risk

Low-to-medium. One example script, fully opt-in, no existing workflow can regress — and the previous round's launch-blocking bug is fixed. The residual exposure is entirely in the "trains cleanly but on subtly wrong data" category, concentrated in over-length records (finding 1) and space-sensitive prompt formats (finding 2), neither of which the reported Nemotron-Nano-3 run would have exercised. Finding 1 is a one-line change.

Still worth noting: the checklist marks no tests. A --sft smoke case in tests/examples/megatron_bridge/test_distill.py over a two-record JSONL fixture would pin the argument wiring and the tokenizer/dataset contract cheaply — and would have caught last round's finding 1 outright.

[yueshen2016]

/claude review

[ChenhanYu]

So this is a single turn SFT support right? I recalled that MBridge should have multi-turn OpenAI message style support with answer_only masking.

[ChenhanYu]

Megatron Bridge current SFT support appears to handle OpenAI/HF-style multi-turn messages / conversation rows, with ChatSFTPreprocessingConfig(loss_mode="assistant") masking every assistant turn. This PR instead restricts the input to a single {input, output} pair and uses the legacy FinetuningDatasetConfig path.

Is there a reason we are not supporting the multi-turn OAI message schema here, with answer-only masking over all assistant responses? If the older materialized JSONL path is a deliberate constraint, could we document that and/or consider the direct-HF chat SFT configuration so distillation matches the dataset formats Bridge already supports?

[kevalmorabia97]

LGTM!

[yueshen2016]

Megatron Bridge current SFT support appears to handle OpenAI/HF-style multi-turn messages / conversation rows, with ChatSFTPreprocessingConfig(loss_mode="assistant") masking every assistant turn. This PR instead restricts the input to a single {input, output} pair and uses the legacy FinetuningDatasetConfig path.

Is there a reason we are not supporting the multi-turn OAI message schema here, with answer-only masking over all assistant responses? If the older materialized JSONL path is a deliberate constraint, could we document that and/or consider the direct-HF chat SFT configuration so distillation matches the dataset formats Bridge already supports?

This current path is targeting the private data that uses the legacy FinetuningDatasetConfig path. As for the the new HF data-config API (DirectHFSFTDatasetConfig /ChatSFTPreprocessingConfig / HFDatasetSourceConfig) that you referred to, I am also ready to push a new a PR supporting that as that could be used for public HF dataset. The reason it's not in this PR is that the new HF data-config API is not supported in 26.06 container, they should be supported in 26.08 which will be release tmr. And by then my new PR will be out for review.

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-08-13 01:10 UTC