From 8abcdb440366e5cc345beef101169da8d267f333 Mon Sep 17 00:00:00 2001 From: James Shen Date: Fri, 7 Aug 2026 08:31:38 -0700 Subject: [PATCH 1/3] 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 16dd37d5f8d..da0c85712bb 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 080f4de42b3dc8c32b6a0718b64602cff8cbbed7 Mon Sep 17 00:00:00 2001 From: James Shen Date: Fri, 7 Aug 2026 09:04:58 -0700 Subject: [PATCH 2/3] 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 2fae7fe3545..253f3e15f2f 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 da0c85712bb..ae0ebaf70a4 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 8c19f38bc26cfc353212f48b430be8479fd3dc13 Mon Sep 17 00:00:00 2001 From: James Shen Date: Fri, 7 Aug 2026 09:21:52 -0700 Subject: [PATCH 3/3] 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 ae0ebaf70a4..956cb07119a 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}",