Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yueshen2016 wants to merge 3 commits into
base: main
Choose a base branch
from
+97 −5
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8abcdb4
feat(megatron-bridge): SFT-masked data support in distillation
yueshen2016 Aug 7, 2026
080f4de
fix(megatron-bridge): let --sft run without pretraining --data_paths;…
yueshen2016 Aug 7, 2026
8c19f38
fix(megatron-bridge): address review on the SFT distillation path
yueshen2016 Aug 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions examples/megatron_bridge/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,13 @@ The distillation script expects pre-tokenized data in Megatron's binary format (
See the **[Dataset Preparation README](../dataset/README.md#tokenizing-for-megatron-frameworks)**
for full instructions on tokenizing JSONL files and Hugging Face datasets and get the list of output prefixes that you can use for `--data_paths` argument.

Alternatively, pass `--sft --sft_dataset_root <dir>` to distill on **raw prompt-completion JSONL**
with the loss masked to the completion. The directory must hold `training.jsonl` and
`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.
Comment on lines +135 to +138

@coderabbitai coderabbitai Bot Aug 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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


### Distillation with Real Data

Example usage to distill a 4B student (HF) from an 8B teacher (HF) on 8 GPUs (TP=8, PP=1):
Expand Down
95 changes: 90 additions & 5 deletions examples/megatron_bridge/distill.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
from megatron.bridge.training.config import (
CheckpointConfig,
ConfigContainer,
FinetuningDatasetConfig,
GPTDatasetConfig,
LoggerConfig,
MockGPTDatasetConfig,
Expand Down Expand Up @@ -125,6 +126,23 @@ def get_args():
parser.add_argument(
"--use_mock_data", action="store_true", help="Use mock data instead of --data_paths"
)
parser.add_argument(
"--sft",
action="store_true",
help="SFT-masked distillation: read raw prompt-completion jsonl from --sft_dataset_root "
"and mask the loss to the completion (assistant response) tokens. Uses "
"FinetuningDatasetConfig and the real HuggingFace tokenizer instead of the pretraining "
"GPTDataset and NullTokenizer.",
)
parser.add_argument(
"--sft_dataset_root",
type=str,
default=None,
help="Directory holding training.jsonl / validation.jsonl of "
'{"input": <prompt>, "output": <response>} records (used with --sft). Both fields are '
"tokenized verbatim: no chat template is applied and no BOS is prepended, so if the model "
"expects role/turn markers or a BOS token, bake them into the fields yourself.",
)
# Training & Eval arguments
parser.add_argument(
"--output_dir", type=str, required=True, help="Folder for logging and checkpoint saving"
Expand Down Expand Up @@ -246,7 +264,7 @@ def get_args():
args = parser.parse_args()

# Sanity checks
if not args.use_mock_data and not args.data_paths:
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.student_hf_model is None:
Expand All @@ -256,6 +274,16 @@ def get_args():
if args.validate_only and args.eval_iters == 0:
raise ValueError("--validate_only requires --eval_iters > 0.")

if args.sft and not args.sft_dataset_root:
raise ValueError(
"--sft requires --sft_dataset_root (a directory with training.jsonl / validation.jsonl)."
)
Comment on lines +277 to +280

@claude claude Bot Aug 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@yueshen2016 yueshen2016 Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

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."
)

Comment thread
coderabbitai[bot] marked this conversation as resolved.
print_args(args)

return args
Expand All @@ -279,6 +307,14 @@ def _build_model_provider(hf_path, load_weights=True):
provider.expert_model_parallel_size = args.ep_size
provider.expert_tensor_parallel_size = 1 # Expert tensor parallelism is not supported
provider.seq_length = args.seq_length
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. This lands on
# both providers (harmless: the teacher's LM loss is zeroed in
# adjust_distillation_model_for_mcore) and must stay in sync with
# ``average_in_collective=not args.sft`` on the shared DistributedDataParallelConfig
# below -- a per-token loss must not be pre-averaged.
provider.calculate_per_token_loss = True
Comment on lines +310 to +317

@claude claude Bot Aug 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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 yueshen2016 Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

if args.recompute_granularity is not None:
provider.recompute_granularity = args.recompute_granularity
provider.recompute_method = args.recompute_method
Expand All @@ -302,6 +338,16 @@ def _build_model_provider(hf_path, load_weights=True):
student_provider.gradient_accumulation_fusion = False
teacher_provider = _build_model_provider(args.teacher_hf_path)

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})."
)

Comment on lines +341 to +350

@coderabbitai coderabbitai Bot Aug 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ 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 tests

Repository: 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.md

Repository: 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.

kd_config = ModelOptDistillConfig(
skip_lm_loss=not args.no_skip_lm_loss, kd_loss_scale=args.kd_loss_scale
)
Expand Down Expand Up @@ -368,7 +414,36 @@ def _restore_student_hook(model_chunks):
"dataloader_type": "single",
"skip_getting_attention_mask_from_dataset": True,
}
if args.use_mock_data:
if args.sft:
# SFT-masked distillation via Bridge's FinetuningDatasetConfig -> NeMo-style GPTSFTDataset.
# `dataset_root` holds training.jsonl / validation.jsonl of {"input", "output"} records.
# 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 only (answer_start_idx == len(context_ids)); truncation_field="input" truncates
# the context when the pair exceeds seq_length.
#
# add_bos=False plus the placeholder-only prompt_template means the records are tokenized
# exactly as written -- no chat template, no BOS, no role markers. Callers whose model
# expects those must bake them into the "input" field; see --sft_dataset_root help.
dataset_config = FinetuningDatasetConfig(
seq_length=args.seq_length,
dataset_root=args.sft_dataset_root,
seed=args.seed,
dataloader_type="batch",
# Honour --eval_iters 0 so a training-only dataset_root does not have to carry a
# dummy validation.jsonl just to satisfy the builder.
do_validation=args.eval_iters > 0,
do_test=False,
Comment on lines +428 to +436

@claude claude Bot Aug 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Suggested change
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,

@yueshen2016 yueshen2016 Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

dataset_kwargs={
"prompt_template": "{input}{output}",
"label_key": "output",
"truncation_field": "input",
"answer_only_loss": True,
"add_bos": False,
"add_eos": True,
},
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Comment on lines +437 to +444

@claude claude Bot Aug 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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 yueshen2016 Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

)
elif args.use_mock_data:
dataset_config = MockGPTDatasetConfig(**dataset_kwargs)
else:
# Convert flat CLI list (e.g. ["1.0", "/path/data"]) to Megatron blend format
Expand Down Expand Up @@ -399,7 +474,7 @@ def _restore_student_hook(model_chunks):
grad_reduce_in_fp32=True,
overlap_grad_reduce=True,
overlap_param_gather=True,
average_in_collective=True,
average_in_collective=not args.sft, # per-token loss must not be pre-averaged
use_distributed_optimizer=True,
),
dataset=dataset_config,
Expand All @@ -412,8 +487,18 @@ def _restore_student_hook(model_chunks):
wandb_entity=args.wandb_entity, # optional
wandb_exp_name=args.wandb_exp_name,
),
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},
)
Comment on lines 489 to +497

@claude claude Bot Aug 7, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@yueshen2016 yueshen2016 Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

if args.sft
else TokenizerConfig(
tokenizer_type="NullTokenizer", vocab_size=distill_provider.vocab_size
)
),
checkpoint=CheckpointConfig(
save_interval=(
Expand Down