From 14ccdf0783a4fbbda88220a3e193198d4f87bb63 Mon Sep 17 00:00:00 2001 From: James Shen Date: Fri, 7 Aug 2026 08:31:38 -0700 Subject: [PATCH 01/11] feat(megatron-bridge): SFT-masked data support in distillation 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": , "output": } 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 --- examples/megatron_bridge/distill.py | 66 +++++++++++++++++++++++++++-- 1 file changed, 62 insertions(+), 4 deletions(-) diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index b35369c40f3..6c6dbb5fbc1 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -34,6 +34,7 @@ from megatron.bridge.training.config import ( CheckpointConfig, ConfigContainer, + FinetuningDatasetConfig, GPTDatasetConfig, LoggerConfig, MockGPTDatasetConfig, @@ -125,6 +126,21 @@ 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": , "output": } records (used with --sft).', + ) # Training & Eval arguments parser.add_argument( "--output_dir", type=str, required=True, help="Folder for logging and checkpoint saving" @@ -256,6 +272,11 @@ 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)." + ) + print_args(args) return args @@ -279,6 +300,10 @@ 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. + provider.calculate_per_token_loss = True if args.recompute_granularity is not None: provider.recompute_granularity = args.recompute_granularity provider.recompute_method = args.recompute_method @@ -368,7 +393,30 @@ 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. + 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_kwargs={ + "prompt_template": "{input}{output}", + "label_key": "output", + "truncation_field": "input", + "answer_only_loss": True, + "add_bos": False, + "add_eos": True, + }, + ) + 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 @@ -399,7 +447,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, @@ -412,8 +460,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}, + ) + if args.sft + else TokenizerConfig( + tokenizer_type="NullTokenizer", vocab_size=distill_provider.vocab_size + ) ), checkpoint=CheckpointConfig( save_interval=( From fb1c84d5c7713ab2c5084c721c2e9fe24bea439f Mon Sep 17 00:00:00 2001 From: James Shen Date: Fri, 7 Aug 2026 09:04:58 -0700 Subject: [PATCH 02/11] fix(megatron-bridge): let --sft run without pretraining --data_paths; 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 " 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 --- examples/megatron_bridge/README.md | 7 +++++++ examples/megatron_bridge/distill.py | 10 ++++++++-- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 030edaf8fd5..ccf60ca7b78 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -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 ` 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": , "output": }` 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. + ### Distillation with Real Data Example usage to distill a 4B student (HF) from an 8B teacher (HF) on 8 GPUs (TP=8, PP=1): diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 6c6dbb5fbc1..d39f764bf66 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -139,7 +139,9 @@ def get_args(): type=str, default=None, help="Directory holding training.jsonl / validation.jsonl of " - '{"input": , "output": } records (used with --sft).', + '{"input": , "output": } 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( @@ -262,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: @@ -400,6 +402,10 @@ def _restore_student_hook(model_chunks): # 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, From cf733d542d54afaca287465a9dac3b4d27e4be3e Mon Sep 17 00:00:00 2001 From: James Shen Date: Fri, 7 Aug 2026 09:21:52 -0700 Subject: [PATCH 03/11] fix(megatron-bridge): address review on the SFT distillation path - 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 --- examples/megatron_bridge/distill.py | 25 +++++++++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index d39f764bf66..c94f6388427 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -278,6 +278,11 @@ def get_args(): raise ValueError( "--sft requires --sft_dataset_root (a directory with training.jsonl / validation.jsonl)." ) + 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." + ) print_args(args) @@ -304,7 +309,11 @@ def _build_model_provider(hf_path, load_weights=True): 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. + # 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 if args.recompute_granularity is not None: provider.recompute_granularity = args.recompute_granularity @@ -329,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})." + ) + kd_config = ModelOptDistillConfig( skip_lm_loss=not args.no_skip_lm_loss, kd_loss_scale=args.kd_loss_scale ) @@ -411,7 +430,9 @@ def _restore_student_hook(model_chunks): dataset_root=args.sft_dataset_root, seed=args.seed, dataloader_type="batch", - do_validation=True, + # 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, dataset_kwargs={ "prompt_template": "{input}{output}", From f015f01060ce390cb1c6801b6389e09492fffbd3 Mon Sep 17 00:00:00 2001 From: James Shen Date: Mon, 10 Aug 2026 22:53:39 -0700 Subject: [PATCH 04/11] fix(megatron-bridge): compare tokenizer mappings, not vocab sizes, and note the appended EOS MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `--sft` guard compared `student_provider.vocab_size` against the teacher's, but equal sizes do not mean equal token->id mappings: Llama-2 and Mistral are both 32000 tokens with different vocabularies, so the check passed while the teacher scored ids it never saw. Compare `get_vocab()` instead. Moved it ahead of provider construction so a mismatch fails in seconds rather than after both models are built. Also note in the help text and README that an EOS token is appended after the response — "tokenized verbatim" was true of the two fields but did not mention `add_eos=True`. Signed-off-by: James Shen --- examples/megatron_bridge/README.md | 2 +- examples/megatron_bridge/distill.py | 36 ++++++++++++++++++++--------- 2 files changed, 26 insertions(+), 12 deletions(-) diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index ccf60ca7b78..56f7aa276b7 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -135,7 +135,7 @@ with the loss masked to the completion. The directory must hold `training.jsonl` `validation.jsonl` of `{"input": , "output": }` 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. +include them in the `"input"` field yourself. An EOS token is appended after the response. ### Distillation with Real Data diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index c94f6388427..8f6194ddb9d 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -141,7 +141,8 @@ def get_args(): help="Directory holding training.jsonl / validation.jsonl of " '{"input": , "output": } 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.", + "expects role/turn markers or a BOS token, bake them into the fields yourself. An EOS " + "token is appended after the response.", ) # Training & Eval arguments parser.add_argument( @@ -293,6 +294,29 @@ def main(args: argparse.Namespace): checkpoint_dir = os.path.join(args.output_dir, "checkpoints") tensorboard_dir = os.path.join(args.output_dir, "tb_logs") + if args.sft: + # SFT tokenizes raw text with the student's tokenizer and scores the teacher on those same + # ids, so both models must agree on what every id means. Equal vocabulary *sizes* are not + # enough -- Llama-2 and Mistral are both 32000 tokens with different mappings -- so compare + # the mappings. Checked before the providers are built so a mismatch costs seconds. + # The pretraining path is structurally immune: NullTokenizer plus pre-tokenized + # --data_paths means one tokenization feeds both models. + from transformers import AutoTokenizer + + _tok = {"trust_remote_code": args.trust_remote_code} + student_vocab = AutoTokenizer.from_pretrained(args.student_hf_path, **_tok).get_vocab() + teacher_vocab = AutoTokenizer.from_pretrained(args.teacher_hf_path, **_tok).get_vocab() + if student_vocab != teacher_vocab: + detail = ( + f"{len(student_vocab)} vs {len(teacher_vocab)} tokens" + if len(student_vocab) != len(teacher_vocab) + else f"both {len(student_vocab)} tokens, but different token->id mappings" + ) + raise ValueError( + "--sft tokenizes with the student's tokenizer and scores the teacher on those " + f"same ids, so student and teacher must share a tokenizer ({detail})." + ) + # Build student and teacher model providers def _build_model_provider(hf_path, load_weights=True): bridge = AutoBridge.from_hf_pretrained(hf_path, trust_remote_code=args.trust_remote_code) @@ -338,16 +362,6 @@ 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})." - ) - kd_config = ModelOptDistillConfig( skip_lm_loss=not args.no_skip_lm_loss, kd_loss_scale=args.kd_loss_scale ) From 805706087ca9ec3805bc659441620aeca9446379 Mon Sep 17 00:00:00 2001 From: James Shen Date: Tue, 11 Aug 2026 00:08:14 -0700 Subject: [PATCH 05/11] feat(megatron-bridge): warn when --sft data lacks a BOS the tokenizer adds at inference MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `--sft` tokenizes both fields verbatim (`add_bos=False`), so the caller owns the BOS. We cannot add one for them — their text may already contain it, and prepending would double it — but a model whose tokenizer sets `add_bos_token=True` and whose data does not carry a BOS is trained without the token it is served with, which is silent train/inference skew. The mismatch is now detected and reported: if the tokenizer prepends a BOS and the first training record does not start with it, warn and name the fix. A missing or malformed file is left to the dataset builder to report. Signed-off-by: James Shen --- examples/megatron_bridge/distill.py | 30 ++++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 8f6194ddb9d..6163f0c1c1d 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -22,6 +22,7 @@ import argparse import contextlib +import json import os import torch @@ -290,6 +291,30 @@ def get_args(): return args +def _warn_if_bos_missing(tokenizer, dataset_root: str) -> None: + """Warn when the tokenizer prepends a BOS at inference but the SFT records do not carry one. + + ``--sft`` tokenizes both fields verbatim, so the caller owns the BOS. Adding one here is not + safe (their text may already have it, which would double it), but training without the BOS the + model is served with is silent train/inference skew. + """ + bos = getattr(tokenizer, "bos_token", None) + if not bos or not getattr(tokenizer, "add_bos_token", False): + return + try: + with open(os.path.join(dataset_root, "training.jsonl")) as f: + first_input = str(json.loads(f.readline()).get("input", "")) + except Exception: + return # a malformed or missing file is the dataset builder's error to report, not ours + if not first_input.startswith(bos): + warn_rank_0( + f"This tokenizer prepends {bos!r} at inference, but the first record in " + f"{dataset_root}/training.jsonl does not start with it. --sft tokenizes the fields " + f"verbatim, so the model would be trained without the BOS it is served with. Include " + f"{bos!r} at the start of the 'input' field." + ) + + def main(args: argparse.Namespace): checkpoint_dir = os.path.join(args.output_dir, "checkpoints") tensorboard_dir = os.path.join(args.output_dir, "tb_logs") @@ -304,7 +329,8 @@ def main(args: argparse.Namespace): from transformers import AutoTokenizer _tok = {"trust_remote_code": args.trust_remote_code} - student_vocab = AutoTokenizer.from_pretrained(args.student_hf_path, **_tok).get_vocab() + student_tokenizer = AutoTokenizer.from_pretrained(args.student_hf_path, **_tok) + student_vocab = student_tokenizer.get_vocab() teacher_vocab = AutoTokenizer.from_pretrained(args.teacher_hf_path, **_tok).get_vocab() if student_vocab != teacher_vocab: detail = ( @@ -317,6 +343,8 @@ def main(args: argparse.Namespace): f"same ids, so student and teacher must share a tokenizer ({detail})." ) + _warn_if_bos_missing(student_tokenizer, args.sft_dataset_root) + # Build student and teacher model providers def _build_model_provider(hf_path, load_weights=True): bridge = AutoBridge.from_hf_pretrained(hf_path, trust_remote_code=args.trust_remote_code) From 44a2f42bc2ee4a57b289da8da3ceb3a1ee081da9 Mon Sep 17 00:00:00 2001 From: James Shen Date: Tue, 11 Aug 2026 00:34:46 -0700 Subject: [PATCH 06/11] fix(megatron-bridge): enforce the verbatim SFT contract in the tokenizer, not just the dataset MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `add_bos=False` only suppresses the dataset's own prepend. The tokenization itself goes through Bridge's HuggingFaceTokenizer, whose `text_to_ids` returns `self.tokenizer(text).input_ids` — i.e. add_special_tokens=True — whenever `include_special_tokens` is set, and it defaults to True. Verified in the container: `build_tokenizer(...)._tokenizer.include_special_tokens` is `True` by default and follows `hf_tokenizer_kwargs`. Because `prompt_template="{input}{output}"` tokenizes the two fields separately, a BOS-adding tokenizer would produce `[BOS, *input, BOS, *output]` — a BOS injected exactly where `answer_only_loss` starts scoring — and the documented "no BOS is prepended" contract would be false. Nemotron-Nano-3.5 never hit this because its tokenizer adds no special tokens either way, so the validated run says nothing about it. Set `include_special_tokens: False`, making the contract true by construction rather than assumed. Also: - `_warn_if_bos_missing` now probes behaviour (`tokenizer("x").input_ids[:1]`) instead of reading `add_bos_token`, which many fast tokenizers never expose even when their post-processor prepends BOS — a missing warning being the worse failure. Docstring records that only the first training record is inspected. - Help text, error message and README: `validation.jsonl` is required only when `--eval_iters > 0`. Signed-off-by: James Shen --- examples/megatron_bridge/README.md | 4 ++-- examples/megatron_bridge/distill.py | 28 ++++++++++++++++++++++++---- 2 files changed, 26 insertions(+), 6 deletions(-) diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 56f7aa276b7..25bcd480ec3 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -131,8 +131,8 @@ See the **[Dataset Preparation README](../dataset/README.md#tokenizing-for-megat 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 ` 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": , "output": }` records, which are tokenized with +with the loss masked to the completion. The directory must hold `training.jsonl` (and +`validation.jsonl` when `--eval_iters > 0`) of `{"input": , "output": }` 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. An EOS token is appended after the response. diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 6163f0c1c1d..d9e5fe980d8 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -139,7 +139,7 @@ def get_args(): "--sft_dataset_root", type=str, default=None, - help="Directory holding training.jsonl / validation.jsonl of " + help="Directory holding training.jsonl (and validation.jsonl when --eval_iters > 0) of " '{"input": , "output": } 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. An EOS " @@ -278,7 +278,8 @@ def get_args(): if args.sft and not args.sft_dataset_root: raise ValueError( - "--sft requires --sft_dataset_root (a directory with training.jsonl / validation.jsonl)." + "--sft requires --sft_dataset_root (a directory with training.jsonl, plus " + "validation.jsonl when --eval_iters > 0)." ) if args.sft and (args.data_paths or args.use_mock_data): raise ValueError( @@ -297,9 +298,20 @@ def _warn_if_bos_missing(tokenizer, dataset_root: str) -> None: ``--sft`` tokenizes both fields verbatim, so the caller owns the BOS. Adding one here is not safe (their text may already have it, which would double it), but training without the BOS the model is served with is silent train/inference skew. + + Bounded on purpose: it inspects only the first record of ``training.jsonl``, so a mixed corpus + whose first record happens to carry a BOS will not be flagged. """ bos = getattr(tokenizer, "bos_token", None) - if not bos or not getattr(tokenizer, "add_bos_token", False): + if not bos: + return + try: + # Probe rather than trusting ``add_bos_token``: many fast tokenizers prepend BOS via the + # post-processor without exposing that attribute, and a missing warning is the worse + # failure here. + if tokenizer("x").input_ids[:1] != [tokenizer.bos_token_id]: + return + except Exception: return try: with open(os.path.join(dataset_root, "training.jsonl")) as f: @@ -535,7 +547,15 @@ def _restore_student_hook(model_chunks): TokenizerConfig( tokenizer_type="HuggingFaceTokenizer", tokenizer_model=args.student_hf_path, - hf_tokenizer_kwargs={"trust_remote_code": args.trust_remote_code}, + hf_tokenizer_kwargs={ + "trust_remote_code": args.trust_remote_code, + # Enforce the verbatim contract here rather than relying on the dataset's + # add_bos/add_eos: text_to_ids adds special tokens when this is left at its + # default of True, and prompt_template tokenizes "{input}" and "{output}" + # separately -- so a BOS-adding tokenizer would inject one at the answer + # boundary, where answer_only_loss starts scoring. + "include_special_tokens": False, + }, ) if args.sft else TokenizerConfig( From 7baef79e4c1a85227b0c623bef8a7cdec07feaec Mon Sep 17 00:00:00 2001 From: James Shen Date: Tue, 11 Aug 2026 01:36:52 -0700 Subject: [PATCH 07/11] fix(megatron-bridge): truncate SFT prompts from the left, and relax the tokenizer check MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - `truncation_method` was left at GPTSFTDataset's default of "right", which truncates the END of "input" — the question and whatever turn marker the caller baked in, i.e. exactly the boundary `answer_only_loss` starts scoring at. An over-length record would train the model to begin responding at a position that never occurs at inference. Set "left" explicitly so the oldest context is dropped and the prompt->response seam survives. Confirmed against Bridge's GPTSFTDataset: `truncation_method: str = "right"` with `"right" -> ids[:expect_length]`. - The "verbatim" contract was slightly stronger than the truth: GPTSFTDataset applies `.strip(" ")` to each template field, so a significant leading or trailing SPACE is lost (a newline is not). Stated in the help text, the code comment and the README, with the recommendation to express separators as newlines. - `AutoTokenizer` moved to the module import block; `transformers` is already imported there, so none of the sanctioned reasons for a deferred import applied. - The teacher/student tokenizer check required full `get_vocab()` equality, which is stricter than what KD needs. Only student ids ever reach the teacher, so the invariant is that the teacher maps every student token to the same id; a teacher whose vocabulary is a strict superset (extra reserved tokens) is fine and now warns instead of failing. Cross-family pairs still raise. Signed-off-by: James Shen --- examples/megatron_bridge/README.md | 8 +++--- examples/megatron_bridge/distill.py | 42 ++++++++++++++++++----------- 2 files changed, 32 insertions(+), 18 deletions(-) diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 25bcd480ec3..232524c38fe 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -133,9 +133,11 @@ for full instructions on tokenizing JSONL files and Hugging Face datasets and ge Alternatively, pass `--sft --sft_dataset_root ` to distill on **raw prompt-completion JSONL** with the loss masked to the completion. The directory must hold `training.jsonl` (and `validation.jsonl` when `--eval_iters > 0`) of `{"input": , "output": }` 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. An EOS token is appended after the response. +the model's own HuggingFace tokenizer. Both fields are tokenized **as written**, except that +leading and trailing spaces on each field are stripped — 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, and express any significant separator as a newline rather than a trailing +space. An EOS token is appended after the response. ### Distillation with Real Data diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index d9e5fe980d8..04b251fd9d7 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -50,7 +50,7 @@ from megatron.core.datasets.utils import get_blend_from_list from megatron.core.distributed import DistributedDataParallelConfig from megatron.core.utils import unwrap_model -from transformers import AutoConfig +from transformers import AutoConfig, AutoTokenizer import modelopt.torch.distill as mtd import modelopt.torch.utils.distributed as dist @@ -141,9 +141,11 @@ def get_args(): default=None, help="Directory holding training.jsonl (and validation.jsonl when --eval_iters > 0) of " '{"input": , "output": } 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. An EOS " - "token is appended after the response.", + "tokenized as written, except that leading and trailing spaces on each field are " + "stripped: 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, and put any " + "significant separator inside the text as a newline rather than a space. An EOS token is " + "appended after the response.", ) # Training & Eval arguments parser.add_argument( @@ -338,21 +340,25 @@ def main(args: argparse.Namespace): # the mappings. Checked before the providers are built so a mismatch costs seconds. # The pretraining path is structurally immune: NullTokenizer plus pre-tokenized # --data_paths means one tokenization feeds both models. - from transformers import AutoTokenizer - _tok = {"trust_remote_code": args.trust_remote_code} student_tokenizer = AutoTokenizer.from_pretrained(args.student_hf_path, **_tok) student_vocab = student_tokenizer.get_vocab() teacher_vocab = AutoTokenizer.from_pretrained(args.teacher_hf_path, **_tok).get_vocab() - if student_vocab != teacher_vocab: - detail = ( - f"{len(student_vocab)} vs {len(teacher_vocab)} tokens" - if len(student_vocab) != len(teacher_vocab) - else f"both {len(student_vocab)} tokens, but different token->id mappings" - ) + # Only student ids ever reach the teacher, so agreement on those is the invariant. A + # teacher vocabulary that is a strict superset (extra reserved tokens, say) is fine. + 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( "--sft tokenizes with the student's tokenizer and scores the teacher on those " - f"same ids, so student and teacher must share a tokenizer ({detail})." + "same ids, so the teacher must map every student token to the same id: " + f"{len(conflicts)} conflicting id(s), {len(missing)} token(s) absent from the " + "teacher. Student and teacher must share a tokenizer." + ) + if len(teacher_vocab) > len(student_vocab): + warn_rank_0( + f"Teacher vocabulary has {len(teacher_vocab) - len(student_vocab)} token(s) the " + "student lacks; ids agree on every student token, so distillation is well-defined." ) _warn_if_bos_missing(student_tokenizer, args.sft_dataset_root) @@ -477,8 +483,10 @@ def _restore_student_hook(model_chunks): # 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. + # as written -- no chat template, no BOS, no role markers. One exception: GPTSFTDataset + # applies .strip(" ") to each field, so a significant trailing/leading SPACE is lost (a + # newline is not). Callers whose model expects markers 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, @@ -492,6 +500,10 @@ def _restore_student_hook(model_chunks): "prompt_template": "{input}{output}", "label_key": "output", "truncation_field": "input", + # GPTSFTDataset defaults to "right", which truncates the END of the prompt -- + # the question and whatever turn marker the caller baked in, i.e. exactly the + # boundary answer_only_loss starts scoring at. Drop the oldest context instead. + "truncation_method": "left", "answer_only_loss": True, "add_bos": False, "add_eos": True, From 34bf2fe69062dd86670f40f92b84c8a5f52bdf91 Mon Sep 17 00:00:00 2001 From: James Shen Date: Tue, 11 Aug 2026 03:14:32 -0700 Subject: [PATCH 08/11] fix(megatron-bridge): check logits width, not tokenizer size, and reject a stray --sft_dataset_root MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The teacher-superset case was downgraded to a warning claiming distillation is well-defined. It is not: LogitsKLLoss ends in F.kl_div(p, q) with p=[s, b, V_student] and q=[s, b, V_teacher], and TopKLogitsKLLoss gathers student logits with the teacher's topk indices — neither tolerates a width mismatch. Token-id agreement is necessary but not sufficient. len(get_vocab()) is also the wrong quantity: the losses operate on the models' padded vocab dimension, not the tokenizer's surface. Two tokenizers of different size can pad to the same width, and the reverse is possible too. So: keep the strict conflicts/missing id check (it still catches cross-family pairs early, before two 31B providers are built), drop the misleading warning, and compare the providers' padded vocab_size once they exist. Also reject --sft_dataset_root without --sft, which was accepted and silently ignored — the mirror of the mutual-exclusion check right above it. Signed-off-by: James Shen --- examples/megatron_bridge/distill.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 04b251fd9d7..6fabf4f17ab 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -288,6 +288,8 @@ def get_args(): "--sft is mutually exclusive with --data_paths / --use_mock_data: the SFT branch wins " "the dataset selection, so those inputs would be silently ignored." ) + if args.sft_dataset_root and not args.sft: + raise ValueError("--sft_dataset_root requires --sft; without it the SFT path is not used.") print_args(args) @@ -344,8 +346,8 @@ def main(args: argparse.Namespace): student_tokenizer = AutoTokenizer.from_pretrained(args.student_hf_path, **_tok) student_vocab = student_tokenizer.get_vocab() teacher_vocab = AutoTokenizer.from_pretrained(args.teacher_hf_path, **_tok).get_vocab() - # Only student ids ever reach the teacher, so agreement on those is the invariant. A - # teacher vocabulary that is a strict superset (extra reserved tokens, say) is fine. + # Every student id must mean the same thing to the teacher. This is necessary but not + # sufficient: the KD losses also need equal logits width, checked on the providers below. 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: @@ -355,11 +357,6 @@ def main(args: argparse.Namespace): f"{len(conflicts)} conflicting id(s), {len(missing)} token(s) absent from the " "teacher. Student and teacher must share a tokenizer." ) - if len(teacher_vocab) > len(student_vocab): - warn_rank_0( - f"Teacher vocabulary has {len(teacher_vocab) - len(student_vocab)} token(s) the " - "student lacks; ids agree on every student token, so distillation is well-defined." - ) _warn_if_bos_missing(student_tokenizer, args.sft_dataset_root) @@ -408,6 +405,13 @@ 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 KD losses compare logits directly, so the padded vocab dimensions must match. + raise ValueError( + "--sft distillation needs student and teacher logits of equal width, but their padded " + f"vocab sizes differ ({student_provider.vocab_size} vs {teacher_provider.vocab_size})." + ) + kd_config = ModelOptDistillConfig( skip_lm_loss=not args.no_skip_lm_loss, kd_loss_scale=args.kd_loss_scale ) From 94675893df9f795e44a1095214d25c5a78042bbc Mon Sep 17 00:00:00 2001 From: James Shen Date: Tue, 11 Aug 2026 05:36:25 -0700 Subject: [PATCH 09/11] fix(megatron-bridge): address review on distillation vocabulary and SFT input The teacher/student token-mapping check moves to get_args() and runs on every run, not just --sft: the KD losses score the teacher on the student's ids whatever the data path, so a disagreement is always fatal, and the pretraining path previously reached it only as a shape error from inside the loss. Loading a tokenizer is warned about rather than fatal, so VLM repos that ship only a processor keep working. The logits-width guard likewise drops its --sft gate, and now compares the padded vocab size it always claimed to check; provider.vocab_size is the raw HF config value, so the old message overstated what it verified. It is not redundant with the mapping check: two models can share a tokenizer and still declare different vocab sizes, and equal sizes do not imply equal mappings. --sft now fails at argparse time when training.jsonl (or validation.jsonl, when evaluating) is absent from --sft_dataset_root, rather than after both checkpoints have loaded onto GPUs. Documents that a record over --seq_length is truncated from the start of "input", dropping any BOS baked in there, and records why truncation_method is "left": "right" would cut the answer boundary and then the answer itself. Notes where Bridge consumes include_special_tokens. Adds an SFT case to tests/examples/megatron_bridge/test_distill.py. Signed-off-by: James Shen --- examples/megatron_bridge/README.md | 7 +- examples/megatron_bridge/distill.py | 112 ++++++++++++------ .../examples/megatron_bridge/test_distill.py | 38 ++++++ 3 files changed, 123 insertions(+), 34 deletions(-) diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 232524c38fe..f1a0950b807 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -137,7 +137,12 @@ the model's own HuggingFace tokenizer. Both fields are tokenized **as written**, leading and trailing spaces on each field are stripped — 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, and express any significant separator as a newline rather than a trailing -space. An EOS token is appended after the response. +space. An EOS token is appended after the response. A record longer than `--seq_length` is +truncated from the **start** of `"input"`, which drops any BOS, system prompt or opening role +marker baked in there, so pre-filter or pre-truncate the corpus if that matters. + +Teacher and student must share a tokenizer — distillation scores the teacher on the student's +token ids, and the KD losses compare the two models' logits elementwise over the vocab dimension. ### Distillation with Real Data diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 6fabf4f17ab..2f871b767ed 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -47,6 +47,7 @@ from megatron.bridge.training.distill import distill from megatron.bridge.training.post_training.checkpointing import has_modelopt_state from megatron.bridge.training.post_training.distillation import ModelOptDistillConfig +from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size from megatron.core.datasets.utils import get_blend_from_list from megatron.core.distributed import DistributedDataParallelConfig from megatron.core.utils import unwrap_model @@ -145,7 +146,9 @@ def get_args(): "stripped: 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, and put any " "significant separator inside the text as a newline rather than a space. An EOS token is " - "appended after the response.", + "appended after the response. A record longer than --seq_length is truncated from the " + "START of 'input', which drops any BOS, system prompt or opening role marker baked in " + "there; pre-filter or pre-truncate the corpus if that matters.", ) # Training & Eval arguments parser.add_argument( @@ -290,13 +293,53 @@ def get_args(): ) if args.sft_dataset_root and not args.sft: raise ValueError("--sft_dataset_root requires --sft; without it the SFT path is not used.") + if args.sft: + # Fail on a mistyped root here rather than after both checkpoints have loaded onto GPUs. + required = ["training.jsonl"] + (["validation.jsonl"] if args.eval_iters > 0 else []) + absent = [f for f in required if not os.path.isfile(os.path.join(args.sft_dataset_root, f))] + if absent: + raise ValueError(f"--sft_dataset_root {args.sft_dataset_root} is missing: {absent}.") + + _check_shared_vocabulary(args) print_args(args) return args -def _warn_if_bos_missing(tokenizer, dataset_root: str) -> None: +def _check_shared_vocabulary(args) -> None: + """Raise when teacher and student do not map every token to the same id. + + The KD losses compare the two models' logits elementwise over the vocab dimension, so a token + must mean the same thing to both regardless of how the data reaches them. Equal vocabulary + *sizes* are not enough -- Llama-2 and Mistral are both 32000 tokens with different mappings -- + so compare the mappings. Run before the providers are built so a mismatch costs seconds. + + A tokenizer that cannot be loaded (some VLM repos ship only a processor) is warned about + rather than fatal: this guard is not worth failing a run that would otherwise work. + """ + _tok = {"trust_remote_code": args.trust_remote_code} + try: + student_vocab = AutoTokenizer.from_pretrained(args.student_hf_path, **_tok).get_vocab() + teacher_vocab = AutoTokenizer.from_pretrained(args.teacher_hf_path, **_tok).get_vocab() + except Exception as e: + warn_rank_0( + f"Could not load both tokenizers to verify teacher and student share a vocabulary: {e}. " + "Distillation needs them to agree on every token id; check it yourself if in doubt." + ) + return + 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( + "Distillation scores the teacher on the student's token ids, so the teacher must map " + f"every student token to the same id: {len(conflicts)} conflicting id(s), " + f"{len(missing)} token(s) absent from the teacher. Student and teacher must share a " + "tokenizer." + ) + + +def _warn_if_bos_missing(args) -> None: """Warn when the tokenizer prepends a BOS at inference but the SFT records do not carry one. ``--sft`` tokenizes both fields verbatim, so the caller owns the BOS. Adding one here is not @@ -306,6 +349,13 @@ def _warn_if_bos_missing(tokenizer, dataset_root: str) -> None: Bounded on purpose: it inspects only the first record of ``training.jsonl``, so a mixed corpus whose first record happens to carry a BOS will not be flagged. """ + dataset_root = args.sft_dataset_root + try: + tokenizer = AutoTokenizer.from_pretrained( + args.student_hf_path, trust_remote_code=args.trust_remote_code + ) + except Exception: + return bos = getattr(tokenizer, "bos_token", None) if not bos: return @@ -336,29 +386,7 @@ def main(args: argparse.Namespace): tensorboard_dir = os.path.join(args.output_dir, "tb_logs") if args.sft: - # SFT tokenizes raw text with the student's tokenizer and scores the teacher on those same - # ids, so both models must agree on what every id means. Equal vocabulary *sizes* are not - # enough -- Llama-2 and Mistral are both 32000 tokens with different mappings -- so compare - # the mappings. Checked before the providers are built so a mismatch costs seconds. - # The pretraining path is structurally immune: NullTokenizer plus pre-tokenized - # --data_paths means one tokenization feeds both models. - _tok = {"trust_remote_code": args.trust_remote_code} - student_tokenizer = AutoTokenizer.from_pretrained(args.student_hf_path, **_tok) - student_vocab = student_tokenizer.get_vocab() - teacher_vocab = AutoTokenizer.from_pretrained(args.teacher_hf_path, **_tok).get_vocab() - # Every student id must mean the same thing to the teacher. This is necessary but not - # sufficient: the KD losses also need equal logits width, checked on the providers below. - 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( - "--sft tokenizes with the student's tokenizer and scores the teacher on those " - "same ids, so the teacher must map every student token to the same id: " - f"{len(conflicts)} conflicting id(s), {len(missing)} token(s) absent from the " - "teacher. Student and teacher must share a tokenizer." - ) - - _warn_if_bos_missing(student_tokenizer, args.sft_dataset_root) + _warn_if_bos_missing(args) # Build student and teacher model providers def _build_model_provider(hf_path, load_weights=True): @@ -405,11 +433,22 @@ 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 KD losses compare logits directly, so the padded vocab dimensions must match. + # The KD losses compare logits elementwise over the vocab dimension, so the two output layers + # must have the same width. That width is the *padded* vocab size, which the providers derive + # from their own ``vocab_size`` and ``make_vocab_size_divisible_by``; a shared tokenizer does + # not imply it, because the two HF configs can declare different vocab sizes. Independent of + # the data path: the pretraining path feeds both models one tokenization, but says nothing + # about how wide their output layers are. + padded = { + name: calculate_padded_vocab_size( + p.vocab_size, p.make_vocab_size_divisible_by, p.tensor_model_parallel_size + ) + for name, p in (("student", student_provider), ("teacher", teacher_provider)) + } + if padded["student"] != padded["teacher"]: raise ValueError( - "--sft distillation needs student and teacher logits of equal width, but their padded " - f"vocab sizes differ ({student_provider.vocab_size} vs {teacher_provider.vocab_size})." + "Distillation needs student and teacher logits of equal width, but their padded vocab " + f"sizes differ ({padded['student']} vs {padded['teacher']})." ) kd_config = ModelOptDistillConfig( @@ -504,9 +543,14 @@ def _restore_student_hook(model_chunks): "prompt_template": "{input}{output}", "label_key": "output", "truncation_field": "input", - # GPTSFTDataset defaults to "right", which truncates the END of the prompt -- - # the question and whatever turn marker the caller baked in, i.e. exactly the - # boundary answer_only_loss starts scoring at. Drop the oldest context instead. + # GPTSFTDataset defaults to "right", which is wrong here twice over: within + # "input" it drops the END of the prompt -- the question tail and whatever turn + # marker the caller baked in, i.e. exactly the boundary answer_only_loss starts + # scoring at -- and once the prompt alone cannot absorb the overflow it walks the + # template back-to-front and eats "output", the only span the loss is computed on. + # "left" drops the oldest context instead. The cost is that a record over + # seq_length loses the head of "input", including any BOS baked in there; see + # --sft_dataset_root help. "truncation_method": "left", "answer_only_loss": True, "add_bos": False, @@ -569,7 +613,9 @@ def _restore_student_hook(model_chunks): # add_bos/add_eos: text_to_ids adds special tokens when this is left at its # default of True, and prompt_template tokenizes "{input}" and "{output}" # separately -- so a BOS-adding tokenizer would inject one at the answer - # boundary, where answer_only_loss starts scoring. + # boundary, where answer_only_loss starts scoring. Consumed by Bridge in + # training/tokenizers/config.py, which reads it out of hf_tokenizer_kwargs; + # it is not forwarded blindly to AutoTokenizer. "include_special_tokens": False, }, ) diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py index 5dee51e2fc0..b9323f68438 100644 --- a/tests/examples/megatron_bridge/test_distill.py +++ b/tests/examples/megatron_bridge/test_distill.py @@ -14,6 +14,7 @@ # limitations under the License. """Tests for prune_minitron.py and distill.py scripts.""" +import json from pathlib import Path import pytest @@ -58,6 +59,43 @@ def test_distill_llm(tmp_path, num_gpus): assert (distilled_hf_path / "config.json").exists() +def test_distill_llm_sft(tmp_path, num_gpus): + """--sft distills from prompt-completion jsonl instead of pre-tokenized --data_paths.""" + teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True) + train_iters = 2 + gbs = 4 + dataset_root = tmp_path / "sft_data" + dataset_root.mkdir() + # More records than train_iters * gbs so the sampler does not run dry. + records = [{"input": f"Q: what follows {i}?\nA:", "output": f" {i + 1}"} for i in range(64)] + for split in ("training", "validation"): + (dataset_root / f"{split}.jsonl").write_text( + "\n".join(json.dumps(r) for r in records) + "\n" + ) + + distill_output_dir = tmp_path / "distill_output" + distill_cmd_parts = extend_cmd_parts( + ["torchrun", f"--nproc_per_node={num_gpus}", "distill.py", "--sft"], + student_hf_path=teacher_hf_path, + teacher_hf_path=teacher_hf_path, + sft_dataset_root=dataset_root, + output_dir=distill_output_dir, + tp_size=num_gpus, + pp_size=1, + seq_length=64, + mbs=1, + gbs=gbs, + train_iters=train_iters, + lr_warmup_iters=1, + eval_interval=train_iters, + eval_iters=1, + log_interval=1, + ) + run_example_command(distill_cmd_parts, example_path="megatron_bridge") + + assert (distill_output_dir / f"checkpoints/iter_{train_iters:07d}").exists() + + def test_distill_validate_only(tmp_path, num_gpus): teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True) train_iters = 2 From 9fd399a41634c3dcde2cc453fb47864967ed9df7 Mon Sep 17 00:00:00 2001 From: James Shen Date: Tue, 11 Aug 2026 06:22:12 -0700 Subject: [PATCH 10/11] fix(megatron-bridge): simplify the tokenizer check, trim comments, add changelog The teacher/student check now compares the two vocabularies for equality rather than counting conflicting and missing tokens. Distillation requires the same tokenizer on both sides -- the KD losses reduce elementwise over the vocab dimension -- so tolerating supersets implied a cross-tokenizer capability that does not exist. Shortens the comments and helper docstrings, and adds a 0.47 changelog entry for the SFT data support. Signed-off-by: James Shen --- CHANGELOG.rst | 1 + examples/megatron_bridge/distill.py | 84 +++++++---------------------- 2 files changed, 21 insertions(+), 64 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 71a084b99eb..964fd8483fc 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -12,6 +12,7 @@ Changelog *Megatron Framework (M-LM / M-Bridge)* +- Add SFT-masked data support to ``examples/megatron_bridge/distill.py``: ``--sft --sft_dataset_root `` distills on raw prompt-completion JSONL (``{"input", "output"}`` records) with the loss masked to the response tokens, using Megatron-Bridge's ``FinetuningDatasetConfig`` and the model's own HuggingFace tokenizer instead of the pretraining ``GPTDataset`` and ``NullTokenizer``. - Add per-expert weight quantization for Transformer Engine ``TEGroupedMLP`` (fused MoE experts): each expert now has its own ``weight_quantizer`` (a ``GroupedQuantizer`` holding one ``TensorQuantizer`` per expert) with an independent ``amax``, instead of a single shared ``amax`` across all experts. Applies to ``mtq.quantize`` calibration, HF / Megatron export, and QAD. - Add opt-in ``torch.compile`` execution for Transformer Engine grouped-linear per-expert weight quantizers while preserving their native checkpoint amax shapes. Set ``MODELOPT_TEGROUPED_COMPILE_WEIGHT_LOOP=1`` before quantized-module conversion; the default path remains eager. diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 2f871b767ed..374698533f2 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -308,15 +308,9 @@ def get_args(): def _check_shared_vocabulary(args) -> None: - """Raise when teacher and student do not map every token to the same id. + """Raise unless teacher and student use the same tokenizer. - The KD losses compare the two models' logits elementwise over the vocab dimension, so a token - must mean the same thing to both regardless of how the data reaches them. Equal vocabulary - *sizes* are not enough -- Llama-2 and Mistral are both 32000 tokens with different mappings -- - so compare the mappings. Run before the providers are built so a mismatch costs seconds. - - A tokenizer that cannot be loaded (some VLM repos ship only a processor) is warned about - rather than fatal: this guard is not worth failing a run that would otherwise work. + Warns instead when a tokenizer cannot be loaded (some VLM repos ship only a processor). """ _tok = {"trust_remote_code": args.trust_remote_code} try: @@ -328,26 +322,17 @@ def _check_shared_vocabulary(args) -> None: "Distillation needs them to agree on every token id; check it yourself if in doubt." ) return - 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: + if student_vocab != teacher_vocab: raise ValueError( - "Distillation scores the teacher on the student's token ids, so the teacher must map " - f"every student token to the same id: {len(conflicts)} conflicting id(s), " - f"{len(missing)} token(s) absent from the teacher. Student and teacher must share a " - "tokenizer." + "Distillation scores the teacher on the student's token ids, so teacher and student " + "must use the same tokenizer." ) def _warn_if_bos_missing(args) -> None: """Warn when the tokenizer prepends a BOS at inference but the SFT records do not carry one. - ``--sft`` tokenizes both fields verbatim, so the caller owns the BOS. Adding one here is not - safe (their text may already have it, which would double it), but training without the BOS the - model is served with is silent train/inference skew. - - Bounded on purpose: it inspects only the first record of ``training.jsonl``, so a mixed corpus - whose first record happens to carry a BOS will not be flagged. + Inspects only the first record of ``training.jsonl``. """ dataset_root = args.sft_dataset_root try: @@ -403,12 +388,8 @@ def _build_model_provider(hf_path, load_weights=True): 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. + # A response-only loss mask needs per-token reduction to combine across CP ranks. + # Must stay in sync with ``average_in_collective=not args.sft`` on the DDP config. provider.calculate_per_token_loss = True if args.recompute_granularity is not None: provider.recompute_granularity = args.recompute_granularity @@ -433,12 +414,8 @@ def _build_model_provider(hf_path, load_weights=True): student_provider.gradient_accumulation_fusion = False teacher_provider = _build_model_provider(args.teacher_hf_path) - # The KD losses compare logits elementwise over the vocab dimension, so the two output layers - # must have the same width. That width is the *padded* vocab size, which the providers derive - # from their own ``vocab_size`` and ``make_vocab_size_divisible_by``; a shared tokenizer does - # not imply it, because the two HF configs can declare different vocab sizes. Independent of - # the data path: the pretraining path feeds both models one tokenization, but says nothing - # about how wide their output layers are. + # The KD losses compare logits elementwise over the vocab dim, so both output layers must have + # the same padded width. A shared tokenizer does not imply it: the HF configs can disagree. padded = { name: calculate_padded_vocab_size( p.vocab_size, p.make_vocab_size_divisible_by, p.tensor_model_parallel_size @@ -455,10 +432,8 @@ def _build_model_provider(hf_path, load_weights=True): skip_lm_loss=not args.no_skip_lm_loss, kd_loss_scale=args.kd_loss_scale ) - # VLM detection convention: HF VLM configs expose a ``vision_config``, and Megatron-Bridge nests - # the text model under the ``language_model`` submodule (used as ``distill_submodule`` below). If a - # future model breaks either convention, the ``getattr(model, "language_model")`` in the provider - # will error loudly rather than silently distilling the wrong module. + # HF VLM configs expose ``vision_config``; Megatron-Bridge nests the text model under + # ``language_model`` (used as ``distill_submodule`` below). is_vlm = hasattr( AutoConfig.from_pretrained(args.student_hf_path, trust_remote_code=args.trust_remote_code), "vision_config", @@ -518,18 +493,9 @@ def _restore_student_hook(model_chunks): "skip_getting_attention_mask_from_dataset": True, } 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 - # as written -- no chat template, no BOS, no role markers. One exception: GPTSFTDataset - # applies .strip(" ") to each field, so a significant trailing/leading SPACE is lost (a - # newline is not). Callers whose model expects markers must bake them into the "input" - # field; see --sft_dataset_root help. + # SFT-masked distillation via Bridge's FinetuningDatasetConfig -> NeMo-style GPTSFTDataset, + # reading {"input", "output"} jsonl. Fields are tokenized as written except that each is + # ``.strip(" ")``-ed; see --sft_dataset_root help. dataset_config = FinetuningDatasetConfig( seq_length=args.seq_length, dataset_root=args.sft_dataset_root, @@ -543,14 +509,8 @@ def _restore_student_hook(model_chunks): "prompt_template": "{input}{output}", "label_key": "output", "truncation_field": "input", - # GPTSFTDataset defaults to "right", which is wrong here twice over: within - # "input" it drops the END of the prompt -- the question tail and whatever turn - # marker the caller baked in, i.e. exactly the boundary answer_only_loss starts - # scoring at -- and once the prompt alone cannot absorb the overflow it walks the - # template back-to-front and eats "output", the only span the loss is computed on. - # "left" drops the oldest context instead. The cost is that a record over - # seq_length loses the head of "input", including any BOS baked in there; see - # --sft_dataset_root help. + # Drop the oldest context. The default "right" would cut the prompt/answer + # boundary and then "output" itself, the only span the loss is computed on. "truncation_method": "left", "answer_only_loss": True, "add_bos": False, @@ -609,13 +569,9 @@ def _restore_student_hook(model_chunks): tokenizer_model=args.student_hf_path, hf_tokenizer_kwargs={ "trust_remote_code": args.trust_remote_code, - # Enforce the verbatim contract here rather than relying on the dataset's - # add_bos/add_eos: text_to_ids adds special tokens when this is left at its - # default of True, and prompt_template tokenizes "{input}" and "{output}" - # separately -- so a BOS-adding tokenizer would inject one at the answer - # boundary, where answer_only_loss starts scoring. Consumed by Bridge in - # training/tokenizers/config.py, which reads it out of hf_tokenizer_kwargs; - # it is not forwarded blindly to AutoTokenizer. + # Default True would make text_to_ids inject a BOS at the answer boundary, + # since "{input}" and "{output}" are tokenized separately. Consumed by Bridge + # in training/tokenizers/config.py. "include_special_tokens": False, }, ) From 8fe9543bda69d757c000afce4b07d9f3967d94d3 Mon Sep 17 00:00:00 2001 From: James Shen Date: Wed, 12 Aug 2026 09:19:21 -0700 Subject: [PATCH 11/11] fix(megatron-bridge): raise on tokenizer failures and add BOS via the framework Review follow-ups on --sft: - Drop the three try/except blocks that turned a failed tokenizer load or a malformed record into a warning. A broken tokenizer is fatal, and a warning scrolls past in a multi-day run. - Move the BOS handling into get_args() with the other sanity checks. - Set add_bos from a tokenizer probe instead of hardcoding False. GPTSFTDataset prepends it after truncation and budgets for it in total_ids, so the BOS survives a record that had to be cut; baking it into 'input' by hand did not, because truncation_method='left' drops the front of that field. The probe is needed because bos_token_id alone is not decisive: Nemotron 3.5 Lightning declares bos_token_id=1 with add_bos_token=False and prepends nothing. - README and --sft_dataset_root help no longer tell users to add the BOS themselves, which would now double it. Signed-off-by: James Shen --- examples/megatron_bridge/README.md | 13 ++--- examples/megatron_bridge/distill.py | 83 +++++++---------------------- 2 files changed, 27 insertions(+), 69 deletions(-) diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index f1a0950b807..f247a0b6137 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -134,12 +134,13 @@ Alternatively, pass `--sft --sft_dataset_root ` to distill on **raw prompt- with the loss masked to the completion. The directory must hold `training.jsonl` (and `validation.jsonl` when `--eval_iters > 0`) of `{"input": , "output": }` records, which are tokenized with the model's own HuggingFace tokenizer. Both fields are tokenized **as written**, except that -leading and trailing spaces on each field are stripped — 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, and express any significant separator as a newline rather than a trailing -space. An EOS token is appended after the response. A record longer than `--seq_length` is -truncated from the **start** of `"input"`, which drops any BOS, system prompt or opening role -marker baked in there, so pre-filter or pre-truncate the corpus if that matters. +leading and trailing spaces on each field are stripped — no chat template is applied. So if your +model expects role/turn markers, include them in the `"input"` field yourself, and express any +significant separator as a newline rather than a trailing space. A BOS token is prepended +automatically when the tokenizer prepends one at inference, so do not add it yourself; an EOS +token is appended after the response. A record longer than `--seq_length` is truncated from the +**start** of `"input"`, which drops any system prompt or opening role marker baked in there, so +pre-filter or pre-truncate the corpus if that matters. Teacher and student must share a tokenizer — distillation scores the teacher on the student's token ids, and the KD losses compare the two models' logits elementwise over the vocab dimension. diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 374698533f2..598c95ecb49 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -22,7 +22,6 @@ import argparse import contextlib -import json import os import torch @@ -131,24 +130,16 @@ def get_args(): 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.", + help="Distill on prompt-completion jsonl from --sft_dataset_root with the loss masked to " + "the completion, instead of pre-tokenized --data_paths.", ) parser.add_argument( "--sft_dataset_root", type=str, default=None, help="Directory holding training.jsonl (and validation.jsonl when --eval_iters > 0) of " - '{"input": , "output": } records (used with --sft). Both fields are ' - "tokenized as written, except that leading and trailing spaces on each field are " - "stripped: 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, and put any " - "significant separator inside the text as a newline rather than a space. An EOS token is " - "appended after the response. A record longer than --seq_length is truncated from the " - "START of 'input', which drops any BOS, system prompt or opening role marker baked in " - "there; pre-filter or pre-truncate the corpus if that matters.", + '{"input": , "output": } records (used with --sft). See the README for ' + "how the fields are tokenized and truncated.", ) # Training & Eval arguments parser.add_argument( @@ -299,6 +290,8 @@ def get_args(): absent = [f for f in required if not os.path.isfile(os.path.join(args.sft_dataset_root, f))] if absent: raise ValueError(f"--sft_dataset_root {args.sft_dataset_root} is missing: {absent}.") + # Decided once here so it reaches print_args and costs a single tokenizer load. + args.sft_add_bos = _tokenizer_prepends_bos(args) _check_shared_vocabulary(args) @@ -308,20 +301,10 @@ def get_args(): def _check_shared_vocabulary(args) -> None: - """Raise unless teacher and student use the same tokenizer. - - Warns instead when a tokenizer cannot be loaded (some VLM repos ship only a processor). - """ + """Raise unless teacher and student use the same tokenizer.""" _tok = {"trust_remote_code": args.trust_remote_code} - try: - student_vocab = AutoTokenizer.from_pretrained(args.student_hf_path, **_tok).get_vocab() - teacher_vocab = AutoTokenizer.from_pretrained(args.teacher_hf_path, **_tok).get_vocab() - except Exception as e: - warn_rank_0( - f"Could not load both tokenizers to verify teacher and student share a vocabulary: {e}. " - "Distillation needs them to agree on every token id; check it yourself if in doubt." - ) - return + student_vocab = AutoTokenizer.from_pretrained(args.student_hf_path, **_tok).get_vocab() + teacher_vocab = AutoTokenizer.from_pretrained(args.teacher_hf_path, **_tok).get_vocab() if student_vocab != teacher_vocab: raise ValueError( "Distillation scores the teacher on the student's token ids, so teacher and student " @@ -329,50 +312,23 @@ def _check_shared_vocabulary(args) -> None: ) -def _warn_if_bos_missing(args) -> None: - """Warn when the tokenizer prepends a BOS at inference but the SFT records do not carry one. +def _tokenizer_prepends_bos(args) -> bool: + """True when the student tokenizer prepends a BOS at inference. - Inspects only the first record of ``training.jsonl``. + Probes an encode: fast tokenizers prepend via a post-processor that exposes no attribute. """ - dataset_root = args.sft_dataset_root - try: - tokenizer = AutoTokenizer.from_pretrained( - args.student_hf_path, trust_remote_code=args.trust_remote_code - ) - except Exception: - return - bos = getattr(tokenizer, "bos_token", None) - if not bos: - return - try: - # Probe rather than trusting ``add_bos_token``: many fast tokenizers prepend BOS via the - # post-processor without exposing that attribute, and a missing warning is the worse - # failure here. - if tokenizer("x").input_ids[:1] != [tokenizer.bos_token_id]: - return - except Exception: - return - try: - with open(os.path.join(dataset_root, "training.jsonl")) as f: - first_input = str(json.loads(f.readline()).get("input", "")) - except Exception: - return # a malformed or missing file is the dataset builder's error to report, not ours - if not first_input.startswith(bos): - warn_rank_0( - f"This tokenizer prepends {bos!r} at inference, but the first record in " - f"{dataset_root}/training.jsonl does not start with it. --sft tokenizes the fields " - f"verbatim, so the model would be trained without the BOS it is served with. Include " - f"{bos!r} at the start of the 'input' field." - ) + tokenizer = AutoTokenizer.from_pretrained( + args.student_hf_path, trust_remote_code=args.trust_remote_code + ) + if not getattr(tokenizer, "bos_token", None): + return False + return tokenizer("x").input_ids[:1] == [tokenizer.bos_token_id] def main(args: argparse.Namespace): checkpoint_dir = os.path.join(args.output_dir, "checkpoints") tensorboard_dir = os.path.join(args.output_dir, "tb_logs") - if args.sft: - _warn_if_bos_missing(args) - # Build student and teacher model providers def _build_model_provider(hf_path, load_weights=True): bridge = AutoBridge.from_hf_pretrained(hf_path, trust_remote_code=args.trust_remote_code) @@ -513,7 +469,8 @@ def _restore_student_hook(model_chunks): # boundary and then "output" itself, the only span the loss is computed on. "truncation_method": "left", "answer_only_loss": True, - "add_bos": False, + # Prepended after truncation, so it survives a record that had to be cut. + "add_bos": args.sft_add_bos, "add_eos": True, }, )