diff --git a/CHANGELOG.rst b/CHANGELOG.rst index bd3e244b485..d3cc67d95bf 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -50,6 +50,7 @@ Changelog *Megatron Framework (M-LM / M-Bridge)* +- Add chat-data distillation to ``examples/megatron_bridge/distill.py``: ``--sft --sft_hf_dataset`` takes a local chat ``.jsonl`` or a Hugging Face Hub dataset, renders each conversation with the model's chat template, and ``--sft_loss_mode`` selects which turns are trained. Requires a Megatron-Bridge with ``DirectHFSFTDatasetConfig``. - Add mergeable LoRA QAD for fake-quantized dense Megatron/Transformer Engine linear layers. Enable it with ``--lora_rank`` in Megatron-Bridge PTQ, train through the normal QAD flow, and export a standalone quantized checkpoint without deployment adapters. - Add optional Top-P (nucleus) truncation to the Megatron top-k logits KD loss (now ``TopLogitsKLLoss``) via ``logit_kl_top_p`` and ``logit_kl_top_p_min_k`` in ``DistillationConfig``: after the global Top-K selection, only the smallest prefix whose cumulative teacher probability reaches ``top_p`` (with a floor of ``min_k`` entries) contributes to the KL, mirroring ``--logits-save-top-p`` / ``--logits-save-top-p-min-k`` in Megatron-LM's logits saver. - Add an end-to-end W4A4 NVFP4 PTQ and QAD tutorial for Qwen3.6-35B-A3B also covering evaluation and vLLM throughput benchmarking. See `examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/README.md `_ for details. diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index fb4bc930f07..3e711e3618d 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -175,6 +175,15 @@ token is appended after the response. A record longer than `--seq_length` is tru **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. +To distill on **chat data** instead, pass `--sft --sft_hf_dataset `, where `` is a +local `.jsonl` of `{"messages": [{"role": ..., "content": ...}, ...]}` records or a Hugging +Face Hub dataset `[:]` (split defaults to `train`, and may be a slice such as +`train[:1000]`). Each conversation is rendered with the model's own chat template, and +`--sft_loss_mode` picks the trained tokens: every assistant turn (`assistant`, the default), only +the last one (`last_turn`), or the whole conversation (`full`). With `--eval_iters > 0`, also pass +`--sft_hf_validation` in the same format. This path needs a Megatron-Bridge with +`DirectHFSFTDatasetConfig`. + 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 f71114cf1f3..e92b188020c 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -48,6 +48,15 @@ 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 + +try: + from megatron.bridge.data.builders import DirectHFSFTDatasetConfig + from megatron.bridge.data.sft_processing import ChatSFTPreprocessingConfig + from megatron.bridge.data.sources import HFDatasetSourceConfig + + HAS_DIRECT_HF_SFT = True +except ImportError: + HAS_DIRECT_HF_SFT = False from megatron.core.datasets.utils import get_blend_from_list from megatron.core.distributed import DistributedDataParallelConfig from megatron.core.utils import unwrap_model @@ -117,6 +126,16 @@ def _nonnegative_int(value: str) -> int: return parsed +def _hf_source(spec: str) -> "HFDatasetSourceConfig": + """``.json[l]`` -> local chat jsonl; otherwise ``[:]`` (split defaults to train).""" + if spec.endswith((".json", ".jsonl")): + return HFDatasetSourceConfig( + path_or_dataset="json", split="train", load_kwargs={"data_files": spec} + ) + dataset, _, split = spec.partition(":") + return HFDatasetSourceConfig(path_or_dataset=dataset, split=split or "train") + + def get_args(): """Parse command-line arguments.""" parser = argparse.ArgumentParser(description="Distillation for Megatron-Bridge.") @@ -181,8 +200,9 @@ def get_args(): 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.", + help="Distill on SFT data with the loss masked to the responses, instead of pre-tokenized " + "--data_paths: prompt-completion jsonl from --sft_dataset_root, or chat data from " + "--sft_hf_dataset.", ) parser.add_argument( "--sft_dataset_root", @@ -192,6 +212,32 @@ def get_args(): '{"input": , "output": } records (used with --sft). See the README for ' "how the fields are tokenized and truncated.", ) + parser.add_argument( + "--sft_hf_dataset", + type=str, + default=None, + help="Chat dataset to distill on (used with --sft): a local chat '.jsonl', or " + "'[:]' (split defaults to train). Only .json/.jsonl are read as local " + "files; anything else is taken as a Hub id. Each conversation is rendered by the " + "model's own chat template. Alternative to --sft_dataset_root, which expects " + '{"input", "output"} records already preprocessed into a dataset root.', + ) + parser.add_argument( + "--sft_hf_validation", + type=str, + default=None, + help="Validation data for --sft_hf_dataset, in the same format (required when " + "--eval_iters > 0), e.g. ':test' or ':train[:1000]'.", + ) + parser.add_argument( + "--sft_loss_mode", + type=str, + default="assistant", + choices=["assistant", "last_turn", "full"], + help="Which tokens --sft_hf_dataset trains on: every assistant turn, only the final one, " + "or the whole conversation. Multi-turn records train all of their responses under " + "'assistant'.", + ) # Training & Eval arguments parser.add_argument( "--output_dir", type=str, required=True, help="Folder for logging and checkpoint saving" @@ -354,11 +400,24 @@ 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: + if args.sft and not (args.sft_dataset_root or args.sft_hf_dataset): raise ValueError( "--sft requires --sft_dataset_root (a directory with training.jsonl, plus " - "validation.jsonl when --eval_iters > 0)." + "validation.jsonl when --eval_iters > 0) or --sft_hf_dataset." ) + if args.sft_dataset_root and args.sft_hf_dataset: + raise ValueError("--sft_dataset_root and --sft_hf_dataset are mutually exclusive.") + if args.sft_hf_dataset: + if not args.sft: + raise ValueError("--sft_hf_dataset requires --sft.") + if not HAS_DIRECT_HF_SFT: + raise ValueError( + "--sft_hf_dataset needs a newer Megatron-Bridge (DirectHFSFTDatasetConfig)." + ) + if args.eval_iters > 0 and not args.sft_hf_validation: + raise ValueError("--sft_hf_dataset with --eval_iters > 0 needs --sft_hf_validation.") + elif args.sft_hf_validation or args.sft_loss_mode != "assistant": + raise ValueError("--sft_hf_validation / --sft_loss_mode require --sft_hf_dataset.") 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 " @@ -366,7 +425,7 @@ 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: + if args.sft and args.sft_dataset_root: # 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))] @@ -557,7 +616,21 @@ def _restore_student_hook(model_chunks): "dataloader_type": "single", "skip_getting_attention_mask_from_dataset": True, } - if args.sft: + if args.sft and args.sft_hf_dataset: + # Chat rows rendered by the model's own chat template, with the loss covering every + # assistant turn, so multi-round records train all of their responses rather than the last. + dataset_config = DirectHFSFTDatasetConfig( + seq_length=args.seq_length, + source=_hf_source(args.sft_hf_dataset), + validation_source=( + _hf_source(args.sft_hf_validation) if args.sft_hf_validation else None + ), + preprocessing=ChatSFTPreprocessingConfig(loss_mode=args.sft_loss_mode), + dataloader_type="batch", + do_validation=args.eval_iters > 0, + do_test=False, + ) + elif 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. diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py index 75c3e430958..e6c25dcb4a2 100644 --- a/tests/examples/megatron_bridge/test_distill.py +++ b/tests/examples/megatron_bridge/test_distill.py @@ -34,6 +34,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "examples" / "megatron_bridge")) +from distill import HAS_DIRECT_HF_SFT, _hf_source from export_distilled_megatron_to_hf import _get_checkpoint_export_paths @@ -120,6 +121,75 @@ def test_distill_llm_sft(tmp_path, num_gpus): assert (distill_output_dir / f"checkpoints/iter_{train_iters:07d}").exists() +@pytest.mark.skipif(not HAS_DIRECT_HF_SFT, reason="Megatron-Bridge lacks DirectHFSFTDatasetConfig") +def test_distill_llm_sft_hf_chat(tmp_path, num_gpus): + """--sft_hf_dataset distills on chat jsonl rendered by the model's chat template.""" + teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True) + train_iters = 2 + gbs = 4 + # More records than train_iters * gbs so the sampler does not run dry; every 4th is multi-turn. + records = [] + for i in range(64): + messages = [ + {"role": "user", "content": f"what follows {i}?"}, + {"role": "assistant", "content": f"{i + 1}"}, + ] + if i % 4 == 0: + messages += [ + {"role": "user", "content": "and then?"}, + {"role": "assistant", "content": f"{i + 2}"}, + ] + records.append({"messages": messages}) + data_files = {} + for split in ("train", "validation"): + data_files[split] = tmp_path / f"chat_{split}.jsonl" + data_files[split].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_hf_dataset=data_files["train"], + sft_hf_validation=data_files["validation"], + 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() + + +@pytest.mark.skipif(not HAS_DIRECT_HF_SFT, reason="Megatron-Bridge lacks DirectHFSFTDatasetConfig") +@pytest.mark.parametrize( + ("spec", "path_or_dataset", "split", "load_kwargs"), + [ + ("data/foo.jsonl", "json", "train", {"data_files": "data/foo.jsonl"}), + ("data/foo.json", "json", "train", {"data_files": "data/foo.json"}), + ("org/ds", "org/ds", "train", None), + ("org/ds:test", "org/ds", "test", None), + # Only the first ':' separates the split, so slice syntax survives. + ("org/ds:train[:1000]", "org/ds", "train[:1000]", None), + ], +) +def test_hf_source_spec(spec, path_or_dataset, split, load_kwargs): + source = _hf_source(spec) + assert (source.path_or_dataset, source.split, source.load_kwargs) == ( + path_or_dataset, + split, + load_kwargs, + ) + + def test_distill_validate_only(tmp_path, num_gpus): teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True) train_iters = 2