From ab441f4020556ff83caf302abcb5a8ca336db2a0 Mon Sep 17 00:00:00 2001 From: Yue Shen Date: Tue, 29 Sep 2026 04:49:42 -0700 Subject: [PATCH 1/5] examples/megatron_bridge: distill on chat data straight from HuggingFace/JSONL distill.py can only read SFT data from --sft_dataset_root, which expects {"input", "output"} records already preprocessed into a Megatron dataset root. This adds the complementary path: point --sft_hf_dataset at a Hub dataset id, or at "json" with --sft_hf_data_files for a local chat jsonl, and each conversation is rendered by the model's own chat template with the loss masked to assistant turns. Multi-turn records train on all of their responses, not only the last. The two sources are mutually exclusive and the existing --sft_dataset_root path is unchanged. The Megatron-Bridge symbols this needs are imported defensively, so an older Bridge still runs every other mode and --sft_hf_dataset reports why it is unavailable instead of failing at import. Flags: --sft_hf_dataset, --sft_hf_split, --sft_hf_data_files, --sft_hf_validation_split, --sft_hf_validation_data_files, --sft_loss_mode. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Yue Shen --- examples/megatron_bridge/distill.py | 94 +++++++++++++++++++++++++++-- 1 file changed, 90 insertions(+), 4 deletions(-) diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index f71114cf1f3..685903a3cd5 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 @@ -192,6 +201,45 @@ 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 Hub dataset id, or 'json' together " + "with --sft_hf_data_files for a local chat jsonl. 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_split", type=str, default="train", help="Split for --sft_hf_dataset" + ) + parser.add_argument( + "--sft_hf_data_files", + type=str, + default=None, + help="Local chat jsonl for --sft_hf_dataset json (HuggingFace load_kwargs.data_files)", + ) + parser.add_argument( + "--sft_hf_validation_split", + type=str, + default=None, + help="Validation split for --sft_hf_dataset (required when --eval_iters > 0)", + ) + parser.add_argument( + "--sft_hf_validation_data_files", + type=str, + default=None, + help="Local chat jsonl for the validation split of --sft_hf_dataset json", + ) + 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 +402,25 @@ 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 Megatron-Bridge providing DirectHFSFTDatasetConfig " + "(added 2026-07-09). Use a newer Bridge or --sft_dataset_root." + ) + if args.eval_iters > 0 and not args.sft_hf_validation_split: + raise ValueError( + "--sft_hf_dataset with --eval_iters > 0 needs --sft_hf_validation_split." + ) 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 +428,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 +619,31 @@ 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. + def _hf_source(split, data_files=None): + return HFDatasetSourceConfig( + path_or_dataset=args.sft_hf_dataset, + split=split, + load_kwargs={"data_files": data_files} if data_files else None, + ) + + dataset_config = DirectHFSFTDatasetConfig( + seq_length=args.seq_length, + source=_hf_source(args.sft_hf_split, args.sft_hf_data_files), + validation_source=( + _hf_source(args.sft_hf_validation_split, args.sft_hf_validation_data_files) + if args.sft_hf_validation_split + else None + ), + preprocessing=ChatSFTPreprocessingConfig(loss_mode=args.sft_loss_mode), + seed=args.seed, + 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. From 0d05131b3cef9503110cf02d583b667ac5f799fb Mon Sep 17 00:00:00 2001 From: Yue Date: Thu, 8 Oct 2026 07:04:42 +0000 Subject: [PATCH 2/5] examples/megatron_bridge: drop seed from DirectHFSFTDatasetConfig Megatron-Bridge's DirectHFSFTDatasetConfig does not accept a seed keyword, so building the --sft_hf_dataset config failed before training started. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yue --- examples/megatron_bridge/distill.py | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 685903a3cd5..7856ecfdd2d 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -638,7 +638,6 @@ def _hf_source(split, data_files=None): else None ), preprocessing=ChatSFTPreprocessingConfig(loss_mode=args.sft_loss_mode), - seed=args.seed, dataloader_type="batch", do_validation=args.eval_iters > 0, do_test=False, From 875aebeaa3552d5ad8a8cfcc5c0c63eb29ae161c Mon Sep 17 00:00:00 2001 From: Yue Date: Thu, 8 Oct 2026 07:44:37 +0000 Subject: [PATCH 3/5] examples/megatron_bridge: take --sft_hf_dataset as one spec, validate it up front Replace --sft_hf_split, --sft_hf_data_files, --sft_hf_validation_split and --sft_hf_validation_data_files with a single format shared by --sft_hf_dataset and the new --sft_hf_validation: a local '.jsonl', or '[:]'. Check in get_args() that --eval_iters > 0 has validation data and that local files exist, rather than failing in the Bridge builder after both checkpoints are on GPUs; reject --sft_hf_validation / --sft_loss_mode without --sft_hf_dataset instead of ignoring them. Note why the DirectHFSFTDatasetConfig import is guarded. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yue --- examples/megatron_bridge/distill.py | 69 +++++++++++++++-------------- 1 file changed, 35 insertions(+), 34 deletions(-) diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 7856ecfdd2d..4bd38694243 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -49,6 +49,8 @@ from megatron.bridge.training.post_training.distillation import ModelOptDistillConfig from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size +# DirectHFSFTDatasetConfig landed in Megatron-Bridge on 2026-07-09; an older Bridge still runs every +# other data mode, and --sft_hf_dataset reports why it is unavailable. try: from megatron.bridge.data.builders import DirectHFSFTDatasetConfig from megatron.bridge.data.sft_processing import ChatSFTPreprocessingConfig @@ -126,6 +128,20 @@ def _nonnegative_int(value: str) -> int: return parsed +def _is_local_json(spec: str) -> bool: + return spec.endswith((".json", ".jsonl")) + + +def _hf_source(spec: str) -> "HFDatasetSourceConfig": + """``.json[l]`` -> local chat jsonl; otherwise ``[:]`` (split defaults to train).""" + if _is_local_json(spec): + 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.") @@ -205,31 +221,17 @@ def get_args(): "--sft_hf_dataset", type=str, default=None, - help="Chat dataset to distill on (used with --sft): a Hub dataset id, or 'json' together " - "with --sft_hf_data_files for a local chat jsonl. Each conversation is rendered by the " + help="Chat dataset to distill on (used with --sft): a local chat '.jsonl', or " + "'[:]' (split defaults to train). 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_split", type=str, default="train", help="Split for --sft_hf_dataset" - ) - parser.add_argument( - "--sft_hf_data_files", - type=str, - default=None, - help="Local chat jsonl for --sft_hf_dataset json (HuggingFace load_kwargs.data_files)", - ) - parser.add_argument( - "--sft_hf_validation_split", - type=str, - default=None, - help="Validation split for --sft_hf_dataset (required when --eval_iters > 0)", - ) - parser.add_argument( - "--sft_hf_validation_data_files", + "--sft_hf_validation", type=str, default=None, - help="Local chat jsonl for the validation split of --sft_hf_dataset json", + 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", @@ -417,10 +419,18 @@ def get_args(): "--sft_hf_dataset needs a Megatron-Bridge providing DirectHFSFTDatasetConfig " "(added 2026-07-09). Use a newer Bridge or --sft_dataset_root." ) - if args.eval_iters > 0 and not args.sft_hf_validation_split: - raise ValueError( - "--sft_hf_dataset with --eval_iters > 0 needs --sft_hf_validation_split." - ) + if args.eval_iters > 0 and not args.sft_hf_validation: + raise ValueError("--sft_hf_dataset with --eval_iters > 0 needs --sft_hf_validation.") + # Fail on a mistyped path here rather than after both checkpoints have loaded onto GPUs. + absent = [ + spec + for spec in (args.sft_hf_dataset, args.sft_hf_validation) + if spec and _is_local_json(spec) and not os.path.isfile(spec) + ] + if absent: + raise ValueError(f"--sft_hf_dataset / --sft_hf_validation files missing: {absent}.") + elif args.sft_hf_validation or args.sft_loss_mode != "assistant": + raise ValueError("--sft_hf_validation and --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 " @@ -622,20 +632,11 @@ def _restore_student_hook(model_chunks): 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. - def _hf_source(split, data_files=None): - return HFDatasetSourceConfig( - path_or_dataset=args.sft_hf_dataset, - split=split, - load_kwargs={"data_files": data_files} if data_files else None, - ) - dataset_config = DirectHFSFTDatasetConfig( seq_length=args.seq_length, - source=_hf_source(args.sft_hf_split, args.sft_hf_data_files), + source=_hf_source(args.sft_hf_dataset), validation_source=( - _hf_source(args.sft_hf_validation_split, args.sft_hf_validation_data_files) - if args.sft_hf_validation_split - else None + _hf_source(args.sft_hf_validation) if args.sft_hf_validation else None ), preprocessing=ChatSFTPreprocessingConfig(loss_mode=args.sft_loss_mode), dataloader_type="batch", From 9de33ac932228384b66ce2c6b1947f81d9cdce0c Mon Sep 17 00:00:00 2001 From: Yue Date: Thu, 8 Oct 2026 08:10:49 +0000 Subject: [PATCH 4/5] examples/megatron_bridge: trim the --sft_hf_dataset argument checks Keep only the checks the HF path needs (--sft, a Bridge that has DirectHFSFTDatasetConfig, validation data when evaluating). Drop the local file-existence check, the guard on HF options without --sft_hf_dataset, and the guarded-import comment. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yue --- examples/megatron_bridge/distill.py | 21 ++------------------- 1 file changed, 2 insertions(+), 19 deletions(-) diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 4bd38694243..41617263bec 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -49,8 +49,6 @@ from megatron.bridge.training.post_training.distillation import ModelOptDistillConfig from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size -# DirectHFSFTDatasetConfig landed in Megatron-Bridge on 2026-07-09; an older Bridge still runs every -# other data mode, and --sft_hf_dataset reports why it is unavailable. try: from megatron.bridge.data.builders import DirectHFSFTDatasetConfig from megatron.bridge.data.sft_processing import ChatSFTPreprocessingConfig @@ -128,13 +126,9 @@ def _nonnegative_int(value: str) -> int: return parsed -def _is_local_json(spec: str) -> bool: - return spec.endswith((".json", ".jsonl")) - - def _hf_source(spec: str) -> "HFDatasetSourceConfig": """``.json[l]`` -> local chat jsonl; otherwise ``[:]`` (split defaults to train).""" - if _is_local_json(spec): + if spec.endswith((".json", ".jsonl")): return HFDatasetSourceConfig( path_or_dataset="json", split="train", load_kwargs={"data_files": spec} ) @@ -416,21 +410,10 @@ def get_args(): raise ValueError("--sft_hf_dataset requires --sft.") if not HAS_DIRECT_HF_SFT: raise ValueError( - "--sft_hf_dataset needs a Megatron-Bridge providing DirectHFSFTDatasetConfig " - "(added 2026-07-09). Use a newer Bridge or --sft_dataset_root." + "--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.") - # Fail on a mistyped path here rather than after both checkpoints have loaded onto GPUs. - absent = [ - spec - for spec in (args.sft_hf_dataset, args.sft_hf_validation) - if spec and _is_local_json(spec) and not os.path.isfile(spec) - ] - if absent: - raise ValueError(f"--sft_hf_dataset / --sft_hf_validation files missing: {absent}.") - elif args.sft_hf_validation or args.sft_loss_mode != "assistant": - raise ValueError("--sft_hf_validation and --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 " From 92f80d699a6e62dfff0a51142afbbe1fff7490f1 Mon Sep 17 00:00:00 2001 From: Yue Date: Fri, 9 Oct 2026 05:00:27 +0000 Subject: [PATCH 5/5] examples/megatron_bridge: test the chat SFT path and reject its flags without it - Raise when --sft_hf_validation or a non-default --sft_loss_mode is given without --sft_hf_dataset, instead of silently ignoring them. - Add test_distill_llm_sft_hf_chat (tiny Qwen3, multi-turn chat jsonl, validation) and test_hf_source_spec; both skip when Megatron-Bridge lacks DirectHFSFTDatasetConfig. - Note in --sft_hf_dataset help that only .json/.jsonl are read as local files; --sft help now names both data sources. - Document the chat path in the README and add a CHANGELOG entry. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yue --- CHANGELOG.rst | 1 + examples/megatron_bridge/README.md | 9 +++ examples/megatron_bridge/distill.py | 10 ++- .../examples/megatron_bridge/test_distill.py | 70 +++++++++++++++++++ 4 files changed, 87 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index f287c409c71..60c46edf97d 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -48,6 +48,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 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. - Add ``--mlflow `` to the ``examples/megatron_bridge`` scripts that write a checkpoint -- ``prune_minitron.py``, ``quantize.py``, ``distill.py``, ``export_quantized_megatron_to_hf.py`` and ``export_distilled_megatron_to_hf.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too). Each run records the invocation, its arguments as searchable params and its log -- a distillation records its training metrics instead of the rank-0 log -- and writes ``.experiment.json`` into the checkpoint it produced, so a pruning, a quantization, the distillation that refines its checkpoint and the export that deploys it can be traced to one another; uploading the checkpoints themselves stays off unless ``--mlflow_log_checkpoints`` is passed. diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index fd7c20a499c..f6015fb8d16 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 41617263bec..e92b188020c 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -200,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", @@ -216,7 +217,8 @@ def get_args(): type=str, default=None, help="Chat dataset to distill on (used with --sft): a local chat '.jsonl', or " - "'[:]' (split defaults to train). Each conversation is rendered by the " + "'[:]' (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.', ) @@ -414,6 +416,8 @@ def get_args(): ) 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 " 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