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/README.md b/examples/megatron_bridge/README.md index 030edaf8fd5..f247a0b6137 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -130,6 +130,21 @@ 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` 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. 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. + ### 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 b35369c40f3..598c95ecb49 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, @@ -45,10 +46,11 @@ 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 -from transformers import AutoConfig +from transformers import AutoConfig, AutoTokenizer import modelopt.torch.distill as mtd import modelopt.torch.utils.distributed as dist @@ -125,6 +127,20 @@ 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="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). See the README for ' + "how the fields are tokenized and truncated.", + ) # Training & Eval arguments parser.add_argument( "--output_dir", type=str, required=True, help="Folder for logging and checkpoint saving" @@ -246,7 +262,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: @@ -256,11 +272,59 @@ 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, plus " + "validation.jsonl when --eval_iters > 0)." + ) + 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." + ) + 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}.") + # 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) + print_args(args) return args +def _check_shared_vocabulary(args) -> None: + """Raise unless teacher and student use the same tokenizer.""" + _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: + raise ValueError( + "Distillation scores the teacher on the student's token ids, so teacher and student " + "must use the same tokenizer." + ) + + +def _tokenizer_prepends_bos(args) -> bool: + """True when the student tokenizer prepends a BOS at inference. + + Probes an encode: fast tokenizers prepend via a post-processor that exposes no attribute. + """ + 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") @@ -279,6 +343,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: + # 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 provider.recompute_method = args.recompute_method @@ -302,14 +370,26 @@ 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 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 + ) + for name, p in (("student", student_provider), ("teacher", teacher_provider)) + } + if padded["student"] != padded["teacher"]: + raise ValueError( + "Distillation needs student and teacher logits of equal width, but their padded vocab " + f"sizes differ ({padded['student']} vs {padded['teacher']})." + ) + kd_config = ModelOptDistillConfig( 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", @@ -368,7 +448,33 @@ 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, + # 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, + 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, + dataset_kwargs={ + "prompt_template": "{input}{output}", + "label_key": "output", + "truncation_field": "input", + # 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, + # Prepended after truncation, so it survives a record that had to be cut. + "add_bos": args.sft_add_bos, + "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 +505,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 +518,24 @@ 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, + # 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, + }, + ) + if args.sft + else TokenizerConfig( + tokenizer_type="NullTokenizer", vocab_size=distill_provider.vocab_size + ) ), checkpoint=CheckpointConfig( save_interval=( 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