diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md index 99366bec7c6..357f16da7c2 100644 --- a/examples/dataset/MEGATRON_DATA_PREP.md +++ b/examples/dataset/MEGATRON_DATA_PREP.md @@ -4,6 +4,7 @@ | :---: | :---: | :---: | | From JSONL files | Tokenize local JSONL files | \[[Link](#from-jsonl-files)\] | | From Hugging Face Hub | Stream or download HF datasets and tokenize | \[[Link](#from-hugging-face-hub)\] | +| Token-budgeted data blends | Prepare weighted subsets for fast experiments | \[[Link](../researcher_guide/README.md#prepare-token-budgeted-data-blends)\] | | `reasoning_content` for Post-Training v3 | Control how chain-of-thought traces are handled | \[[Link](#reasoning_content-for-post-training-v3-datasets)\] | | Nemotron Pre/Post-Training Datasets | Ready-to-run commands for all Nemotron datasets | \[[Link](#ready-to-run-tokenization-commands)\] | @@ -11,6 +12,9 @@ The distillation and pre-training scripts in Megatron-Bridge or Megatron-LM expe Use the `megatron_preprocess_data` utility to tokenize any JSONL or Hugging Face dataset. The tokenization scripts below print the list of output prefixes (e.g. `tokenized_qwen3/data1_text`) that you can use for the `data_paths` argument (with relative weights on different files) in Megatron training scripts. +For iterative research, use the [token-budgeted data blend workflow](../researcher_guide/README.md#prepare-token-budgeted-data-blends) +to prepare smaller weighted datasets before scaling to a full distillation run. + **Important Notes:** - For Pretraining / raw-text data (`text` key) — use `--append_eod` so Megatron can tell where documents end when concatenating them into long sequences. diff --git a/examples/dataset/prepare_data_blend.py b/examples/dataset/prepare_data_blend.py new file mode 100644 index 00000000000..b5cc2607aa8 --- /dev/null +++ b/examples/dataset/prepare_data_blend.py @@ -0,0 +1,169 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Prepare a weighted Megatron data blend from a YAML configuration.""" + +import argparse +import os +import shutil +from pathlib import Path +from typing import Any, cast + +import yaml +from huggingface_hub import hf_hub_download + +from modelopt.torch.utils.plugins.megatron_preprocess_data import megatron_preprocess_data + + +def load_config(path: Path) -> dict[str, Any]: + """Load a data-blend YAML configuration as a dictionary. + + For example, this YAML:: + + tokenizer: /models/Qwen3-8B + output_dir: /datasets/qwen3-blend + target_tokens: 1000000 # Optional; omit to prepare every source in full. + sources: + - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 + config: Nemotron-SFT-General + split: train + content_field: text + weight: 60 + - hf_dataset: nvidia/Nemotron-SFT-Competitive-Programming-v2 + files: + - data/competitive_programming_python_00.jsonl + content_field: messages + weight: 40 + + returns a dictionary with ``tokenizer``, ``output_dir``, and ``sources`` keys, plus + optional ``target_tokens``. Each source has ``hf_dataset``, ``content_field``, and + ``weight``; it uses ``split`` with optional ``config`` and ``max_samples``, or + selects repository ``files``. + """ + with path.open(encoding="utf-8") as stream: + return cast("dict[str, Any]", yaml.safe_load(stream)) + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", type=Path, required=True, help="Path to the blend YAML file") + return parser + + +def _write_data_blend(path: Path, blend: list[tuple[float, str]]) -> None: + """Write the weighted Megatron dataset paths.""" + content = "\n".join(f"{weight:g} {prefix}" for weight, prefix in blend) + "\n" + path.write_text(content, encoding="utf-8") + + +def _copy_config(source: Path, destination: Path) -> None: + """Copy the input configuration alongside the generated blend.""" + if source.resolve() != destination.resolve(): + shutil.copyfile(source, destination) + + +def _prepare_sources( + sources: list[dict[str, Any]], + output_dir: Path, + tokenizer: str, + total_tokens: int | None, +) -> list[tuple[float, str]]: + """Tokenize all sources and return their weighted output paths.""" + workers = min(32, os.cpu_count() or 1) + blend: list[tuple[float, str]] = [] # (weight, shared .bin/.idx path without extension) + allocated_tokens = 0 + + for index, source in enumerate(sources): + weight = float(source["weight"]) + if total_tokens is None: + source_tokens = None + elif index == len(sources) - 1: + source_tokens = total_tokens - allocated_tokens + else: + source_tokens = round(total_tokens * weight / 100) + allocated_tokens += source_tokens + + dataset = source["hf_dataset"] + source_dir = output_dir / f"{index:02d}_{dataset.replace('/', '--')}" + content_field = source["content_field"] + input_args: dict[str, Any] + if "files" in source: + raw_dir = output_dir.parent / "raw" / dataset.replace("/", "--") + paths = [ + hf_hub_download( + repo_id=dataset, + filename=file, + repo_type="dataset", + local_dir=raw_dir, + ) + for file in source["files"] + ] + input_args = {"jsonl_paths": paths} + else: + input_args = { + "hf_dataset": dataset, + "hf_name": source.get("config"), + "hf_split": source["split"], + "hf_max_samples_per_split": source.get("max_samples"), + "hf_streaming": True, + } + + # Each prefix is the path shared by a tokenized Megatron .bin/.idx file pair. + prefixes = megatron_preprocess_data( + **input_args, + output_dir=source_dir, + tokenizer_name_or_path=tokenizer, + json_keys=content_field, + # Plain text lacks chat-template boundary tokens, so terminate each document with EOS. + append_eod=content_field == "text", + # Join lines in text documents by replacing each newline with a space. + strip_newlines=content_field == "text", + reasoning_content="inline" if content_field == "messages" else "strip", + # Guard against pathological records by capping each tokenized document at 256K tokens. + max_sequence_length=256_000, + max_tokens=source_tokens, + workers=workers, + ) + prefix_weight = weight / len(prefixes) + blend.extend((prefix_weight, prefix) for prefix in prefixes) + + return blend + + +def prepare_data_blend(config_path: Path) -> list[tuple[float, str]]: + """Download and tokenize the configured weighted data sources.""" + config = load_config(config_path) + output_dir = Path(config["output_dir"]) + output_dir.mkdir(parents=True, exist_ok=True) + target_tokens = config.get("target_tokens") + total_tokens = None if target_tokens is None else int(target_tokens) + tokenizer = str(config["tokenizer"]) + + blend = _prepare_sources(config["sources"], output_dir, tokenizer, total_tokens) + _write_data_blend(output_dir / "data_blend.txt", blend) + _copy_config(config_path, output_dir / "config.yaml") + return blend + + +def main() -> None: + """Prepare a data blend from the supplied configuration.""" + parser = _build_parser() + args = parser.parse_args() + blend = prepare_data_blend(args.config) + print(f"Prepared {len(blend)} data paths. See data_blend.txt and config.yaml in the output.") + + +if __name__ == "__main__": + main() diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index 60dfed7fc7d..5e7b7cfa7b4 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -159,6 +159,9 @@ Tensorboard logging is enabled by default and logs are saved to `/te To use Weights & Biases for logging, set the `WANDB_API_KEY` environment variable and pass the `--wandb_project` argument. Optionally, you can also pass `--wandb_entity` and `--wandb_exp_name` arguments to group runs under a project and experiment name. +To measure the initial student's CE and distillation losses, add `--validate_only` to the command. +This skips training and evaluates the student at iteration 0. + To see all available arguments: ```bash @@ -225,6 +228,23 @@ torchrun --nnodes 1 --nproc_per_node 8 distill.py \ `--student_hf_model` should match the base architecture of the student (used as a template for export). For non-Puzzletron (i.e. standard) models, it should be same as `--student_hf_path`. +To export the live student periodically during validation, use `--hf_validation_export_path`: + +```bash +torchrun --nnodes 1 --nproc_per_node 8 distill.py \ + ... \ + --eval_interval 20 \ + --hf_validation_export_path /path/to/save/validation_checkpoints \ + --hf_validation_export_interval 200 \ + --student_hf_model Qwen/Qwen3-4B +``` + +The example validates every 20 steps and exports at iterations 200, 400, and so on. The final +iteration is always exported. `--hf_validation_export_path` and +`--hf_validation_export_interval` must be specified together. These directories contain only +HuggingFace model artifacts and can be evaluated independently while distillation retains its +normal resumable Megatron checkpoints. + **Separate conversion** -- convert any saved iteration using the Megatron-Bridge conversion script: ```bash diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 283504f1ec2..4c3c34e0451 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -24,6 +24,7 @@ import contextlib import os from dataclasses import fields +from pathlib import Path import torch from megatron.bridge import AutoBridge @@ -34,6 +35,7 @@ from megatron.bridge.recipes.utils.optimizer_utils import ( distributed_fused_adam_with_cosine_annealing, ) +from megatron.bridge.training.callbacks import Callback from megatron.bridge.training.config import ( CheckpointConfig, ConfigContainer, @@ -45,10 +47,13 @@ TrainingConfig, ) from megatron.bridge.training.distill import distill +from megatron.bridge.training.gpt_step import forward_step_modelopt from megatron.bridge.training.post_training.checkpointing import has_modelopt_state from megatron.bridge.training.post_training.distillation import ModelOptDistillConfig +from megatron.bridge.training.pretrain import pretrain 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 import modelopt.torch.distill as mtd @@ -159,6 +164,63 @@ def _distill_provide_with_megatron_student( DistillationProvider.provide = _distill_provide_with_megatron_student +class _HFValidationExportCallback(Callback): + """Export the live student to Hugging Face format at selected validation stages.""" + + def __init__( + self, + export_dir: str, + student_hf_model: str, + student_hf_path: str, + trust_remote_code: bool, + export_interval: int, + ) -> None: + self.export_dir = Path(export_dir) + self.student_hf_path = student_hf_path + self.trust_remote_code = trust_remote_code + self.export_interval = export_interval + self._last_exported_iteration: int | None = None + self.bridge = AutoBridge.from_hf_pretrained( + student_hf_model, trust_remote_code=trust_remote_code + ) + + def on_eval_end(self, context) -> None: + """Export the student at the iteration that was just validated.""" + iteration = context.state.train_state.step + train_iters = context.state.cfg.train.train_iters + if iteration % self.export_interval != 0 and iteration != train_iters: + return + # The final iteration can be validated both on its regular interval and after training. + # Avoid exporting and overwriting the same Hugging Face checkpoint twice. + if iteration == self._last_exported_iteration: + return + output_path = self.export_dir / f"iter_{iteration:07d}" + print_rank_0(f"Exporting validation checkpoint {iteration} to {output_path}") + + # DistillationModel is the student with teacher and KD-loss modules attached. Hide the + # auxiliary modules temporarily so the Hugging Face export contains only student weights. + with contextlib.ExitStack() as stack: + for model_chunk in unwrap_model(context.model): + if isinstance(model_chunk, mtd.DistillationModel): + stack.enter_context(model_chunk.hide_teacher_model()) + stack.enter_context(model_chunk.hide_loss_modules()) + self.bridge.save_hf_pretrained( + context.model, + output_path, + show_progress=True, + strict=True, + ) + + if dist.rank() == 0: + # Preserve the student architecture from student_hf_path, including heterogeneous + # layer changes; AutoConfig supports both local paths and Hugging Face model IDs. + AutoConfig.from_pretrained( + self.student_hf_path, trust_remote_code=self.trust_remote_code + ).save_pretrained(output_path) + torch.distributed.barrier() + self._last_exported_iteration = iteration + + def get_args(): """Parse command-line arguments.""" parser = argparse.ArgumentParser(description="Distillation for Megatron-Bridge.") @@ -233,6 +295,11 @@ def get_args(): parser.add_argument("--lr", type=float, default=1e-4, help="Peak learning rate") parser.add_argument("--min_lr", type=float, default=1e-5, help="Minimum learning rate") parser.add_argument("--lr_warmup_iters", type=int, default=50, help="Number of LR warmup steps") + parser.add_argument( + "--reset_optimizer", + action="store_true", + help="Do not restore optimizer or scheduler state when resuming from a checkpoint", + ) parser.add_argument( "--recompute_granularity", type=str, @@ -267,6 +334,11 @@ def get_args(): parser.add_argument( "--eval_iters", type=int, default=32, help="Number of batches per validation stage" ) + parser.add_argument( + "--validate_only", + action="store_true", + help="Run validation on the initial student without training", + ) # Logging arguments parser.add_argument("--log_interval", type=int, default=10, help="Write to log every steps") parser.add_argument( @@ -284,6 +356,24 @@ def get_args(): "If provided, exports last iteration checkpoint to HF format after distillation." ), ) + parser.add_argument( + "--hf_validation_export_path", + type=str, + default=None, + help=( + "Path where HuggingFace checkpoints are exported after each validation stage. " + "Each checkpoint is saved in an iter_ subdirectory." + ), + ) + parser.add_argument( + "--hf_validation_export_interval", + type=int, + default=None, + help=( + "Export a HuggingFace checkpoint every training steps. " + "The final iteration is always exported." + ), + ) parser.add_argument( "--student_hf_model", type=str, @@ -298,8 +388,19 @@ def get_args(): if not args.use_mock_data and not args.data_paths: raise ValueError("Must provide either --data_paths or set --use_mock_data.") - if args.hf_export_path and not args.student_hf_model: - raise ValueError("Must provide --student_hf_model if --hf_export_path is provided.") + if args.validate_only and args.hf_export_path: + raise ValueError("--hf_export_path cannot be used with --validate_only.") + if (args.hf_export_path or args.hf_validation_export_path) and not args.student_hf_model: + raise ValueError( + "Must provide --student_hf_model when HuggingFace checkpoint export is enabled." + ) + if args.hf_validation_export_interval is not None and not args.hf_validation_export_path: + raise ValueError("--hf_validation_export_interval requires --hf_validation_export_path.") + if args.hf_validation_export_path: + if args.hf_validation_export_interval is None: + raise ValueError( + "--hf_validation_export_path requires --hf_validation_export_interval." + ) print_args(args) @@ -400,6 +501,7 @@ def _build_model_provider(hf_path, load_weights=True): train_iters=args.train_iters, eval_interval=args.eval_interval, eval_iters=args.eval_iters, + skip_train=args.validate_only, global_batch_size=args.gbs, micro_batch_size=args.mbs, manual_gc=True, @@ -433,7 +535,8 @@ def _build_model_provider(hf_path, load_weights=True): checkpoint=CheckpointConfig( save_interval=args.eval_interval, save=checkpoint_dir, - load=checkpoint_dir, # Resume from this directory (if exists) + load=None if args.validate_only else checkpoint_dir, + load_optim=not args.reset_optimizer, most_recent_k=5, # Keeps 5 most recent checkpoints (not metric-based) ckpt_format="torch_dist", async_save=True, @@ -444,12 +547,32 @@ def _build_model_provider(hf_path, load_weights=True): ) print_rank_0("\nStarting distillation...") - distill(config) - print_rank_0( - f"\nDistillation done! Saved checkpoint to {checkpoint_dir}" - " in megatron distributed checkpoint format.\n" - ) - + if args.hf_validation_export_path: + assert isinstance(config.model, DistillationProvider), ( + "Distillation requires a DistillationProvider" + ) + callback = _HFValidationExportCallback( + export_dir=args.hf_validation_export_path, + student_hf_model=args.student_hf_model, + student_hf_path=args.student_hf_path, + trust_remote_code=args.trust_remote_code, + export_interval=args.hf_validation_export_interval, + ) + # TODO: Use distill(..., callbacks=[callback]) once Megatron-Bridge supports callbacks. + pretrain(config, forward_step_modelopt, callbacks=[callback]) + else: + distill(config) + if args.validate_only: + print_rank_0("\nValidation-only run done!\n") + else: + print_rank_0( + f"\nDistillation done! Saved checkpoint to {checkpoint_dir}" + " in megatron distributed checkpoint format.\n" + ) + + # TODO: Extend _HFValidationExportCallback to export --hf_export_path from the in-memory + # student at the end of training. This avoids reloading the newly saved Megatron checkpoint + # and duplicating large disk I/O. if args.hf_export_path: print_rank_0(f"Exporting final distilled ckpt to HF format to {args.hf_export_path}") # Save rank before destroying process group (dist.rank() won't work after destruction) @@ -470,7 +593,8 @@ def _build_model_provider(hf_path, load_weights=True): show_progress=True, strict=True, ) - # Copy config.json from student_hf_path (handles both local paths and HF model IDs) + # Preserve the student architecture from student_hf_path, including heterogeneous + # layer changes; AutoConfig supports both local paths and Hugging Face model IDs. AutoConfig.from_pretrained( args.student_hf_path, trust_remote_code=args.trust_remote_code ).save_pretrained(args.hf_export_path) diff --git a/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md b/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md index 22a31ff9db7..a19511af1e7 100644 --- a/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md +++ b/examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md @@ -69,7 +69,14 @@ Distillation uses the **30% Pretraining (Code 5, General 20, MATH 5) + 70% Post- ### 1. Data Preparation -See [examples/dataset/MEGATRON_DATA_PREP.md](../../../dataset/MEGATRON_DATA_PREP.md) for tokenization commands for all datasets used in this blend. +Prepare this blend with the +[token-budgeted data blend workflow](../../../researcher_guide/README.md#prepare-token-budgeted-data-blends). +The complete blend listed below contains approximately 142B tokens. For an initial experiment, set +`target_tokens: 1000000000` to prepare a 1B-token subset with the same source weights, avoiding the time and +storage needed to preprocess the complete blend. Omit `target_tokens` to prepare every configured source in +full, subject to any per-source `max_samples` setting. See +[examples/dataset/MEGATRON_DATA_PREP.md](../../../dataset/MEGATRON_DATA_PREP.md) for additional dataset +tokenization commands. For this experiment: `TOKENIZER=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16`, `OUTPUT_DIR=tokenized_nemotron_3`. diff --git a/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md b/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md index 55d4706175d..5fab5aebbee 100644 --- a/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md +++ b/examples/pruning/minitron/NVIDIA-Nemotron-Nano-9B-v2/README.md @@ -63,6 +63,10 @@ Distillation uses the **30% Pretraining (Code 5, General 20, MATH 5) + 70% Post- ### 1. Data Preparation See [examples/dataset/MEGATRON_DATA_PREP.md](../../../dataset/MEGATRON_DATA_PREP.md) for tokenization commands for all datasets used in this blend. +To prepare a token-limited subset, follow the +[token-budgeted data blend workflow](../../../researcher_guide/README.md#prepare-token-budgeted-data-blends), +but create a custom YAML configuration using this tutorial's tokenizer, sources, and weights below. The +researcher guide's example configuration targets Nemotron 3 and should not be reused unchanged. For this experiment: `TOKENIZER=nvidia/NVIDIA-Nemotron-Nano-9B-v2`, `OUTPUT_DIR=tokenized_nemotron_v2`. diff --git a/examples/researcher_guide/README.md b/examples/researcher_guide/README.md new file mode 100644 index 00000000000..00b009be685 --- /dev/null +++ b/examples/researcher_guide/README.md @@ -0,0 +1,263 @@ +# ModelOpt for Researchers: Fast Experimentation Workflows + +Model optimization research depends on short feedback loops: test a hypothesis cheaply, compare candidates +reproducibly, and spend full-scale compute only on the most promising experiments. This guide collects practical +ModelOpt workflows for that iterative research process. + +Current workflows include: + +- [Efficient model evaluation](#efficient-evaluation-with-lm-eval-harness) with smaller benchmark subsets. +- [Downstream evaluation over time during distillation](#track-downstream-quality-over-time-during-distillation) + with validation checkpoints. +- [Efficient data blend preparation](#prepare-token-budgeted-data-blends) for distillation experiments. + +The guide will grow as additional research workflows are documented. It complements the feature-specific +[examples](../) by connecting them into experimentation strategies rather than replacing their detailed +instructions. + +## Efficient evaluation with LM-Eval Harness + +[LM-Eval Harness](../llm_eval/README.md) supports many accuracy benchmarks, but full runs are often too slow for +every iteration of model pruning, distillation, or quantization. Use progressively larger evaluation subsets to +reject weak candidates quickly and reserve full runs for the most promising models. + +In LM-Eval, `--limit N` evaluates the first `N` samples of each individual task. For task groups such as MMLU and +MMLU-Pro, the limit applies to every subject, not to the group as a whole. + +The following table gives a practical progression for LM-Eval's MMLU-Pro task group, which contains 14 subjects +and 12,032 questions. Example times assume Qwen3-8B, a batch size of 4, and subject-level parallelism on eight +H100 80GB GPUs: + +| Limit per subject | Questions evaluated | Worst-case 95% margin of error | Example time | +|-------------------|--------------------:|--------------------------------:|-------------:| +| `10` | 140 | ±8.3 percentage points | ~3 minutes | +| `50` | 700 | ±3.7 percentage points | ~14 minutes | +| `100` | 1,400 | ±2.6 percentage points | ~28 minutes | +| `200` | 2,800 | ±1.9 percentage points | ~56 minutes | +| None | 12,032 | ±0.9 percentage points | 4 hours | + +The example times scale an approximately four-hour full run by the fraction of questions evaluated. Actual time +depends on the model, hardware, batch size, and parallelism. + +The margins of error are conservative planning estimates. They use 50% accuracy, the normal approximation for a +[binomial proportion confidence interval](https://en.wikipedia.org/wiki/Binomial_proportion_confidence_interval#Normal_approximation_interval). + +These estimates treat benchmark questions as independent random samples from a broader population of possible +questions. Because `--limit` selects the first samples, limited scores may also be affected by dataset ordering +and should not be reported as final benchmark results. + +Add `--log_samples` for paired per-question analysis. When multiple GPUs are available, use data parallelism to +split samples across model copies; see the [LM-Eval examples](../llm_eval/README.md) for commands. + +## Track downstream quality over time during distillation + +Validation KD and CE losses show whether the student is fitting the teacher and validation data, but they do not +necessarily predict downstream accuracy. Export the live student at each validation interval and evaluate the +resulting checkpoints to see when downstream quality improves, plateaus, or regresses. + +Use `--hf_validation_export_path` to choose the output directory and +`--hf_validation_export_interval` to choose how often `distill.py` saves a Hugging Face checkpoint, as described in +the [Megatron-Bridge distillation guide](../megatron_bridge/README.md#converting-to-hugging-face-format-optional). + +The export path contains one loadable Hugging Face checkpoint per validation iteration: + +```text +hf_validation/ +├── iter_0000100/ +├── iter_0000200/ +└── iter_0000300/ +``` + +Evaluate the teacher, pruned student, and each exported checkpoint. Follow the +[LM-Eval Harness instructions](../llm_eval/README.md#lm-eval-harness) and use the +[efficient evaluation workflow](#efficient-evaluation-with-lm-eval-harness) to choose limits. + +The following experiment pruned Qwen3-8B to 0.7x and distilled the same student for approximately 100 million +tokens using four data recipes. All runs used a global batch size of 8 and sequence length of 4,096. MMLU used 25 +questions per subject (1,425 total), and MMLU-Pro used 50 per subject (700 total). The tables show representative +checkpoints; token counts are derived from the consumed fixed-length training sequences. + +Measured compute per data recipe on eight H100 80GB GPUs: + +| Stage | Checkpoints | Time | GPU use | +|-------|------------:|-----:|---------| +| Distillation to 100M tokens | - | ~2h10m | 8 GPUs | +| MMLU trajectory | 21 | ~51m | 8 GPUs | +| MMLU-Pro trajectory | 13 | ~3h50m | Two checkpoints in parallel, 4 GPUs each | +| Total | - | ~6h50m | Excludes Slurm queue and worker setup | + +| Baseline model | MMLU | MMLU-Pro | +|----------------|-----:|---------:| +| Teacher: Qwen3-8B | 74.93% (full) | 58.62% (full) | +| Pruned 0.7x student | 48.69% (full) | 23.09% (full) | + +### WikiText + +- Dataset: Salesforce/wikitext (`wikitext-103-v1`) +- Teacher CE: 2.6834 + +| Training tokens | Validation KD | Validation CE | MMLU | MMLU-Pro | +|----------------:|--------------:|--------------:|-----:|---------:| +| 0 | 0.8261 | 3.3458 | 48.69% (full) | 23.09% (full) | +| 0.7M | 0.3031 | 2.6570 | 59.72% | 25.00% | +| 3.3M | 0.2343 | 2.6091 | 63.58% | 29.29% | +| 39.3M | 0.1495 | 2.5665 | 65.89% | 39.86% | +| 78.6M | 0.1315 | 2.5699 | 66.46% | 39.57% | +| 100.0M | 0.1291 | 2.5863 | 67.30% | 40.57% | + +### Nemotron v2 + +- Dataset: nvidia/Nemotron-Post-Training-Dataset-v2 (math and stem) +- Teacher CE: 1.1566 + +| Training tokens | Validation KD | Validation CE | MMLU | MMLU-Pro | +|----------------:|--------------:|--------------:|-----:|---------:| +| 0 | 0.5187 | 1.4739 | 48.69% (full) | 23.09% (full) | +| 0.7M | 0.1919 | 1.0931 | 58.74% | 13.14% | +| 3.3M | 0.1342 | 1.0550 | 60.56% | 14.29% | +| 39.3M | 0.0675 | 1.0296 | 64.63% | 6.14% | +| 78.6M | 0.0613 | 1.0773 | 65.61% | 7.71% | +| 100.0M | 0.0582 | 1.0516 | 65.75% | 11.29% | + +### 50/50 WikiText and Nemotron v2 blend + +- Dataset: 50/50 blend of WikiText and Nemotron v2 math and stem +- Teacher CE: 1.9025 + +| Training tokens | Validation KD | Validation CE | MMLU | MMLU-Pro | +|----------------:|--------------:|--------------:|-----:|---------:| +| 0 | 0.6662 | 2.3780 | 48.69% (full) | 23.09% (full) | +| 0.7M | 0.2479 | 1.8363 | 57.89% | 12.57% | +| 3.3M | 0.1824 | 2.0265 | 62.46% | 23.14% | +| 39.3M | 0.1164 | 1.9157 | 67.44% | 33.86% | +| 78.6M | 0.0973 | 1.8503 | 67.72% | 41.57% | +| 100.0M | 0.0916 | 1.7680 | 68.28% | 41.71% | + +### Nemotron 3 + +- Dataset: Nemotron 3 Nano [distillation blend](#prepare-token-budgeted-data-blends) +- Teacher CE: 1.4702 + +| Training tokens | Validation KD | Validation CE | MMLU | MMLU-Pro | +|----------------:|--------------:|--------------:|-----:|---------:| +| 0 | 0.6395 | 1.9113 | 48.69% (full) | 23.09% (full) | +| 0.7M | 0.2424 | 1.5910 | 57.05% | 24.86% | +| 3.3M | 0.1604 | 1.5190 | 62.46% | 36.86% | +| 39.3M | 0.0978 | 1.4144 | 67.23% | 45.00% | +| 78.6M | 0.0890 | 1.4112 | 67.93% | 47.14% | +| 100.0M | 0.0845 | 1.4656 | 67.37% | 47.71% | + +Interesting observations include: + +- All four data recipes recover MMLU to about 66% to 68% by 100 million tokens. The 50/50 blend is numerically + highest at 68.28%. +- Nemotron 3 produces the strongest MMLU-Pro trajectory, reaching 47.71%. +- Although Nemotron v2 performs poorly alone, its 50/50 blend with WikiText slightly outperforms WikiText alone + on both final benchmarks. +- Nemotron v2 KD continues to decrease, but its MMLU-Pro score remains below the pruned baseline. Understanding + this discrepancy could help design better distillation corpora. + +## Prepare token-budgeted data blends + +Full distillation datasets are often unnecessarily large for testing a pruning or distillation hypothesis. Use +[`prepare_data_blend.py`](../dataset/prepare_data_blend.py) to prepare a smaller weighted blend with a shared token +budget. The utility supports Hugging Face configurations and splits as well as specific JSONL files stored in a +Hugging Face dataset repository. + +Define the tokenizer, output directory, and source weights in YAML. Set the optional `target_tokens` field to +prepare a weighted subset, or omit it to prepare every source in full. This example scales the +[Nemotron 3 Nano distillation blend](../megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md#1-data-preparation) +down to one billion tokens while preserving its source weights: + +```yaml +tokenizer: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 +output_dir: /path/to/nemotron_3_nano_30b_distillation_blend_1b +# Optional; omit this field to prepare every source in full. +target_tokens: 1000000000 +sources: + - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 + config: Nemotron-SFT-Code + split: train + max_samples: 10000000 + content_field: text + weight: 5 + - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 + config: Nemotron-SFT-General + split: train + max_samples: 10000000 + content_field: text + weight: 20 + - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 + config: Nemotron-SFT-MATH + split: train + max_samples: 10000000 + content_field: text + weight: 5 + - hf_dataset: nvidia/Nemotron-Math-v2 + split: high_part00 + content_field: messages + weight: 10 + - hf_dataset: nvidia/Nemotron-SFT-Math-v3 + files: + - data/train.jsonl + content_field: messages + weight: 17 + - hf_dataset: nvidia/Nemotron-SFT-Competitive-Programming-v2 + files: + - data/competitive_programming_python_00.jsonl + content_field: messages + weight: 15 + - hf_dataset: nvidia/Nemotron-SFT-Competitive-Programming-v2 + files: + - data/competitive_programming_cpp_00.jsonl + content_field: messages + weight: 5 + - hf_dataset: nvidia/Nemotron-Post-Training-Dataset-v1 + config: default + split: stem + max_samples: 5000000 + content_field: messages + weight: 8 + - hf_dataset: nvidia/Nemotron-Science-v1 + files: + - data/MCQ.jsonl + content_field: messages + weight: 3 + - hf_dataset: nvidia/Nemotron-Science-v1 + files: + - data/RQA.jsonl + content_field: messages + weight: 2 + - hf_dataset: nvidia/Nemotron-SFT-Instruction-Following-Chat-v2 + files: + - data/reasoning_on.jsonl + content_field: messages + weight: 3 + - hf_dataset: nvidia/Nemotron-SFT-Instruction-Following-Chat-v2 + files: + - data/reasoning_off.jsonl + content_field: messages + weight: 2 + - hf_dataset: nvidia/Nemotron-Agentic-v1 + files: + - data/tool_calling.jsonl + content_field: messages + weight: 5 +``` + +Run from the repository root: + +```bash +python examples/dataset/prepare_data_blend.py --config blend.yaml +``` + +The output contains tokenized Megatron `.bin`/`.idx` files, `data_blend.txt` with the weighted paths for training, +and `config.yaml` recording how the blend was generated. The final token count can slightly exceed the target +because the final document from each source is kept whole. See the +[Megatron data preparation guide](../dataset/MEGATRON_DATA_PREP.md) for dataset-specific details. + +## Planned topics + +Future additions can cover: + +- Iterative pruning and distillation workflows diff --git a/modelopt/torch/utils/plugins/megatron_preprocess_data.py b/modelopt/torch/utils/plugins/megatron_preprocess_data.py index 52dae7c140b..04e59b0b6ae 100644 --- a/modelopt/torch/utils/plugins/megatron_preprocess_data.py +++ b/modelopt/torch/utils/plugins/megatron_preprocess_data.py @@ -85,6 +85,7 @@ import argparse import gzip +import itertools import json import multiprocessing import time @@ -261,7 +262,26 @@ def _print_processing_stats( flush=True, ) - def _encode_docs(self, encoder: "_Encoder", lines): + @staticmethod + def _encode_in_batches(pool, encoder: "_Encoder", lines, batch_size: int): + """Encode finite batches for a run bounded by ``max_tokens``. + + Once the token target is reached, the caller stops without consuming the remaining input. + Batches run one at a time, while documents within each batch are encoded in parallel. Unlike + ``imap``, ``map`` collects every result from the current batch before returning. Therefore, + no worker is writing a pending result when the caller stops at the token target. + The final batch may encode documents that the caller does not use after reaching its limit. + """ + lines = iter(lines) + while True: + batch = list(itertools.islice(lines, batch_size)) + if not batch: + break + + encoded_batch = pool.map(encoder.encode, batch, chunksize=1) + yield from encoded_batch + + def _encode_docs(self, encoder: "_Encoder", lines, may_stop_early: bool = False): """Tokenize ``lines``, forking worker processes only when ``workers > 1``. ``multiprocessing.Pool`` always ``fork()``s, even for a single worker. Forking a @@ -269,16 +289,28 @@ def _encode_docs(self, encoder: "_Encoder", lines): this is called in-process after GPU work) is unsafe and can segfault the children. The single-worker path avoids the fork entirely by tokenizing inline in this process. + When ``may_stop_early`` is true, wait for finite batches so no worker results are pending + if the caller stops consuming documents after reaching its token limit. + Returns ``(pool, encoded_docs)``; ``pool`` is ``None`` in the inline case. """ if self.workers == 1: encoder.initializer() return None, map(encoder.encode, lines) pool = multiprocessing.Pool(self.workers, initializer=encoder.initializer) - return pool, pool.imap(encoder.encode, lines, 32) + if may_stop_early: + batch_size = self.workers * 4 # Balance throughput against unused final-batch work. + encoded_docs = self._encode_in_batches(pool, encoder, lines, batch_size) + else: + encoded_docs = pool.imap(encoder.encode, lines, 32) + return pool, encoded_docs def process_json_file( - self, input_file_name: str | Path, output_dir: str | Path, encoder: _Encoder + self, + input_file_name: str | Path, + output_dir: str | Path, + encoder: _Encoder, + max_tokens: int | None = None, ) -> tuple[int, list[str]]: input_path = Path(input_file_name) stem = input_path.stem if input_path.suffix != ".gz" else Path(input_path.stem).stem @@ -291,7 +323,8 @@ def process_json_file( else: fin = open(input_path, encoding="utf-8") - pool, encoded_docs = self._encode_docs(encoder, fin) + # Workers encode asynchronously; iterating encoded_docs waits for results in input order. + pool, encoded_docs = self._encode_docs(encoder, fin, may_stop_early=max_tokens is not None) output_bin_files = {} output_idx_files = {} @@ -312,7 +345,8 @@ def process_json_file( return 0, prefixes start_time = time.time() - total_doc_len, total_enc_len, final_enc_len = 0, 0, 0 + total_doc_len, total_enc_len = 0, 0 + final_enc_len = 0 # Tokens written to output, including appended EOD tokens. for i, (doc, sentence_lens, (doc_len, enc_len)) in enumerate(encoded_docs, start=1): total_doc_len += doc_len total_enc_len += enc_len @@ -320,6 +354,8 @@ def process_json_file( for key in doc: builders[key].add_document(doc[key], sentence_lens[key]) self._print_processing_stats(i, total_doc_len, total_enc_len, start_time) + if max_tokens is not None and final_enc_len >= max_tokens: + break self._print_processing_stats(i, total_doc_len, total_enc_len, start_time, force_print=True) fin.close() @@ -345,6 +381,7 @@ def process_hf_split( config: str | None, split: str, max_samples: int | None = None, + max_tokens: int | None = None, streaming: bool = False, ) -> tuple[int, list[str]]: """Load a HF dataset split and tokenize directly without writing an intermediate JSONL. @@ -357,11 +394,12 @@ def process_hf_split( """ print(f"\nLoading HF dataset {dataset_name=}, {config=}, {split=}, {streaming=}") ds = load_dataset(path=dataset_name, name=config, split=split, streaming=streaming) - if max_samples is not None: + if max_samples is not None or max_tokens is not None: # Shuffle first so the selected subset is random, not a biased prefix. # Non-streaming: global index shuffle (memory-mapped, efficient) then .select(N). # Streaming: buffer shuffle (approximate) then .take(N). ds = ds.shuffle(seed=42) + if max_samples is not None: if streaming: ds = ds.take(max_samples) else: @@ -378,6 +416,7 @@ def process_hf_split( safe_name = dataset_name.replace("/", "--") sample_tag = f"_max{max_samples}" if max_samples is not None else "" + token_tag = f"_tokens{max_tokens}" if max_tokens is not None else "" output_prefix = Path(output_dir) / f"{safe_name}_{config}_{split}" prefixes = [] @@ -385,7 +424,7 @@ def process_hf_split( output_idx_files = {} builders = {} for key in self.json_keys: - prefix = f"{output_prefix}_{key}{sample_tag}" + prefix = f"{output_prefix}_{key}{sample_tag}{token_tag}" prefixes.append(prefix) output_bin_files[key] = f"{prefix}.bin" output_idx_files[key] = f"{prefix}.idx" @@ -400,10 +439,14 @@ def process_hf_split( print(f"\t[SKIP] Output files for {dataset_name} {config}/{split} already exist") return 0, prefixes - pool, encoded_docs = self._encode_docs(encoder, self._iter_hf_as_json(ds)) + # Workers encode asynchronously; iterating encoded_docs waits for results in input order. + pool, encoded_docs = self._encode_docs( + encoder, self._iter_hf_as_json(ds), may_stop_early=max_tokens is not None + ) start_time = time.time() - total_doc_len, total_enc_len, final_enc_len = 0, 0, 0 + total_doc_len, total_enc_len = 0, 0 + final_enc_len = 0 # Tokens written to output, including appended EOD tokens. i = 0 for i, (doc, sentence_lens, (doc_len, enc_len)) in enumerate(encoded_docs, start=1): total_doc_len += doc_len @@ -412,6 +455,8 @@ def process_hf_split( for key in doc: builders[key].add_document(doc[key], sentence_lens[key]) self._print_processing_stats(i, total_doc_len, total_enc_len, start_time) + if max_tokens is not None and final_enc_len >= max_tokens: + break if i: self._print_processing_stats( @@ -462,6 +507,7 @@ def megatron_preprocess_data( hf_split: str | None = None, hf_max_samples_per_split: int | None = None, hf_streaming: bool = False, + max_tokens: int | None = None, # Other arguments output_dir: str | Path, tokenizer_name_or_path: str, @@ -488,6 +534,8 @@ def megatron_preprocess_data( downloaded — useful for very large pretraining datasets or datasets with complex nested message schemas that cause Arrow type-cast errors in non-streaming mode. Note: streaming does not cache to disk, so re-runs re-download. Defaults to False. + max_tokens: Stop after processing at least this many tokens across the source files or + selected Hugging Face split. The final document may make the result slightly larger. output_dir: Path to directory to save binary output files. tokenizer_name_or_path: Name or path of the Hugging Face tokenizer to use. json_keys: Key or list of keys to extract from json. Defaults to ["text"]. @@ -515,11 +563,16 @@ def megatron_preprocess_data( raise ValueError( "Exactly one of `input_dir`, `jsonl_paths`, or `hf_dataset` must be provided." ) - if hf_streaming and hf_max_samples_per_split is None and _is_main_or_first_worker(): + if ( + hf_streaming + and hf_max_samples_per_split is None + and max_tokens is None + and _is_main_or_first_worker() + ): warnings.warn( - "--hf_streaming is set but --hf_max_samples_per_split is not. " - "Streaming without a sample cap re-downloads the full dataset on every run with no " - "disk cache, which is slower than the cached non-streaming path.", + "--hf_streaming is set but neither --hf_max_samples_per_split nor --max_tokens is " + "set. Streaming without a sample or token cap re-downloads the full dataset on " + "every run with no disk cache, which is slower than the cached non-streaming path.", stacklevel=2, ) @@ -536,12 +589,16 @@ def megatron_preprocess_data( ) partition = _Partition(vocab_size, json_keys, log_interval, workers) + # Tokens written across all input files or Hugging Face splits. final_enc_len = 0 all_prefixes: list[str] = [] overall_start = time.time() if hf_dataset is not None: for config, split in _enumerate_hf_splits(hf_dataset, hf_name, hf_split): + remaining_tokens = None if max_tokens is None else max_tokens - final_enc_len + if remaining_tokens is not None and remaining_tokens <= 0: + break enc_len, prefixes = partition.process_hf_split( output_dir, encoder, @@ -549,6 +606,7 @@ def megatron_preprocess_data( config, split, hf_max_samples_per_split, + remaining_tokens, hf_streaming, ) final_enc_len += enc_len @@ -566,7 +624,12 @@ def megatron_preprocess_data( file_names = list(jsonl_paths) # type: ignore[arg-type] for name in file_names: - enc_len, prefixes = partition.process_json_file(name, output_dir, encoder) + remaining_tokens = None if max_tokens is None else max_tokens - final_enc_len + if remaining_tokens is not None and remaining_tokens <= 0: + break + enc_len, prefixes = partition.process_json_file( + name, output_dir, encoder, remaining_tokens + ) final_enc_len += enc_len all_prefixes.extend(prefixes) @@ -633,6 +696,12 @@ def main(): parser.add_argument( "--max_sequence_length", type=int, default=None, help="Maximum sequence length" ) + parser.add_argument( + "--max_tokens", + type=int, + default=None, + help="Stop after processing at least this many tokens", + ) parser.add_argument("--workers", type=int, default=8, help="Number of worker processes") parser.add_argument("--log_interval", type=int, default=100000, help="Log interval") parser.add_argument( @@ -675,6 +744,7 @@ def main(): json_keys=args.json_keys, append_eod=args.append_eod, max_sequence_length=args.max_sequence_length, + max_tokens=args.max_tokens, workers=args.workers, log_interval=args.log_interval, reasoning_content=args.reasoning_content, diff --git a/tests/examples/dataset/test_prepare_data_blend.py b/tests/examples/dataset/test_prepare_data_blend.py new file mode 100644 index 00000000000..c71f1c0b801 --- /dev/null +++ b/tests/examples/dataset/test_prepare_data_blend.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import sys +from pathlib import Path +from unittest.mock import Mock + +import pytest +import yaml + +# examples/dataset is not a package; add it to the path to import and test the script in-process. +sys.path.insert(0, str(Path(__file__).parents[3] / "examples/dataset")) + +import prepare_data_blend + + +def _setup_test( + tiny_qwen3_path: str, tmp_path: Path, monkeypatch, target_tokens: int | None +) -> tuple[Path, Path, Mock]: + output_dir = tmp_path / "tokenized" + config = { + "tokenizer": tiny_qwen3_path, + "output_dir": str(output_dir), + "sources": [ + { + "hf_dataset": "nanotron/minipile_100_samples", + "split": "train", + "max_samples": 100, + "content_field": "text", + "weight": 60, + }, + { + "hf_dataset": "nvidia/Nemotron-SFT-Competitive-Programming-v2", + "files": ["data/competitive_programming_python_00.jsonl"], + "content_field": "messages", + "weight": 40, + }, + ], + } + if target_tokens is not None: + config["target_tokens"] = target_tokens + config_path = tmp_path / "config.yaml" + config_path.write_text(yaml.safe_dump(config), encoding="utf-8") + jsonl_path = tmp_path / "competitive_programming_python_00.jsonl" + conversation = { + "messages": [ + {"role": "user", "content": "Write a Python function that adds two integers."}, + {"role": "assistant", "content": "def add(a, b):\n return a + b"}, + ] + } + jsonl_path.write_text( + "".join(json.dumps(conversation) + "\n" for _ in range(20)), encoding="utf-8" + ) + download = Mock(return_value=str(jsonl_path)) + monkeypatch.setattr(prepare_data_blend, "hf_hub_download", download) + monkeypatch.setattr(sys, "argv", ["prepare_data_blend.py", "--config", str(config_path)]) + return output_dir, config_path, download + + +@pytest.mark.parametrize("target_tokens", [1_000, None], ids=["token-budget", "all-data"]) +def test_prepare_data_blend_with_split_and_files_sources( + tiny_qwen3_path: str, tmp_path: Path, monkeypatch, target_tokens: int | None +): + output_dir, config_path, download = _setup_test( + tiny_qwen3_path, tmp_path, monkeypatch, target_tokens + ) + + # Run in-process so the mocked NVIDIA download is visible; run_example_command uses a subprocess. + prepare_data_blend.main() + + download.assert_called_once_with( + repo_id="nvidia/Nemotron-SFT-Competitive-Programming-v2", + filename="data/competitive_programming_python_00.jsonl", + repo_type="dataset", + local_dir=tmp_path / "raw/nvidia--Nemotron-SFT-Competitive-Programming-v2", + ) + blend = [ + line.split(maxsplit=1) + for line in (output_dir / "data_blend.txt").read_text(encoding="utf-8").splitlines() + ] + assert [weight for weight, _ in blend] == ["60", "40"] + for _, prefix in blend: + assert Path(prefix + ".bin").exists() + assert Path(prefix + ".idx").exists() + assert ("_tokens" in prefix) is (target_tokens is not None) + assert (output_dir / "config.yaml").read_bytes() == config_path.read_bytes() diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py index 8b5e206567d..f9b90e80edd 100644 --- a/tests/examples/megatron_bridge/test_distill.py +++ b/tests/examples/megatron_bridge/test_distill.py @@ -25,9 +25,10 @@ def test_distill_and_convert(tmp_path: Path, num_gpus): teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True) - train_iters = 2 + train_iters = 3 distill_output_dir = tmp_path / "distill_output" distilled_hf_path = tmp_path / "distilled_hf" + validation_exports = tmp_path / "validation_exports" distill_cmd_parts = extend_cmd_parts( ["torchrun", f"--nproc_per_node={num_gpus}", "distill.py", "--use_mock_data"], student_hf_path=teacher_hf_path, @@ -40,16 +41,53 @@ def test_distill_and_convert(tmp_path: Path, num_gpus): gbs=4, train_iters=train_iters, lr_warmup_iters=1, - eval_interval=train_iters, + eval_interval=1, eval_iters=1, log_interval=1, hf_export_path=distilled_hf_path, + hf_validation_export_path=validation_exports, + hf_validation_export_interval=2, student_hf_model=teacher_hf_path, ) run_example_command(distill_cmd_parts, example_path="megatron_bridge") assert (distill_output_dir / f"checkpoints/iter_{train_iters:07d}").exists() assert (distilled_hf_path / "config.json").exists() + assert (validation_exports / "iter_0000002/config.json").exists() + assert (validation_exports / "iter_0000003/config.json").exists() + assert not (validation_exports / "iter_0000001").exists() + + +def test_distill_validate_only(tmp_path: Path, num_gpus): + teacher_hf_path = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True) + output_dir = tmp_path / "validation_output" + validation_exports = tmp_path / "validation_exports" + cmd_parts = extend_cmd_parts( + [ + "torchrun", + f"--nproc_per_node={num_gpus}", + "distill.py", + "--use_mock_data", + "--validate_only", + ], + student_hf_path=teacher_hf_path, + teacher_hf_path=teacher_hf_path, + output_dir=output_dir, + tp_size=num_gpus, + pp_size=1, + seq_length=16, + mbs=1, + gbs=4, + train_iters=1, + eval_iters=1, + hf_validation_export_path=validation_exports, + hf_validation_export_interval=1, + student_hf_model=teacher_hf_path, + ) + run_example_command(cmd_parts, example_path="megatron_bridge") + + assert (validation_exports / "iter_0000000/config.json").exists() + assert not (output_dir / "checkpoints/iter_0000001").exists() def test_distill_puzzletron_anymodel(tmp_path: Path, num_gpus): diff --git a/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py b/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py index c0453484056..9cb6a82ada3 100644 --- a/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py +++ b/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py @@ -53,6 +53,61 @@ def test_megatron_preprocess_data_with_jsonl_path(tmp_path): assert Path(prefixes[0] + ".idx").stat().st_size > 0, "Index file should not be empty" +def test_megatron_preprocess_data_jsonl_stops_at_max_tokens(tmp_path): + input_path = tmp_path / "data.jsonl" + input_path.write_text( + "".join(json.dumps({"text": f"Document {index} " * 100}) + "\n" for index in range(10)), + encoding="utf-8", + ) + + common_args = { + "jsonl_paths": input_path, + "tokenizer_name_or_path": "gpt2", + "json_keys": "text", + "workers": 2, + } + limited_prefix = megatron_preprocess_data( + **common_args, + output_dir=tmp_path / "limited", + max_tokens=100, + )[0] + full_prefix = megatron_preprocess_data( + **common_args, + output_dir=tmp_path / "full", + )[0] + + # The .bin file stores token IDs, so its byte size reflects how many tokens were written. + limited_size = Path(limited_prefix + ".bin").stat().st_size + full_size = Path(full_prefix + ".bin").stat().st_size + assert limited_size < full_size + + +def test_megatron_preprocess_data_hf_split_stops_at_max_tokens(tmp_path): + common_args = { + "hf_dataset": "nanotron/minipile_100_samples", + "hf_split": "train", + "hf_max_samples_per_split": 100, + "hf_streaming": True, + "tokenizer_name_or_path": "gpt2", + "json_keys": "text", + "workers": 2, + } + limited_prefix = megatron_preprocess_data( + **common_args, + output_dir=tmp_path / "limited", + max_tokens=100, + )[0] + full_prefix = megatron_preprocess_data( + **common_args, + output_dir=tmp_path / "full", + )[0] + + # The .bin file stores token IDs, so its byte size reflects how many tokens were written. + limited_size = Path(limited_prefix + ".bin").stat().st_size + full_size = Path(full_prefix + ".bin").stat().st_size + assert limited_size < full_size + + @pytest.mark.parametrize( ("hf_dataset", "hf_split", "json_keys"), [