diff --git a/.github/actions/run-example/action.yml b/.github/actions/run-example/action.yml index 09a07297a8f..7138fb057dd 100644 --- a/.github/actions/run-example/action.yml +++ b/.github/actions/run-example/action.yml @@ -81,7 +81,17 @@ runs: # The scheduled and on-demand runs also run the tests marked nightly (tests/conftest.py). pytest_args+=(--run-nightly) fi - python3 -m pytest tests/examples/${{ inputs.example }} --cov "${pytest_args[@]}" || status=$? + readme_cli_tests="tests/examples/${{ inputs.example }}/test_${{ inputs.example }}_readme_cli.py" + if [[ -f "$readme_cli_tests" ]]; then + echo "Checking README CLI options before example execution" + python3 -m pytest "$readme_cli_tests" --no-cov -x "${pytest_args[@]}" || status=$? + pytest_args+=(--ignore="$readme_cli_tests") + fi + if (( status == 0 )); then + python3 -m pytest tests/examples/${{ inputs.example }} --cov "${pytest_args[@]}" || status=$? + else + echo "Skipping example execution because README CLI checks failed" + fi fi echo "status=$status" >> "$GITHUB_OUTPUT" - name: Verdict diff --git a/.github/workflows/_example_tests_runner.yml b/.github/workflows/_example_tests_runner.yml index e1463965a08..0d2a7ef5129 100644 --- a/.github/workflows/_example_tests_runner.yml +++ b/.github/workflows/_example_tests_runner.yml @@ -94,6 +94,18 @@ jobs: echo "PIP_CONSTRAINT=/tmp/pip-constraints.txt" >> "$GITHUB_ENV" fi + # Preserve the QAT container's matching Torch/FlashAttention binaries. + if [[ "${{ inputs.example }}" == "llm_qat" ]]; then + python3 -c 'import torch; import flash_attn_2_cuda' + + export PIP_CONSTRAINT="$RUNNER_TEMP/qat-constraints.txt" + python3 -c ' + from importlib.metadata import version + for name in ("torch", "flash-attn"): + print(f"{name}=={version(name)}") + ' > "$PIP_CONSTRAINT" + fi + # Use `python3 -m pip` instead of `pip` to avoid conflicts with system pip for nemo containers # Editable install so example scripts launched as subprocesses resolve modelopt to the same source path as the test process python3 -m pip install -e ".${{ inputs.pip_install_extras }}" diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml index 722819135c2..8f931792b51 100644 --- a/.github/workflows/example_tests.yml +++ b/.github/workflows/example_tests.yml @@ -76,6 +76,8 @@ jobs: - examples/dataset/** # data prep for llm_qat and speculative_decoding - examples/hf_ptq/** # the speculative_decoding test drives hf_ptq - "!**/*.{md,rst,png,ipynb}" + torch_docs: + - examples/llm_qat/README.md vllm: - "{examples,tests/examples}/vllm_serve/**" - examples/hf_ptq/convert_ggml_to_int8.py # tests/examples/vllm_serve serves its output @@ -87,6 +89,8 @@ jobs: megatron: - "{examples,tests/examples}/megatron_bridge/**" - "!**/*.{md,rst,png,ipynb}" + megatron_docs: + - examples/megatron_bridge/README.md diffusers: - "{examples,tests/examples}/diffusers/**" - "!**/*.{md,rst,png,ipynb}" @@ -101,10 +105,10 @@ jobs: # Stacked PRs (base is another PR's branch) skip GPU tests; only the bottom PR runs them. STACKED: ${{ steps.base.outputs.stacked }} COMMON: ${{ steps.changed.outputs.common_any_modified }} - TORCH: ${{ steps.changed.outputs.torch_any_modified }} + TORCH: ${{ steps.changed.outputs.torch_any_modified == 'true' || steps.changed.outputs.torch_docs_any_modified == 'true' }} VLLM: ${{ steps.changed.outputs.vllm_any_modified }} TRTLLM: ${{ steps.changed.outputs.trtllm_any_modified }} - MEGATRON: ${{ steps.changed.outputs.megatron_any_modified }} + MEGATRON: ${{ steps.changed.outputs.megatron_any_modified == 'true' || steps.changed.outputs.megatron_docs_any_modified == 'true' }} DIFFUSERS: ${{ steps.changed.outputs.diffusers_any_modified }} ONNX: ${{ steps.changed.outputs.onnx_any_modified }} run: | diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index 4f22e19a019..183fce7eb9f 100644 --- a/.github/workflows/unit_tests.yml +++ b/.github/workflows/unit_tests.yml @@ -11,6 +11,10 @@ on: - "noxfile.py" - "pyproject.toml" - "tests/unit/**" + - "tests/_test_utils/**" + - "tests/examples/**/test_*_readme.py" + - "examples/**/README*.md" + - "examples/**/readme*.md" - "tools/launcher/**" - "tools/mcp/**" - "tools/resource_monitor.py" @@ -52,6 +56,9 @@ jobs: pyproject.toml tests/unit/** tests/_test_utils/** + tests/examples/**/test_*_readme.py + examples/**/README*.md + examples/**/readme*.md tools/launcher/** tools/mcp/** tools/resource_monitor.py diff --git a/examples/llm_qat/README.md b/examples/llm_qat/README.md index ccdc2facd09..bf16abea65b 100644 --- a/examples/llm_qat/README.md +++ b/examples/llm_qat/README.md @@ -27,11 +27,15 @@ Please refer to [hf_ptq/README.md](../hf_ptq/README.md#pre-requisites) for conta recommendations and base ModelOpt installation guidance. For this QAT/QAD example, install the Hugging Face dependencies and the example-specific requirements: + ```bash pip install -U nvidia-modelopt[hf] -pip install -r examples/llm_qat/requirements.txt +pip install --no-build-isolation -r examples/llm_qat/requirements.txt ``` +`--no-build-isolation` lets FlashAttention build against the PyTorch installed by +the first command. + The Qwen3-8B example below requires a minimum of **2 x 80GB GPUs**. > ModelOpt provides accelerated quantization kernels using Triton for NVFP4 QAT. See the [installation guide](https://nvidia.github.io/Model-Optimizer/getting_started/_installation_for_Linux.html#accelerated-quantization-with-triton-kernels). @@ -45,6 +49,8 @@ All arguments can be set via YAML, CLI, or both (CLI overrides YAML). See Recipes are declarative YAML files that specify the quantization configuration. Built-in recipes are available in [`modelopt_recipes/`](../../modelopt_recipes/): + + ```sh # From the Model-Optimizer repository root, list available built-in recipes ls modelopt_recipes/general/ptq/ @@ -54,9 +60,13 @@ See [custom calibration](https://nvidia.github.io/Model-Optimizer/guides/_pytorc ### QAT -Quantize, fine-tune on labeled data, and export: +Quantize, fine-tune on labeled data, and export. Start from the repository root: + + ```sh +cd examples/llm_qat + # 1. Quantize python quantize.py \ --model_name_or_path Qwen/Qwen3-8B \ @@ -78,6 +88,7 @@ python export.py --pyt_ckpt_path qwen3-8b-qat-nvfp4 --export_path qwen3-8b-qat-d Quantize, recover accuracy using the original model as teacher, and export: + ```sh # 1. Quantize python quantize.py \ @@ -102,10 +113,12 @@ Exported checkpoints can be deployed on [TensorRT-LLM](https://github.com/NVIDIA > [!NOTE] > For a minimal end-to-end demo (quantize + train + save in one script), see [simple_qat_train.py](simple_qat_train.py). It runs on a **single GPU** only and is intended as a quick introduction to the QAT flow (without transformer trainer)—not for distributed training. > +> > ```sh > python simple_qat_train.py --model-path meta-llama/Llama-3.2-3B --recipe general/ptq/nvfp4_default-kv_fp8 > ``` > +> > For multi-GPU training (FSDP2, DDP, DeepSpeed), use [train.py](train.py) with `accelerate launch` as shown in the [commands](#qat) above. > [!TIP] @@ -154,6 +167,7 @@ Replace `--config-file configs/accelerate/fsdp2.yaml` with the desired backend c [QLoRA](https://arxiv.org/pdf/2305.14314) reduces training memory by quantizing LoRA backbone weights with real quantization via `mtq.compress()`. + ```sh # 1. Quantize with compression python quantize.py \ @@ -174,9 +188,15 @@ python export.py \ --pyt_ckpt_path qwen3-8b-fp4-qlora \ --export_path qwen3-8b-fp4-qlora-hf +``` + +Serve the exported adapter with vLLM: + + +```sh # 4. Serve with vLLM vllm serve qwen3-8b-fp4-qlora-hf/base_model --enable-lora \ - --lora-modules adapter=qwen3-8b-fp4-qlora-hf --port 8000 \ + --lora-modules adapter=qwen3-8b-fp4-qlora-hf --port "${PORT:-8000}" \ --tokenizer qwen3-8b-fp4-qlora-hf ``` @@ -186,6 +206,7 @@ vllm serve qwen3-8b-fp4-qlora-hf/base_model --enable-lora \ ### Quantize and Fine-Tune with Python + ```python import modelopt.torch.quantization as mtq from modelopt.recipe import load_recipe @@ -205,6 +226,7 @@ trainer.save_model() `QATTrainer` is a drop-in replacement for HuggingFace's `Trainer` that handles quantization-aware training seamlessly with various distributed backends (FSDP2, DeepSpeed, DDP): + ```python from modelopt.torch.quantization.plugins.transformers_trainer import QATTrainer @@ -218,15 +240,19 @@ trainer.train() trainer.save_model() ``` -`QADTrainer` extends `QATTrainer` with distillation. Pass the teacher model and a `DistillArguments` instance: +`QADTrainer` extends `QATTrainer` with distillation. Load the teacher first and pass a `DistillArgsWithTeacherModel` instance: + ```python -from modelopt.torch.distill.plugins.huggingface import DistillArguments +from transformers import AutoModelForCausalLM + +from modelopt.torch.distill.plugins.huggingface import DistillArgsWithTeacherModel from modelopt.torch.quantization.plugins.transformers_trainer import QADTrainer -distill_args = DistillArguments( +teacher = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-8B", dtype=model.dtype) +distill_args = DistillArgsWithTeacherModel( distill=True, - teacher_model="Qwen/Qwen3-8B", + teacher_model=teacher, criterion="logits_loss", ) @@ -250,14 +276,17 @@ You can either: 1. **Override via CLI** (recommended for one-off runs): + ```sh accelerate launch --config-file configs/accelerate/fsdp2.yaml \ --fsdp_transformer_layer_cls_to_wrap LlamaDecoderLayer \ - train.py --config configs/train/qat_nvfp4.yaml ... + train.py --config configs/train/qat_nvfp4.yaml \ + --model_name_or_path llama-quantized --output_dir llama-qat ``` 2. **Create a custom config** (recommended for repeated use): + ```sh cp configs/accelerate/fsdp2.yaml configs/accelerate/fsdp2_llama.yaml # Edit fsdp2_llama.yaml: change Qwen3DecoderLayer -> LlamaDecoderLayer @@ -282,10 +311,12 @@ There are two types of configs: `quantize.py` only needs `--dataset_config` and `--recipe`. `train.py` uses a full training config via `--config`. All arguments can be specified via YAML, CLI flags, or both (CLI overrides YAML). See [ARGUMENTS.md](ARGUMENTS.md) for the full reference, regenerated with `python_pwd examples/llm_qat/arguments.py --generate_docs examples/llm_qat/ARGUMENTS.md`. + ```sh # YAML + CLI override accelerate launch --config-file configs/accelerate/fsdp2.yaml train.py \ - --config configs/train/qat_nvfp4.yaml --learning_rate 5e-5 + --config configs/train/qat_nvfp4.yaml --learning_rate 5e-5 \ + --model_name_or_path qwen3-8b-quantized --output_dir qwen3-8b-qat-override ``` See [Dataset Configuration](configs/dataset/README.md) for custom dataset blends and adding new datasets. @@ -297,6 +328,7 @@ See [Dataset Configuration](configs/dataset/README.md) for custom dataset blends You can pre-tokenize and cache the dataset before training using `dataset_utils.py`. This is useful for large blends or multi-node setups where you want to build the cache once and reuse it across experiments. + ```sh python dataset_utils.py \ --dataset_config configs/dataset/blend.yaml \ @@ -315,13 +347,14 @@ See [Qwen3-8B evaluation results](evaluation/README.md) for the teacher, PTQ, an ModelOpt quantized models can be saved and restored without exporting to a deployment platform. This is useful for fast evaluation with fake quantization using standard LLM benchmarks (MMLU, WikiText, etc.). See [HuggingFace checkpointing](https://nvidia.github.io/Model-Optimizer/guides/2_save_load.html#modelopt-save-restore-using-huggingface-checkpointing-apis) for details. + ```sh cd ../llm_eval python lm_eval_hf.py --model hf \ --tasks mmlu,wikitext \ --model_args pretrained=../llm_qat/qwen3-8b-qat-nvfp4 \ - --batch_size 4 + --batch_size 4 --output_path "${EVAL_OUTPUT:-eval-results}" ``` See [llm_eval/README.md](../llm_eval/README.md) for supported tasks. diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md index c16a5c6e147..39b334ede45 100644 --- a/examples/megatron_bridge/README.md +++ b/examples/megatron_bridge/README.md @@ -26,6 +26,7 @@ Running these examples requires many additional dependencies to be installed (e. To get the ModelOpt examples scripts, mount your Model-Optimizer repo to the container as follows: + ```bash export MODELOPT_DIR=${PWD}/Model-Optimizer # or set to your local Model-Optimizer repository path if you have cloned it if [ ! -d "${MODELOPT_DIR}" ]; then @@ -68,8 +69,9 @@ and add these mounts to the `docker run` command above: You also need to login with your HuggingFace token to download gated datasets / models. Note that the default dataset for pruning and quantization is [`nemotron-post-training-dataset-v2`](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2), which is gated. + ```bash -hf auth login --token +hf auth login --token "$HF_TOKEN" ``` ### Importing a HuggingFace Checkpoint (optional) @@ -79,15 +81,22 @@ checkpoint (e.g. to reuse across runs), convert it with Megatron-Bridge's conver `--tp` / `--pp` / `--ep` to shard a model that does not fit on one GPU (`--ep` for MoE models), or `--device cpu` to convert in a single process without GPUs: + ```bash bash /opt/Megatron-Bridge/scripts/conversion/convert.sh import \ --executor local \ --device gpu \ --gpus-per-node 8 \ --hf-model Qwen/Qwen3-8B \ - --megatron-path /tmp/Qwen3-8B-megatron + --megatron-path "${OUTPUT_ROOT:-/tmp}/Qwen3-8B-megatron" ``` +The commands accept optional environment overrides for `MODEL`, `OUTPUT_ROOT`, +`CALIB_DATA`, `CALIB_SAMPLES`, `SEQ_LENGTH`, `GPUS`, `GLOBAL_BATCH`, `TRAIN_ITERS`, +`EVAL_INTERVAL`, `EVAL_ITERS`, `WARMUP_ITERS`, and `PRUNE_CONFIG` where shown. +Unset variables retain the full-scale defaults. README tests use local tiny models +and data through these same variables; see the [test contract](../../tests/_test_utils/doc_tests/README.md). + ## Post-Training Quantization This section shows how to quantize a HuggingFace model using ModelOpt in the Megatron-Bridge framework. Quantization is a two-step flow: @@ -99,14 +108,17 @@ This section shows how to quantize a HuggingFace model using ModelOpt in the Meg **Step 1 — quantize** Qwen3-8B to NVFP4 on 2 GPUs (Tensor Parallelism = 2) using 1024 samples from default dataset (Mix of [`cnn_dailymail`](https://huggingface.co/datasets/abisee/cnn_dailymail) and [`nemotron-post-training-dataset-v2`](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2)) for calibration (sequence length = 4096): + ```bash torchrun --nproc_per_node 2 quantize.py \ - --hf_model_name_or_path Qwen/Qwen3-8B \ + --hf_model_name_or_path "${MODEL:-Qwen/Qwen3-8B}" \ --quant_cfg nvfp4 \ --tp_size 2 \ --calib_batch_size 1 \ - --seq_length 4096 \ - --export_megatron_path /tmp/Qwen3-8B-NVFP4-megatron + --seq_length "${SEQ_LENGTH:-4096}" \ + --calib_dataset_name "${CALIB_DATA:-cnn_nemotron_v2_mix}" \ + --calib_num_samples "${CALIB_SAMPLES:-1024}" \ + --export_megatron_path "${OUTPUT_ROOT:-/tmp}/Qwen3-8B-NVFP4-megatron" ``` > [!NOTE] @@ -114,12 +126,13 @@ torchrun --nproc_per_node 2 quantize.py \ **Step 2 — export** the Megatron checkpoint to a deployable HuggingFace checkpoint: + ```bash torchrun --nproc_per_node 2 export_quantized_megatron_to_hf.py \ - --hf_model_name_or_path Qwen/Qwen3-8B \ - --megatron_path /tmp/Qwen3-8B-NVFP4-megatron \ + --hf_model_name_or_path "${MODEL:-Qwen/Qwen3-8B}" \ + --megatron_path "${OUTPUT_ROOT:-/tmp}/Qwen3-8B-NVFP4-megatron" \ --pp_size 2 \ - --export_unified_hf_path /tmp/Qwen3-8B-NVFP4-hf + --export_unified_hf_path "${OUTPUT_ROOT:-/tmp}/Qwen3-8B-NVFP4-hf" ``` > [!NOTE] @@ -144,13 +157,14 @@ For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `quantize.py` automati Set MLflow's own `MLFLOW_TRACKING_URI`, or pass `--mlflow `, to record a run on an MLflow server. Every script here that writes a checkpoint takes the flag — [`prune_minitron.py`](#pruning), `quantize.py`, [`distill.py`](#distillation), `export_quantized_megatron_to_hf.py` and `export_distilled_megatron_to_hf.py` — and they share one experiment-name convention, so a pruning, a quantization, the distillation that refines its checkpoint and the export that deploys it can be found together. (One gap remains: `distill.py --hf_export_path` writes its HuggingFace checkpoint from rank 0, which is not the rank that owns that run, so it carries no pointer.) + ```bash torchrun --nproc_per_node 2 quantize.py \ - --hf_model_name_or_path Qwen/Qwen3-8B \ + --hf_model_name_or_path "${MODEL:-Qwen/Qwen3-8B}" \ --recipe general/ptq/nvfp4_default-kv_fp8 \ --tp_size 2 \ - --export_megatron_path /tmp/Qwen3-8B-NVFP4-megatron \ - --mlflow https:/// + --export_megatron_path "${OUTPUT_ROOT:-/tmp}/Qwen3-8B-NVFP4-megatron" \ + --mlflow "${MLFLOW_TRACKING_URI:?set the tracking URI}" ``` The run opens *before* the model loads, so a bad URI fails in seconds rather than after a full calibration. Only the master rank uploads: the invocation, every argument as a searchable param, the resolved recipe, that rank's log and the quantizer summary — plus `.experiment.json` written into `--export_megatron_path` once the checkpoint is saved, so a checkpoint on disk names the run that produced it. A failed run is still recorded, with its traceback. @@ -226,13 +240,14 @@ token ids, and the KD losses compare the two models' logits elementwise over the Example usage to distill a 4B student (HF) from an 8B teacher (HF) on 8 GPUs (TP=8, PP=1): + ```bash torchrun --nnodes 1 --nproc_per_node 8 distill.py \ --tp_size 8 \ --teacher_hf_path Qwen/Qwen3-8B \ --student_hf_path Qwen/Qwen3-4B \ - --data_paths 1.0 tokenized_qwen3/data1_text_document 1.0 tokenized_qwen3/data2_text_document \ - --data_path_to_cache /path/to/cache/dataset_indices_qwen3 \ + --data_paths 1.0 "${DATA_PREFIX_1:?set a prepared Megatron data prefix}" 1.0 "${DATA_PREFIX_2:?set a prepared Megatron data prefix}" \ + --data_path_to_cache "${OUTPUT_ROOT:-/tmp}/dataset_indices_qwen3" \ --seq_length 8192 \ --mbs 1 \ --gbs 768 \ @@ -243,7 +258,7 @@ torchrun --nnodes 1 --nproc_per_node 8 distill.py \ --eval_interval 100 \ --eval_iters 32 \ --log_interval 10 \ - --output_dir /output/qwen3_8b_to_4b_distill + --output_dir "${OUTPUT_ROOT:-/output}/qwen3_8b_to_4b_distill" ``` Tensorboard logging is enabled by default and logs are saved to `/tensorboard` directory. @@ -255,6 +270,7 @@ This skips training and evaluates the student at iteration 0. To see all available arguments: + ```bash torchrun --nproc_per_node 1 distill.py --help ``` @@ -263,25 +279,28 @@ torchrun --nproc_per_node 1 distill.py --help Example usage with mock data for quick testing (no pre-tokenized data needed): + ```bash -torchrun --nproc_per_node 8 distill.py \ - --tp_size 8 \ - --teacher_hf_path Qwen/Qwen3-0.6B \ - --student_hf_path Qwen/Qwen3-0.6B \ +torchrun --nproc_per_node "${GPUS:-8}" distill.py \ + --tp_size "${GPUS:-8}" \ + --teacher_hf_path "${MODEL:-Qwen/Qwen3-0.6B}" \ + --student_hf_path "${MODEL:-Qwen/Qwen3-0.6B}" \ --use_mock_data \ - --seq_length 512 \ + --seq_length "${SEQ_LENGTH:-512}" \ --mbs 1 \ - --gbs 8 \ - --train_iters 100 \ - --eval_interval 10 \ - --eval_iters 4 \ - --output_dir /tmp/test_distill + --gbs "${GLOBAL_BATCH:-8}" \ + --train_iters "${TRAIN_ITERS:-100}" \ + --eval_interval "${EVAL_INTERVAL:-10}" \ + --eval_iters "${EVAL_ITERS:-4}" \ + --lr_warmup_iters "${WARMUP_ITERS:-50}" \ + --output_dir "${OUTPUT_ROOT:-/tmp}/test_distill" ``` ### Vision-Language Models (VLMs) For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `distill.py` distills only the **language model** (on text data) and leaves the vision tower and projector untouched — matching the pruning and quantization behavior. It composes with pruning and QAD (`--student_megatron_path`) exactly as for LLMs, and the HF export reuses `--student_hf_path` (no `--student_hf_model` needed). + ```bash torchrun --nproc_per_node 8 distill.py \ --tp_size 8 \ @@ -296,6 +315,7 @@ A **non-quantized** distilled checkpoint (LLM or VLM) is saved in Megatron distr **Inline** -- add `--hf_export_path` to the `distill.py` command to automatically convert the **final** checkpoint after distillation: + ```bash torchrun --nnodes 1 --nproc_per_node 8 distill.py \ ... \ @@ -306,6 +326,7 @@ torchrun --nnodes 1 --nproc_per_node 8 distill.py \ **Separate conversion** -- convert **any** saved iteration (intermediate or final) with [export_distilled_megatron_to_hf.py](export_distilled_megatron_to_hf.py): + ```bash torchrun --nproc_per_node 1 export_distilled_megatron_to_hf.py \ --student_hf_path \ @@ -321,16 +342,18 @@ to keep all saved checkpoints. Then export all retained checkpoints, with one Hugging Face checkpoint written per `iter_` subdirectory: + ```bash torchrun --nproc_per_node 1 export_distilled_megatron_to_hf.py \ - --student_hf_path \ - --megatron_path /checkpoints \ - --hf_export_path /path/to/save/hf_validation_checkpoints \ + --student_hf_path "${MODEL:-Qwen/Qwen3-0.6B}" \ + --megatron_path "${OUTPUT_ROOT:-/tmp}/test_distill/checkpoints" \ + --hf_export_path "${OUTPUT_ROOT:-/tmp}/hf_validation" \ --export_iterations all ``` The export path contains one loadable Hugging Face checkpoint per exported iteration: + ```text hf_validation/ ├── iter_0000100/ @@ -353,20 +376,21 @@ Recalculate the budget when changing those settings, and keep the prepared data We also use a smaller learning rate for QAD: + ```bash torchrun --nproc_per_node 8 distill.py \ --tp_size 8 \ --teacher_hf_path Qwen/Qwen3-8B \ --student_hf_path Qwen/Qwen3-8B \ - --student_megatron_path /tmp/Qwen3-8B-NVFP4-megatron \ - --data_paths 1.0 tokenized_qwen3/data1_text_document 1.0 tokenized_qwen3/data2_text_document \ - --data_path_to_cache /path/to/cache/dataset_indices_qwen3 \ + --student_megatron_path "${QUANTIZED_CHECKPOINT:?set the PTQ Megatron checkpoint path}" \ + --data_paths 1.0 "${DATA_PREFIX_1:?set a prepared Megatron data prefix}" 1.0 "${DATA_PREFIX_2:?set a prepared Megatron data prefix}" \ + --data_path_to_cache "${OUTPUT_ROOT:-/tmp}/dataset_indices_qwen3" \ --seq_length 8192 \ --gbs 768 \ --train_iters 1000 \ --lr 1e-5 \ --min_lr 5e-6 \ - --output_dir /output/qwen3_8b_nvfp4_qad + --output_dir "${OUTPUT_ROOT:-/output}/qwen3_8b_nvfp4_qad" ``` With pipeline parallelism, the last stage also runs the teacher and student output layers and the KD loss, and @@ -391,7 +415,7 @@ This keeps an extra BF16 copy of the student's quantized weights on every GPU. T uncached one. It needs a straight-through weight gradient (`pass_through_bwd`, the default). The more micro-batches per step, and the more expensive the weight format, the more it saves. It needs a [newer Megatron-Bridge](#pre-requisites) than `nemo:26.08`. -The distilled checkpoint retains the ModelOpt quantization state, so it can be converted to a deployable HuggingFace checkpoint with [export_quantized_megatron_to_hf.py](export_quantized_megatron_to_hf.py) (point `--megatron_path` at `/output/qwen3_8b_nvfp4_qad/checkpoints`), exactly like the PTQ checkpoint in [step 2 above](#post-training-quantization). +The distilled checkpoint retains the ModelOpt quantization state, so it can be converted to a deployable HuggingFace checkpoint with [export_quantized_megatron_to_hf.py](export_quantized_megatron_to_hf.py) (point `--megatron_path` at `${OUTPUT_ROOT:-/output}/qwen3_8b_nvfp4_qad/checkpoints`), exactly like the PTQ checkpoint in [step 2 above](#post-training-quantization). #### Evaluating a QAD checkpoint in vLLM with fake quantization @@ -418,6 +442,7 @@ Weight-only formats leave no quantizer enabled after folding, so no sidecars are Add `--lora_rank` to PTQ to train low-rank weight updates with the usual QAD flow: + ```bash torchrun --nproc_per_node 2 quantize.py \ --hf_model_name_or_path Qwen/Qwen3-0.6B \ @@ -475,49 +500,58 @@ Multiple NAS targets can be combined — e.g. `--prune_target_params 6e9 --prune at-most 20% depth (`num_layers`) and 40% width is pruned per prunable hparam (`hidden_size`, `ffn_hidden_size`, ...), top-10 candidates are evaluated for MMLU score (5% sampled data) to select the best model. + ```bash torchrun --nproc_per_node 2 prune_minitron.py \ --pp_size 2 \ - --hf_model_name_or_path Qwen/Qwen3-8B \ + --hf_model_name_or_path "${MODEL:-Qwen/Qwen3-8B}" \ --prune_target_params 6e9 \ --hparams_to_skip num_attention_heads \ - --output_hf_path /tmp/Qwen3-8B-Pruned-6B + --output_hf_path "${OUTPUT_ROOT:-/tmp}/Qwen3-8B-Pruned-6B" ``` **Prune by active parameter count** — useful for MoE models where most experts are inactive per token (e.g. prune Nemotron-3-Nano-30B-A3B-BF16 (3.6B active params) to 3B active params): + ```bash torchrun --nproc_per_node 2 prune_minitron.py \ --pp_size 2 \ --hf_model_name_or_path nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 \ --prune_target_active_params 3e9 \ - --output_hf_path /tmp/Nemotron-3-Nano-30B-A3B-BF16-Pruned-3B-Active + --output_hf_path "${OUTPUT_ROOT:-/tmp}/Nemotron-3-Nano-30B-A3B-BF16-Pruned-3B-Active" ``` **Prune by memory footprint** — prune to fit a target GPU memory budget (weights + KV-cache at the given sequence length and batch size, assuming BF16): + ```bash torchrun --nproc_per_node 2 prune_minitron.py \ --pp_size 2 \ - --hf_model_name_or_path Qwen/Qwen3-8B \ + --hf_model_name_or_path "${MODEL:-Qwen/Qwen3-8B}" \ --prune_target_memory_mb 12288 \ --seq_length 4096 \ --calib_batch_size 1 \ - --output_hf_path /tmp/Qwen3-8B-Pruned-12GB + --output_hf_path "${OUTPUT_ROOT:-/tmp}/Qwen3-8B-Pruned-12GB" ``` **Manual pruning** — prune directly to a specified architecture (no NAS, no score evaluation): + ```bash +PRUNE_CONFIG=${PRUNE_CONFIG:-'{"hidden_size":3584,"ffn_hidden_size":9216}'} torchrun --nproc_per_node 2 prune_minitron.py \ --pp_size 2 \ - --hf_model_name_or_path Qwen/Qwen3-8B \ - --prune_export_config '{"hidden_size": 3584, "ffn_hidden_size": 9216}' \ - --output_hf_path /tmp/Qwen3-8B-Pruned-6B-manual + --hf_model_name_or_path "${MODEL:-Qwen/Qwen3-8B}" \ + --prune_export_config "$PRUNE_CONFIG" \ + --calib_dataset_name "${CALIB_DATA:-nemotron-post-training-dataset-v2}" \ + --calib_num_samples "${CALIB_SAMPLES:-1024}" \ + --seq_length "${SEQ_LENGTH:-4096}" \ + --output_hf_path "${OUTPUT_ROOT:-/tmp}/Qwen3-8B-Pruned-6B-manual" ``` To see the full usage for advanced configurations, run: + ```bash torchrun --nproc_per_node 1 prune_minitron.py --help ``` @@ -540,24 +574,26 @@ For a vision-language model (e.g. Qwen3.5-VL, Gemma3-VL), `prune_minitron.py` au - The `--prune_target_params` / `--prune_target_active_params` / `--prune_target_memory_mb` targets (and `export_config` dimensions) apply to the **language model only** — the (unpruned) vision tower's parameters are *not* counted, so the full saved VLM will be larger than the target. - `hidden_size` is never pruned for VLMs (it is shared with the vision projector). + ```bash torchrun --nproc_per_node 2 prune_minitron.py \ --pp_size 2 \ --hf_model_name_or_path Qwen/Qwen3.5-4B \ --prune_target_params 3e9 \ - --output_hf_path /tmp/Qwen3.5-4B-Pruned-3B + --output_hf_path "${OUTPUT_ROOT:-/tmp}/Qwen3.5-4B-Pruned-3B" ``` ## Sanity-Check Generation [generate_vllm.py](generate_vllm.py) runs a quick generation check on an exported HuggingFace checkpoint using vLLM — a useful smoke test for a **quantized**, **pruned**, or **distilled** model to confirm it still produces coherent text. For quantized checkpoints, vLLM auto-detects the ModelOpt quantization from the exported `hf_quant_config.json`, so no extra flags are needed: + ```bash # Quantized model python generate_vllm.py --model nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4 --trust_remote_code # Pruned model -python generate_vllm.py --model /tmp/Qwen3-8B-Pruned-6B +python generate_vllm.py --model "${PRUNED_CHECKPOINT:?set the pruned Hugging Face checkpoint path}" ``` > [!NOTE] diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 8b233e05916..4430746e34c 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -998,8 +998,8 @@ def _per_stage_recompute_hook(model_chunks): if __name__ == "__main__": - dist.setup() args = get_args() + dist.setup() try: # Entered inside the try: opening the run is fatal by design, and the peers of a rank # that exits without dist.abort() stay blocked on the first collective. diff --git a/examples/megatron_bridge/export_distilled_megatron_to_hf.py b/examples/megatron_bridge/export_distilled_megatron_to_hf.py index 26c481c6de5..7cd7c8f50a6 100644 --- a/examples/megatron_bridge/export_distilled_megatron_to_hf.py +++ b/examples/megatron_bridge/export_distilled_megatron_to_hf.py @@ -349,8 +349,8 @@ def main(args: argparse.Namespace): if __name__ == "__main__": - dist.setup() args = get_args() + dist.setup() try: # Entered inside the try: opening the run is fatal by design, and the peers of a rank # that exits without dist.abort() stay blocked on the first collective. diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py index 3b7871e5ebb..f69af6c2dbc 100644 --- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py +++ b/examples/megatron_bridge/export_quantized_megatron_to_hf.py @@ -297,8 +297,8 @@ def main(args: argparse.Namespace): if __name__ == "__main__": - dist.setup() args = get_args() + dist.setup() try: # Entered inside the try: opening the run is fatal by design, and the peers of a rank # that exits without dist.abort() stay blocked on the first collective. diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index 28e86c8baa1..e8f73d56704 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -794,8 +794,8 @@ def score_func(m): if __name__ == "__main__": - dist.setup() args = get_args() + dist.setup() try: # Entered inside the try: opening the run is fatal by design, and the peers of a rank # that exits without dist.abort() stay blocked on the first collective. diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py index efdab627b36..effa57567e8 100644 --- a/examples/megatron_bridge/quantize.py +++ b/examples/megatron_bridge/quantize.py @@ -538,8 +538,8 @@ def forward_loop(_model=None): if __name__ == "__main__": - dist.setup() args = get_args() + dist.setup() try: # Entered inside the try: opening the run is fatal by design, and the peers of a rank # that exits without dist.abort() stay blocked on the first collective. diff --git a/modelopt/torch/distill/plugins/huggingface.py b/modelopt/torch/distill/plugins/huggingface.py index 80675a3f4aa..afdc193c342 100644 --- a/modelopt/torch/distill/plugins/huggingface.py +++ b/modelopt/torch/distill/plugins/huggingface.py @@ -91,7 +91,7 @@ class DistillArguments(ModelOptHFArguments): "The CE (cross-entropy) loss weight is (1 - kd_loss_alpha). " "Set to 1.0 (default) for pure KD loss. " "Set to a value in (0, 1) to blend KD and CE losses, " - "e.g. 0.7 for 70% KD loss + 30% CE loss." + "e.g. 0.7 for 70%% KD loss + 30%% CE loss." ) }, ) diff --git a/modelopt/torch/opt/plugins/transformers.py b/modelopt/torch/opt/plugins/transformers.py index 83155be5344..cfefc51bf01 100644 --- a/modelopt/torch/opt/plugins/transformers.py +++ b/modelopt/torch/opt/plugins/transformers.py @@ -15,6 +15,7 @@ """ModelOpt plugin for enabling automatic save/restore of ModelOpt state for HuggingFace models.""" +import argparse import dataclasses import fnmatch import gc @@ -288,6 +289,26 @@ def __init__(self, *args, docs_header_extra: str | None = None, **kwargs): super().__init__(*args, **kwargs) self._docs_header_extra = docs_header_extra + def format_help(self): + """Include options consumed before HuggingFace's dataclass parsing.""" + options = argparse.ArgumentParser(add_help=False) + options.add_argument( + "--config", + metavar="PATH", + help="Load YAML defaults; command-line arguments override them.", + ) + options.add_argument( + "--generate_docs", + nargs="?", + metavar="PATH", + help="Generate argument documentation and exit (default: ARGUMENTS.md).", + ) + formatter = self._get_formatter() + formatter.start_section("ModelOpt options") + formatter.add_arguments(options._actions) + formatter.end_section() + return super().format_help() + "\n" + formatter.format_help() + def parse_args_into_dataclasses(self, args=None, **kwargs): """Parse args with optional YAML config defaults and doc generation.""" if args is None: diff --git a/tests/_test_utils/doc_tests/README.md b/tests/_test_utils/doc_tests/README.md new file mode 100644 index 00000000000..e3397516b94 --- /dev/null +++ b/tests/_test_utils/doc_tests/README.md @@ -0,0 +1,280 @@ +# Executable Markdown tests + +The QAT and Megatron Bridge READMEs use the same parser, subprocess runner and +pytest adapter. Example-specific code only prepares inputs and verifies results. +Standalone snippets execute as written; existing end-to-end tests substitute explicit +CI inputs in the extracted command arguments. Neither path mocks model operations. +These are trusted repository tests with the same privileges as pytest. + +## Why a snippet runner + +Pytest provides collection, fixtures, assertions and manual/resource markers, while +pytest-timeout provides an outer test deadline. The snippet layer adds Markdown +extraction, complete ID accounting, source locations and a worker deadline covering +setup, execution and verification, with process-group cleanup. + +`tests/_test_utils/examples/run_command.py` runs individual example commands and +supports streaming output and Hub retries. It does not extract README fences or +preserve one shell/Python namespace across selected snippets, and it imports +Torch-specific distributed utilities. Keeping this small adapter separate lets +CPU-only documentation checks run without importing the ML stack. The migrated CLI workflows use that existing runner with arguments extracted +from README fences. + +Stdlib doctest checks Python prompt/output examples. These READMEs also contain +shell pipelines, distributed training, background services and artifact assertions, +which need the scenario/handler contract described below. + +## Fixture scope + +The parser, runner and pytest adapter are generic. Helpers under `fixtures/` +implement inputs and output checks for the remaining standalone/manual scenarios: use `llm_qat` or `bridge` +in helper names, and prefix standalone QAT fixture modules with `llm_qat_`. +Installation copies QAT requirements and checks its dependencies; serving expects +the QAT README's QLoRA checkpoint layout and adapter name. These are not general +installation or serving APIs. Add a shared abstraction only when another example +needs the same behavior with a clear parameter contract. + +## Snippet IDs and Python scenarios + +A README contains only one ID marker immediately before each fence (blank lines +are allowed). IDs are unique within the document and remain stable when prose or +snippets move. Code is extracted directly, without substitutions: + +````markdown + +```bash +printf '%s' "$MESSAGE" > result.txt +``` +```` + +The pytest module owns scenario definitions, execution order, resource metadata, +and exclusions. One snippet can be reused in multiple scenarios: + +```python +SCENARIOS = [ + Scenario( + "example-smoke", + ["example-command"], + timeout_seconds=30, + handler="_test_utils.doc_tests.fixtures.example_readme:smoke", + ), +] +EXCLUDED = {"interactive-container": "Provisioned outside pytest by the CI job."} + + +@pytest.mark.parametrize( + "scenario", scenario_parameters(README, SCENARIOS, excluded=EXCLUDED) +) +def test_readme(scenario, tmp_path, monkeypatch): + execute_scenario(scenario, REPO, tmp_path, monkeypatch) +``` + +An optional importable handler performs setup, calls `ctx.run()` exactly once, +then verifies results. Put it in an example-specific module under `fixtures/`: + +```python +def smoke(ctx): + ctx.cwd = ctx.tmp + ctx.env["MESSAGE"] = "hello" + ctx.run() + assert (ctx.cwd / "result.txt").read_text() == "hello" +``` + +- `parse_markdown(path)` extracts named fences and source locations without + executing them. `resolve_scenarios` requires an ID on **every** fence and + validates that each ID is referenced or explicitly excluded with a reason. + Missing/duplicate IDs, unknown references, stale exclusions, and tested/excluded + overlaps fail collection. Accounting does **not** mean manual/excluded code passed. + Scenario metadata is trusted Python test configuration; the resolver does not + duplicate type or format checks for resource fields and handler names. +- `Scenario` requires a unique scenario ID and a nonempty ordered list of snippet + IDs. Defaults are `profile="cpu"` and `timeout_seconds=60`. Optional metadata: + `profile="gpu"`, `min_gpus=2`, `requires=["megatron.bridge"]`, `manual=True`, + `background=True`, and `handler="module:function"`. +- Handlers are imported only by the worker. Collection does not import their ML + dependencies or discover GPUs. Keep optional ML imports in the worker path so + CPU-only scenarios can still run independently. Without a handler, the worker + simply runs the selected snippets. +- `ctx.repo` and `ctx.tmp` are absolute paths. `ctx.cwd` starts at the repository; + `ctx.env` is mutable. Select a temporary workspace when isolation is needed. +- Shell snippets run in test-defined order in one Bash process with + `set -Eeuo pipefail`. Variables, directory changes, and exported environment + persist; the handler sees final `ctx.cwd` and `ctx.env` after `ctx.run()`. +- Python snippets share the dictionary supplied to `ctx.run(namespace)`, which is + also returned for assertions. The worker supplies `ctx`, `__name__="__main__"`, + and `__file__` pointing to the README. Before execution it applies `ctx.cwd` and + `ctx.env`; afterward it records their final values. Keep Python and shell + snippets in separate scenarios. +- `background=True` starts the **last shell snippet** asynchronously. A handler + must verify readiness and behavior after `ctx.run()`. `ctx.background_pid` + identifies its shell process. Background changes to directory/environment do + not propagate back; configure server settings before running the snippets. +- The parent passes a JSON execution plan to the worker; the worker does not + reparse the README or import pytest modules. The deadline covers handler import, + setup, snippet execution, and verification. The process group is terminated on + success, failure, or timeout. A zero exit status cannot bypass completion checks. +- Failures include the scenario ID and Markdown source location; execution logs + identify each snippet ID. Python snippets preserve README traceback lines; + handler errors point to the actual Python test helper. Internal control files + are removed while workspace artifacts remain in pytest's temporary directory. +- Commands must not detach, disable error checking, or replace/exit the shell. + These are trusted repository tests, not a security sandbox. +- The runner prepends the checkout and test utilities to `PYTHONPATH`. The separate + installation test removes it and uses a disposable virtualenv to test the + installed package without replacing the checkout environment. +- Manual scenarios use the existing `--run-manual` option. Missing declared + modules or insufficient GPUs produce explicit skips. A GPU profile requires + at least one device. Manual prerequisites remain the operator's responsibility. + +The old `modelopt-doc-test:begin/setup/run/verify/end/skip` format is unsupported. +Test-only setup, assertions, metadata, and exclusion reasons belong in Python. + +## Coverage and limits + +| README | Default smoke scenarios | Opt-in integration scenarios | Exclusions | +| --- | --- | --- | --- | +| `llm_qat` | QAT/QAD/QLoRA in existing end-to-end tests; recipes; simple Llama demo; three Python APIs; FSDP2 layer override; config copy; CLI learning-rate override; dataset cache | Installation; vLLM adapter serving; MMLU/WikiText evaluation | Quickstart fences are externally covered by `test_llm_qat.py`, not duplicated as scenarios | +| `megatron_bridge` | PTQ/export, mock-data distill/export, parameter-target/manual pruning in existing end-to-end tests; CLI preflight | Login; checkpoint import; MLflow PTQ; real-data distillation/QAD; active-parameter/memory NAS targets; VLM pruning; vLLM generation | Migrated fences are externally covered; Docker setup, LoRA integration, incomplete templates and illustrative output remain excluded | + +QAT/QAD/QLoRA quickstarts and dataset pre-tokenization run through `test_llm_qat.py` and its existing shared +PTQ model, dataset cache, backends and export checks. Commands are read from their +README fences, with explicit local model/output paths, test dataset inputs and +two training steps. The remaining QAT-only fragments retain their small workspace +and namespace setup: simple Llama, Python APIs, layer/learning-rate overrides, +configuration copy. The three Python API scenarios expose only the first device +in `CUDA_VISIBLE_DEVICES` (or device 0 when unset) to their worker. This prevents +Trainer from using threaded DataParallel, which races shared attention patching; +existing end-to-end tests retain distributed coverage. Dataset pre-tokenization +uses the existing dataset CLI test. +Manual installation, serving and evaluation +checks remain separate because the existing CLI tests do not exercise them. + +Bridge PTQ/export uses the `readme` case in `test_quantize_export.py`; mock-data +distillation and all-iteration export use `test_distill_llm_hf_export`; pruning +uses the existing `qwen3` case plus a `readme-manual` case in `test_prune_minitron.py`. +These tests use the existing `run_example_command` runner and its fixture inputs, +checkpoint assertions and GPU-count handling. There is no separate Bridge smoke +workspace or second distillation prerequisite. Other architecture and recipe +cases retain their existing commands. Full-scale manual integrations retain their +resource metadata and artifact checks; CLI help is checked by the preflight. + +`_test_utils.examples.readme.readme_command` selects one Python invocation from an +existing named fence and returns an argument list for the existing command runner. +It substitutes explicit test overrides and literal shell defaults, preserving +other published flags. It neither evaluates shell expressions nor runs commands; +unresolved variables fail with the README location. Overrides are single-value +options; general shell workflows and Python fragments still use the snippet runner. +No new README markers are needed for these migrations. + +Checks cover quantization state, training progress/finite losses, changed QAT/QAD +weights, packed NVFP4 exports, adapter artifacts, dataset contents and loadable +Bridge outputs. These tests establish command/API composition; they do not establish +full-model accuracy, convergence, memory requirements or serving performance. +A copied FSDP2 config test checks the copy, not the subsequent human edit. + +## Run and collect + +Name README test modules `test__readme.py` so different examples can be +collected together without module-name collisions. + +Run inside the appropriate example environment: PyTorch 26.07 with the QAT +requirements, or NeMo 26.08 for Megatron Bridge, as in the example CI workflow. +Keep the container's matching Torch/FlashAttention binaries when installing test +dependencies. Install the checkout and its `dev-test` dependencies first. + +```bash +# Parser/runner tests: no Torch, GPU or repository-wide conftest required. +python -m pytest tests/unit/doc_tests --confcutdir=tests/unit/doc_tests -o addopts='' -q + +# Collect and inspect skip reasons in the normal test environment. +python -m pytest tests/examples/llm_qat/test_llm_qat_readme.py --collect-only --no-cov +python -m pytest tests/examples/megatron_bridge/test_megatron_bridge_readme.py --collect-only --no-cov + +# Bounded defaults; manual integrations are skipped. +CUDA_VISIBLE_DEVICES=0 python -m pytest tests/examples/llm_qat/test_llm_qat_readme.py --no-cov -rs +python -m pytest tests/examples/llm_qat/test_llm_qat.py --no-cov -k 'qat_nvfp4 or qad_nvfp4 or qlora_nvfp4' -rs +python -m pytest tests/examples/megatron_bridge/test_quantize_export.py --no-cov -k readme -rs +python -m pytest tests/examples/megatron_bridge/test_distill.py --no-cov -k test_distill_llm_hf_export -rs +python -m pytest tests/examples/megatron_bridge/test_prune_minitron.py --no-cov -k 'qwen3 or readme-manual' -rs + +# Select ONE manual integration with its prerequisites already supplied. +# MODELOPT_DOC_QLORA_CHECKPOINT: exported adapter directory containing base_model/. +python -m pytest tests/examples/llm_qat/test_llm_qat_readme.py --no-cov --run-manual -k llm-qlora-serve +# MODELOPT_DOC_EVAL_CHECKPOINT: fake-quantized QAT checkpoint; lm_eval and dataset access required. +python -m pytest tests/examples/llm_qat/test_llm_qat_readme.py --no-cov --run-manual -k llm-qat-evaluation +``` + +Installation downloads packages into a temporary virtualenv and may compile GPU +extensions. Bridge manual tests need the full models, gated-data access and the +GPU count declared in their scenario. Login requires `HF_TOKEN` and writes into a +temporary `HF_HOME`; no token belongs in Markdown. Real-data training needs +`DATA_PREFIX_1` and `DATA_PREFIX_2`; QAD also needs `QUANTIZED_CHECKPOINT`. MLflow +requires `MLFLOW_TRACKING_URI`; generation requires `PRUNED_CHECKPOINT`. +Manual tests can run for hours and are never enabled by the example CI command. + +## CPU CLI drift checks + +`test_readme_cli_options` compares the flag names in README Python script commands +against each script's `--help`. It checks direct `python`/`python3` commands and +extracts the script arguments after `torchrun` or `accelerate launch` options. +Help runs once per script per pytest process, with a 60-second timeout and CUDA +visibility disabled. Missing scripts, failed help commands and unknown flags +report the README path, snippet ID and command line number. + +Only existing ID-marked fences are inspected; these checks add no README markers +and do not enforce fence coverage. The checks run in the installed example environment but do not launch training, +models or distributed workers. Imports still require the example dependencies. +The QAT evaluation and Bridge vLLM scripts need their optional manual environments; +their CLI checks use `--run-manual`. Values, config contents, launcher flags, +non-Python commands and general shell control flow are outside this check's scope. +Existing execution scenarios continue to validate behavior and output artifacts. + +CI runs the CLI test module first and starts the example suite only if it passes. +A failed preflight preserves the failure status and skips the suite. The main run +excludes the CLI module, so help checks do not run twice. Direct pytest invocations +remain independent; use `&&` locally to apply the same gate: + +```bash +python -m pytest tests/examples/llm_qat/test_llm_qat_readme_cli.py --no-cov -x && \ + python -m pytest tests/examples/llm_qat/test_llm_qat_readme.py --no-cov -v -rs +``` + +Use the corresponding `megatron_bridge` paths in the Bridge container. + +```bash +python -m pytest tests/examples/llm_qat/test_llm_qat_readme_cli.py -k cli_options --no-cov -rs +python -m pytest tests/examples/megatron_bridge/test_megatron_bridge_readme_cli.py -k cli_options --no-cov -rs +``` + +## CI triggers + +Both README test modules are collected by the existing example jobs. The +workflow runs nightly, on demand, and on pushes to copied PR branches +`pull-request/` after unit tests pass. README-only changes explicitly +activate their respective Torch or Megatron lane despite the general Markdown +exclusion. Changes to shared test utilities activate every lane. Migrated Bridge +commands use the existing end-to-end tests' GPU-count fixture, including single-GPU +PR runners and two-GPU nightly runners. Optional full-scale scenarios retain their +own GPU requirements. + +## File responsibilities + +| File | Purpose in this extension | +| --- | --- | +| `parser.py` | Extract named fences and check references, fence coverage, and executable languages. | +| `runner.py` | Execute resolved plans with worker-side handlers, deadlines, shared state, and cleanup. | +| `pytest_utils.py` | Common collection, manual markers and resource checks, independent of the example. | +| `tests/_test_utils/examples/readme.py` | Extract README command arguments for existing end-to-end tests; no separate command runner. | +| `fixtures/llm_qat_readme.py`, `fixtures/megatron_bridge_readme.py` | Worker-side scenario setup and verification moved out of Markdown. | +| `fixtures/llm_qat.py` | Inputs and assertions for the remaining QAT Python/API and standalone snippets. | +| `fixtures/megatron_bridge.py` | Working-directory and export checks for full-scale manual Bridge integrations. | +| `fixtures/llm_qat_installation.py` | QAT requirements in a disposable venv and QAT dependency import verification. | +| `fixtures/llm_qat_serving.py` | QAT QLoRA checkpoint layout, adapter-server readiness and a real completion request. | +| `tests/unit/doc_tests/test_markdown.py` | Focused fence parsing, failure-location, and timeout-cleanup regressions. | +| `tests/examples/llm_qat/test_llm_qat_readme.py` | Collect QAT scenarios through the common pytest adapter. | +| `tests/examples/megatron_bridge/test_megatron_bridge_readme.py` | Collect Bridge scenarios through the same adapter. | +| `examples/llm_qat/README.md` | Annotate the remaining snippets; correct QAD teacher construction and incomplete training commands. | +| `examples/megatron_bridge/README.md` | Annotate smoke/manual scenarios, expose input/scale overrides and account for non-executable templates. | +| `.github/workflows/example_tests.yml` | Trigger the Megatron lane for README-only changes. | +| `tests/examples/README.md` | Point example authors to the shared shell/Python test mechanism. | +| `tests/_test_utils/doc_tests/README.md` | Document the contract, scope, prerequisites, limits and CI behavior. | diff --git a/tests/_test_utils/doc_tests/cli.py b/tests/_test_utils/doc_tests/cli.py new file mode 100644 index 00000000000..64fa0a1304a --- /dev/null +++ b/tests/_test_utils/doc_tests/cli.py @@ -0,0 +1,143 @@ +# 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. + +"""Check README Python CLI option names against each script's public help.""" + +import os +import re +import shlex +import subprocess +import sys +from dataclasses import dataclass +from functools import cache +from pathlib import Path + +from .parser import DocTestError, parse_markdown + +__all__ = ["PythonCommand", "check_cli_options", "python_commands"] + + +@dataclass(frozen=True) +class PythonCommand: + """A README invocation, without executing its arguments or launcher.""" + + snippet: str + line: int + script: Path + flags: tuple[str, ...] + argv: tuple[str, ...] = () + + +def python_commands(path, *, cwd, cwd_overrides=None): + """Extract Python script invocations from simple shell commands and continuations.""" + commands = [] + for snippet in parse_markdown(path).values(): + if snippet.kind != "run": + continue + directory = Path((cwd_overrides or {}).get(snippet.id, cwd)) + logical, start = "", snippet.line + for offset, line in enumerate(snippet.code.splitlines()): + if not logical: + start = snippet.line + offset + logical += line.rstrip().removesuffix("\\") + " " + if line.rstrip().endswith("\\"): + continue + lexer = shlex.shlex(logical, posix=True, punctuation_chars=";&|") + lexer.whitespace_split = True + tokens = list(lexer) + logical = "" + if not tokens: + continue + if tokens[0] == "cd" and len(tokens) == 2: + directory = (directory / tokens[1]).resolve() + continue + launcher = Path(tokens[0]).name + if not ( + re.fullmatch(r"python(?:[23](?:\.\d+)?)?", launcher) + or launcher == "torchrun" + or tokens[:2] == ["accelerate", "launch"] + ): + continue + index = next( + (i for i, token in enumerate(tokens[1:], 1) if token.endswith(".py")), None + ) + if index is None: + continue + if any(token in {";", "&&", "||", "|", "&"} for token in tokens): + raise DocTestError(f"{path}:{start}: {snippet.id}: compound commands unsupported") + flags = [] + for token in tokens[index + 1 :]: + if token == "--": + break + if re.fullmatch(r"--?[a-zA-Z][\w-]*(?:=.*)?", token): + flags.append(token.split("=", 1)[0]) + commands.append( + PythonCommand( + snippet.id, + start, + (directory / tokens[index]).resolve(), + tuple(flags), + tuple(tokens), + ) + ) + return commands + + +@cache +def _help_options(script, timeout): + if not script.is_file(): + raise DocTestError(f"Python script does not exist: {script}") + try: + result = subprocess.run( + [sys.executable, str(script), "--help"], + cwd=script.parent, + env=dict( + os.environ, + CUDA_VISIBLE_DEVICES="", + NO_COLOR="1", + HF_HUB_OFFLINE="1", + HF_DATASETS_OFFLINE="1", + ), + capture_output=True, + text=True, + timeout=timeout, + check=False, + ) + except subprocess.TimeoutExpired as error: + raise DocTestError(f"{script.name} --help timed out after {timeout}s") from error + if result.returncode: + raise DocTestError( + f"{script.name} --help exited with {result.returncode}:\n{result.stdout}{result.stderr}" + ) + options = set() + for line in result.stdout.splitlines(): + if re.match(r"^ -[a-zA-Z-]", line): + definition = re.split(r"\s{2,}", line.strip(), maxsplit=1)[0] + options.update(re.findall(r"(? Path: + """Copy scripts/configs, build local inputs and retain real NVFP4/FSDP2 execution.""" + root = tmp / "checkout" + example = root / "examples" / "llm_qat" + example.mkdir(parents=True) + source = repo / "examples" / "llm_qat" + for script in source.glob("*.py"): + shutil.copy2(script, example / script.name) + shutil.copytree(source / "configs", example / "configs") + (root / "modelopt_recipes").symlink_to(repo / "modelopt_recipes", target_is_directory=True) + model = (create_tiny_llama_dir if llama else create_tiny_qwen3_dir)( + tmp / "model", + with_tokenizer=True, + hidden_size=512, + intermediate_size=512, + num_hidden_layers=2, + ) + (example / "Qwen").mkdir() + (example / "Qwen" / "Qwen3-8B").symlink_to(model, target_is_directory=True) + (example / "meta-llama").mkdir() + (example / "meta-llama/Llama-3.2-3B").symlink_to(model, target_is_directory=True) + data = tmp / "data" + Dataset.from_dict( + {"text": ["The quick brown fox jumps over the lazy dog. " * 4] * 20} + ).save_to_disk(str(data)) + blend = { + "blend_size": 20, + "splits": {"train": 0.8, "eval": 0.2}, + "sources": [ + {"hf_path": str(data), "split": "train", "ratio": 1, "apply_chat_template": False} + ], + } + (example / "configs/dataset/blend.yaml").write_text(yaml.safe_dump(blend)) + for training_path in (example / "configs/train").glob("*_nvfp4.yaml"): + training = yaml.safe_load(training_path.read_text()) + training.update( + model_max_length=128, + train_samples=16, + eval_samples=4, + max_steps=2, + per_device_train_batch_size=1, + per_device_eval_batch_size=1, + gradient_accumulation_steps=1, + eval_steps=1, + save_steps=1, + logging_steps=1, + num_proc=1, + report_to="none", + ) + training_path.write_text(yaml.safe_dump(training)) + devices = int(os.environ.get("MODELOPT_DOC_TEST_GPUS", "1")) + if devices not in (1, 2): + raise ValueError("MODELOPT_DOC_TEST_GPUS must be 1 or 2") + for name in ("fsdp2", "ddp"): + accelerate_path = example / f"configs/accelerate/{name}.yaml" + accelerate = yaml.safe_load(accelerate_path.read_text()) + accelerate.update(num_processes=devices, main_process_port=0) + accelerate_path.write_text(yaml.safe_dump(accelerate)) + print( + f"QAT smoke inputs: 2 layers, hidden/intermediate=512, 20 local text samples; " + f"PTQ length=8192, training length=128, max_steps=2, FSDP2 ranks={devices}; " + "NVFP4 recipe, FlashAttention and Liger unchanged.", + flush=True, + ) + return root + + +def prepare_llm_qat_quantized_input(ctx, *, llama=False): + """Produce a real prerequisite checkpoint for standalone training snippets.""" + root = prepare_llm_qat_workspace(ctx.repo, ctx.tmp, llama=llama) + ctx.cwd = root / "examples/llm_qat" + subprocess.run( + [ + sys.executable, + "quantize.py", + "--model_name_or_path", + "Qwen/Qwen3-8B", + "--dataset_config", + "configs/dataset/blend.yaml", + "--recipe", + "general/ptq/nvfp4_default-kv_fp8", + "--output_dir", + "qwen3-8b-quantized", + ], + cwd=ctx.cwd, + env=ctx.env, + check=True, + ) + return root + + +def prepare_llm_qat_python_api(ctx, *, quantized): + """Supply the model, data and Trainer inputs omitted by the API fragments.""" + from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + TrainingArguments, + default_data_collator, + ) + + import modelopt.torch.opt as mto + import modelopt.torch.quantization as mtq + from modelopt.recipe import load_recipe + from modelopt.torch.quantization.plugins.transformers_trainer import QATTrainer + + root = prepare_llm_qat_workspace(ctx.repo, ctx.tmp) + ctx.cwd = root / "examples/llm_qat" + mto.enable_huggingface_checkpointing() + model = AutoModelForCausalLM.from_pretrained( + ctx.cwd / "Qwen/Qwen3-8B", torch_dtype=torch.bfloat16 + ).cuda() + tokenizer = AutoTokenizer.from_pretrained(ctx.cwd / "Qwen/Qwen3-8B") + tokens = tokenizer("The quick brown fox jumps over the lazy dog.", return_tensors="pt") + + def forward_loop(m): + with torch.no_grad(): + m(**{key: value.cuda() for key, value in tokens.items()}) + + if quantized: + mtq.quantize(model, load_recipe("general/ptq/nvfp4_default-kv_fp8").quantize, forward_loop) + record = {key: value[0].tolist() for key, value in tokens.items()} + record["labels"] = record["input_ids"].copy() + data_module = { + "train_dataset": Dataset.from_list([record] * 4), + "data_collator": default_data_collator, + } + training_args = TrainingArguments( + output_dir=str(ctx.tmp / "api-trained"), + max_steps=2, + per_device_train_batch_size=1, + learning_rate=1e-4, + logging_steps=1, + save_strategy="no", + report_to="none", + bf16=True, + ) + trainer = ( + None + if quantized + else QATTrainer(model=model, processing_class=tokenizer, args=training_args, **data_module) + ) + initial_weights = { + name: parameter.detach().cpu().clone() for name, parameter in model.named_parameters() + } + return { + "initial_weights": initial_weights, + "model": model, + "tokenizer": tokenizer, + "training_args": training_args, + "data_module": data_module, + "forward_loop": forward_loop, + "trainer": trainer, + } + + +def verify_llm_qat_python_api(trainer, initial_weights): + """Assert optimization completed and checkpointing retained quantization state.""" + assert trainer.state.global_step == 2 + losses = [entry["loss"] for entry in trainer.state.log_history if "loss" in entry] + assert losses and all(math.isfinite(value) for value in losses) + assert any( + not torch.equal(initial_weights[name], parameter.detach().cpu()) + for name, parameter in trainer.model.named_parameters() + if name in initial_weights + ) + output = Path(trainer.args.output_dir) + assert (output / "modelopt_state.pth").is_file() + assert (output / "model.safetensors").stat().st_size > 0 diff --git a/tests/_test_utils/doc_tests/fixtures/llm_qat_installation.py b/tests/_test_utils/doc_tests/fixtures/llm_qat_installation.py new file mode 100644 index 00000000000..853ebd31d99 --- /dev/null +++ b/tests/_test_utils/doc_tests/fixtures/llm_qat_installation.py @@ -0,0 +1,55 @@ +# 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. + +"""QAT README installation checks in a disposable environment.""" + +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path + +__all__ = ["prepare_llm_qat_installation", "verify_llm_qat_installation"] + + +def prepare_llm_qat_installation(ctx): + """Prepare the QAT requirements in a disposable venv without altering pytest.""" + venv = ctx.tmp / "venv" + subprocess.run([sys.executable, "-m", "venv", str(venv)], check=True) + ctx.cwd = ctx.tmp + target = ctx.tmp / "examples/llm_qat" + target.mkdir(parents=True) + shutil.copy2(ctx.repo / "examples/llm_qat/requirements.txt", target / "requirements.txt") + ctx.env["PATH"] = str(venv / "bin") + os.pathsep + ctx.env["PATH"] + ctx.env.pop("PYTHONPATH", None) + ctx.env.pop("VIRTUAL_ENV", None) + + +def verify_llm_qat_installation(ctx): + """Import the QAT dependencies and require ModelOpt to come from the new venv.""" + result = subprocess.run( + [ + str(ctx.tmp / "venv/bin/python"), + "-c", + "import json, modelopt, accelerate, peft; print(json.dumps(modelopt.__file__))", + ], + cwd=ctx.cwd, + env=ctx.env, + check=True, + capture_output=True, + text=True, + ) + assert Path(json.loads(result.stdout.strip().splitlines()[-1])).is_relative_to(ctx.tmp / "venv") diff --git a/tests/_test_utils/doc_tests/fixtures/llm_qat_readme.py b/tests/_test_utils/doc_tests/fixtures/llm_qat_readme.py new file mode 100644 index 00000000000..825e8f86b05 --- /dev/null +++ b/tests/_test_utils/doc_tests/fixtures/llm_qat_readme.py @@ -0,0 +1,168 @@ +# 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. + +"""Worker-side setup and assertions for llm_qat README scenarios.""" + +import json +import math +import os +import shutil +from pathlib import Path + +__all__ = [ + "llm_qad_trainer", + "llm_qat_cli_override", + "llm_qat_copy_config", + "llm_qat_evaluation", + "llm_qat_install", + "llm_qat_llama_wrap", + "llm_qat_python", + "llm_qat_simple", + "llm_qat_trainer", + "llm_qlora_serve", +] + + +def llm_qat_install(ctx): + """Run and verify the llm-qat-install scenario.""" + # Optional ML dependencies are imported only inside the scenario worker. + from _test_utils.doc_tests.fixtures.llm_qat_installation import ( + prepare_llm_qat_installation, + verify_llm_qat_installation, + ) + + prepare_llm_qat_installation(ctx) + ctx.run() + verify_llm_qat_installation(ctx) + + +def llm_qat_simple(ctx): + """Run and verify the llm-qat-simple scenario.""" + # Optional ML dependencies are imported only inside the scenario worker. + from _test_utils.doc_tests.fixtures.llm_qat import prepare_llm_qat_workspace + + workspace = prepare_llm_qat_workspace(ctx.repo, ctx.tmp, llama=True) + ctx.cwd = workspace / "examples/llm_qat" + ctx.env.update(HF_HUB_OFFLINE="1", HF_DATASETS_OFFLINE="1", TOKENIZERS_PARALLELISM="false") + ctx.run() + assert (ctx.cwd / "qat_model/modelopt_state.pth").is_file() + assert (ctx.cwd / "qat_model/model.safetensors").stat().st_size > 0 + + +def llm_qlora_serve(ctx): + """Run and verify the llm-qlora-serve scenario.""" + # Optional ML dependencies are imported only inside the scenario worker. + from _test_utils.doc_tests.fixtures.llm_qat_serving import ( + prepare_llm_qat_adapter_server, + verify_llm_qat_adapter_server, + ) + + prepare_llm_qat_adapter_server(ctx) + ctx.run() + verify_llm_qat_adapter_server(ctx) + + +def llm_qat_python(ctx): + """Run and verify the llm-qat-python scenario.""" + # Optional ML dependencies are imported only inside the scenario worker. + from _test_utils.doc_tests.fixtures.llm_qat import ( + prepare_llm_qat_python_api, + verify_llm_qat_python_api, + ) + + namespace = prepare_llm_qat_python_api(ctx, quantized=False) + ctx.run(namespace) + verify_llm_qat_python_api(namespace["trainer"], namespace["initial_weights"]) + + +def llm_qat_trainer(ctx): + """Run and verify the llm-qat-trainer scenario.""" + # Optional ML dependencies are imported only inside the scenario worker. + from _test_utils.doc_tests.fixtures.llm_qat import ( + prepare_llm_qat_python_api, + verify_llm_qat_python_api, + ) + + namespace = prepare_llm_qat_python_api(ctx, quantized=True) + ctx.run(namespace) + verify_llm_qat_python_api(namespace["trainer"], namespace["initial_weights"]) + + +def llm_qad_trainer(ctx): + """Run and verify the llm-qad-trainer scenario.""" + # Optional ML dependencies are imported only inside the scenario worker. + from _test_utils.doc_tests.fixtures.llm_qat import ( + prepare_llm_qat_python_api, + verify_llm_qat_python_api, + ) + + namespace = prepare_llm_qat_python_api(ctx, quantized=True) + ctx.run(namespace) + verify_llm_qat_python_api(namespace["trainer"], namespace["initial_weights"]) + + +def llm_qat_llama_wrap(ctx): + """Run and verify the llm-qat-llama-wrap scenario.""" + # Optional ML dependencies are imported only inside the scenario worker. + from _test_utils.doc_tests.fixtures.llm_qat import prepare_llm_qat_quantized_input + + prepare_llm_qat_quantized_input(ctx, llama=True) + (ctx.cwd / "llama-quantized").symlink_to( + ctx.cwd / "qwen3-8b-quantized", target_is_directory=True + ) + ctx.run() + assert json.loads((ctx.cwd / "llama-qat/trainer_state.json").read_text())["global_step"] == 2 + + +def llm_qat_copy_config(ctx): + """Run and verify the llm-qat-copy-config scenario.""" + ctx.cwd = ctx.tmp + shutil.copytree(ctx.repo / "examples/llm_qat/configs", ctx.tmp / "configs") + ctx.run() + assert (ctx.cwd / "configs/accelerate/fsdp2_llama.yaml").read_bytes() == ( + ctx.cwd / "configs/accelerate/fsdp2.yaml" + ).read_bytes() + + +def llm_qat_cli_override(ctx): + """Run and verify the llm-qat-cli-override scenario.""" + # Optional ML dependencies are imported only inside the scenario worker. + import torch + from _test_utils.doc_tests.fixtures.llm_qat import prepare_llm_qat_quantized_input + + prepare_llm_qat_quantized_input(ctx) + ctx.run() + output = ctx.cwd / "qwen3-8b-qat-override" + state = json.loads((output / "trainer_state.json").read_text()) + assert state["global_step"] == 2 + scheduler = torch.load( + output / "checkpoint-2/scheduler.pt", map_location="cpu", weights_only=True + ) + base_rates = scheduler["base_lrs"] + assert base_rates and all(math.isclose(rate, 5e-5) for rate in base_rates), base_rates + + +def llm_qat_evaluation(ctx): + """Run and verify the llm-qat-evaluation scenario.""" + checkpoint = Path(os.environ["MODELOPT_DOC_EVAL_CHECKPOINT"]).resolve() + assert (checkpoint / "modelopt_state.pth").is_file() + ctx.cwd = ctx.tmp / "examples/llm_qat" + ctx.cwd.mkdir(parents=True) + (ctx.cwd / "qwen3-8b-qat-nvfp4").symlink_to(checkpoint, target_is_directory=True) + shutil.copytree(ctx.repo / "examples/llm_eval", ctx.cwd.parent / "llm_eval") + ctx.env["EVAL_OUTPUT"] = str(ctx.tmp / "eval-results") + ctx.run() + files = list((ctx.tmp / "eval-results").rglob("results*.json")) + assert files and json.loads(files[0].read_text())["results"] diff --git a/tests/_test_utils/doc_tests/fixtures/llm_qat_serving.py b/tests/_test_utils/doc_tests/fixtures/llm_qat_serving.py new file mode 100644 index 00000000000..4e720ffe85b --- /dev/null +++ b/tests/_test_utils/doc_tests/fixtures/llm_qat_serving.py @@ -0,0 +1,62 @@ +# 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. + +"""QAT README QLoRA checkpoint setup and vLLM adapter verification.""" + +import json +import socket +import time +import urllib.error +import urllib.request +from pathlib import Path + +__all__ = ["prepare_llm_qat_adapter_server", "verify_llm_qat_adapter_server"] + + +def prepare_llm_qat_adapter_server(ctx): + """Link the QLoRA export under the checkpoint name used by the QAT README.""" + checkpoint = Path(ctx.env["MODELOPT_DOC_QLORA_CHECKPOINT"]).resolve() + assert (checkpoint / "adapter_config.json").is_file() + assert (checkpoint / "base_model/config.json").is_file() + ctx.cwd = ctx.tmp + (ctx.cwd / "qwen3-8b-fp4-qlora-hf").symlink_to(checkpoint, target_is_directory=True) + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + ctx.env["PORT"] = str(sock.getsockname()[1]) + + +def verify_llm_qat_adapter_server(ctx): + """Require readiness and a completion from the QAT README's "adapter" model.""" + url = f"http://127.0.0.1:{ctx.env['PORT']}" + deadline = time.monotonic() + 600 + while True: + try: + with urllib.request.urlopen(url + "/health", timeout=5) as response: + assert response.status == 200 + break + except OSError: + if time.monotonic() >= deadline: + raise TimeoutError("vLLM did not become ready") from None + time.sleep(1) + request = urllib.request.Request( + url + "/v1/completions", + data=json.dumps( + {"model": "adapter", "prompt": "The capital of France is", "max_tokens": 4} + ).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=60) as response: + result = json.load(response) + assert result["choices"] and result["usage"]["completion_tokens"] > 0 diff --git a/tests/_test_utils/doc_tests/fixtures/megatron_bridge.py b/tests/_test_utils/doc_tests/fixtures/megatron_bridge.py new file mode 100644 index 00000000000..7efd33cce41 --- /dev/null +++ b/tests/_test_utils/doc_tests/fixtures/megatron_bridge.py @@ -0,0 +1,38 @@ +# 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. + +"""Local inputs and artifact checks for real Megatron Bridge README commands.""" + +import json +from pathlib import Path + +__all__ = [ + "prepare_bridge_full_scale", + "verify_bridge_hf_export", +] + + +def prepare_bridge_full_scale(ctx): + """Set the Bridge working/output directories; callers supply full-scale inputs.""" + ctx.cwd = ctx.repo / "examples/megatron_bridge" + ctx.env["OUTPUT_ROOT"] = str(ctx.tmp / "outputs") + Path(ctx.env["OUTPUT_ROOT"]).mkdir() + + +def verify_bridge_hf_export(path): + """Check config and weight-file presence for full-scale Bridge HF exports.""" + assert json.loads((path / "config.json").read_text())["model_type"] + weights = list(path.glob("*.safetensors")) + assert weights and all(weight.stat().st_size > 0 for weight in weights) diff --git a/tests/_test_utils/doc_tests/fixtures/megatron_bridge_readme.py b/tests/_test_utils/doc_tests/fixtures/megatron_bridge_readme.py new file mode 100644 index 00000000000..2e86b8ca3b6 --- /dev/null +++ b/tests/_test_utils/doc_tests/fixtures/megatron_bridge_readme.py @@ -0,0 +1,121 @@ +# 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. + +"""Worker-side setup and assertions for megatron_bridge README scenarios.""" + +import json +import os +from pathlib import Path + +from _test_utils.doc_tests.fixtures.megatron_bridge import ( + prepare_bridge_full_scale, + verify_bridge_hf_export, +) + +__all__ = [ + "bridge_generation", + "bridge_import", + "bridge_login", + "bridge_mlflow", + "bridge_prune_active", + "bridge_prune_memory", + "bridge_prune_vlm", + "bridge_qad", + "bridge_real_data", +] + + +def bridge_login(ctx): + """Run and verify the bridge-login scenario.""" + assert os.environ.get("HF_TOKEN"), "Set HF_TOKEN for the manual login test" + ctx.cwd = ctx.tmp + ctx.env["HF_HOME"] = str(ctx.tmp / "hf-home") + ctx.run() + assert (ctx.tmp / "hf-home/token").stat().st_size > 0 + + +def bridge_import(ctx): + """Run and verify the bridge-import scenario.""" + + prepare_bridge_full_scale(ctx) + output = Path(ctx.env["OUTPUT_ROOT"]) + ctx.run() + assert (output / "Qwen3-8B-megatron").is_dir() + + +def bridge_mlflow(ctx): + """Run and verify the bridge-mlflow scenario.""" + + prepare_bridge_full_scale(ctx) + ctx.run() + pointer = Path(ctx.env["OUTPUT_ROOT"]) / "Qwen3-8B-NVFP4-megatron/.experiment.json" + assert json.loads(pointer.read_text()) + + +def bridge_real_data(ctx): + """Run and verify the bridge-real-data scenario.""" + + prepare_bridge_full_scale(ctx) + ctx.run() + checkpoints = Path(ctx.env["OUTPUT_ROOT"]) / "qwen3_8b_to_4b_distill/checkpoints" + assert (checkpoints / "latest_checkpointed_iteration.txt").is_file() + assert list(checkpoints.glob("iter_*")) + + +def bridge_qad(ctx): + """Run and verify the bridge-qad scenario.""" + # Import the optional Megatron dependency only in this worker. + from _test_utils.torch.megatron.modelopt_state import assert_has_modelopt_state + + prepare_bridge_full_scale(ctx) + ctx.run() + checkpoints = Path(ctx.env["OUTPUT_ROOT"]) / "qwen3_8b_nvfp4_qad/checkpoints" + assert (checkpoints / "latest_checkpointed_iteration.txt").is_file() + assert list(checkpoints.glob("iter_*")) + + assert_has_modelopt_state(checkpoints) + + +def bridge_prune_active(ctx): + """Run and verify the bridge-prune-active scenario.""" + + prepare_bridge_full_scale(ctx) + output = Path(ctx.env["OUTPUT_ROOT"]) + ctx.run() + verify_bridge_hf_export(output / "Nemotron-3-Nano-30B-A3B-BF16-Pruned-3B-Active") + + +def bridge_prune_memory(ctx): + """Run and verify the bridge-prune-memory scenario.""" + + prepare_bridge_full_scale(ctx) + output = Path(ctx.env["OUTPUT_ROOT"]) + ctx.run() + verify_bridge_hf_export(output / "Qwen3-8B-Pruned-12GB") + + +def bridge_prune_vlm(ctx): + """Run and verify the bridge-prune-vlm scenario.""" + + prepare_bridge_full_scale(ctx) + output = Path(ctx.env["OUTPUT_ROOT"]) + ctx.run() + verify_bridge_hf_export(output / "Qwen3.5-4B-Pruned-3B") + + +def bridge_generation(ctx): + """Run and verify the bridge-generation scenario.""" + ctx.cwd = ctx.repo / "examples/megatron_bridge" + ctx.run() diff --git a/tests/_test_utils/doc_tests/parser.py b/tests/_test_utils/doc_tests/parser.py new file mode 100644 index 00000000000..3b5c19cf51f --- /dev/null +++ b/tests/_test_utils/doc_tests/parser.py @@ -0,0 +1,165 @@ +# 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. + +"""Extract named Markdown snippets and resolve Python-owned test scenarios.""" + +import re +import textwrap +from dataclasses import dataclass, field, replace +from pathlib import Path + +__all__ = ["DocTestError", "Scenario", "Snippet", "parse_markdown", "resolve_scenarios"] + +_ID = r"[a-z0-9][a-z0-9-]*" + + +class DocTestError(ValueError): + """A Markdown snippet or scenario is malformed or failed to execute.""" + + +@dataclass +class Snippet: + """A named code fence with its original Markdown location.""" + + id: str + kind: str + code: str + line: int + + +@dataclass +class Scenario: + """An ordered snippet selection with worker-side orchestration and resource limits.""" + + id: str + snippets: list[str] + profile: str = "cpu" + timeout_seconds: int = 60 + handler: str | None = None + manual: bool = False + min_gpus: int = 0 + requires: list[str] = field(default_factory=list) + background: bool = False + path: Path | None = None + steps: list[Snippet] = field(default_factory=list) + + @property + def line(self): + return self.steps[0].line + + +def parse_markdown(path: Path, *, require_ids: bool = False) -> dict[str, Snippet]: + """Extract ID-marked fences without executing code; optionally require IDs on every fence.""" + path = Path(path).resolve() + lines = path.read_text().splitlines(keepends=True) + snippets = {} + pending = None + i = 0 + + def fail(message, line): + raise DocTestError(f"{path}:{line}: {message}") + + def unquote(text): + return re.sub(r"^ {0,3}> ?", "", text) + + while i < len(lines): + line = i + 1 + text = unquote(lines[i]).strip() + fence = re.fullmatch(r"(`{3,}|~{3,})([^`~]*)", text) + if fence: + delimiter, language = fence.groups() + quoted = lines[i].lstrip().startswith(">") + start = i + 1 + i += 1 + while i < len(lines) and not re.fullmatch( + re.escape(delimiter[0]) + "{" + str(len(delimiter)) + r",}\s*", + (unquote(lines[i]) if quoted else lines[i]).strip(), + ): + i += 1 + if i == len(lines) and (pending or require_ids): + fail("unclosed snippet fence", line) + if pending: + code = textwrap.dedent( + "".join(unquote(part) if quoted else part for part in lines[start:i]) + ) + language = language.strip() + kind = {"bash": "run", "sh": "run", "python": "python"}.get( + language, "unsupported:" + language + ) + snippets[pending] = Snippet(pending, kind, code, start + 1) + pending = None + elif require_ids: + fail("unaccounted fence: add a snippet ID and a test or exclusion", line) + i += 1 + continue + if pending and text: + fail("snippet ID must be immediately followed by a fence (blank lines allowed)", line) + if "modelopt-doc-test:" in text: + fail( + "legacy doc-test directives are unsupported; use snippet IDs and Python tests", line + ) + if "modelopt-snippet:" in text: + match = re.fullmatch(r"", text) + if not match: + fail("invalid snippet ID marker", line) + pending = match.group(1) + if pending in snippets: + fail(f"duplicate snippet ID: {pending}", line) + i += 1 + if pending: + fail(f"snippet ID {pending} missing its fence", len(lines)) + return snippets + + +def resolve_scenarios(path, scenarios, *, excluded=None): + """Validate references and account for every fence before any scenario can execute.""" + path = Path(path).resolve() + snippets = parse_markdown(path, require_ids=True) + excluded = {} if excluded is None else excluded + unknown = excluded.keys() - snippets.keys() + if unknown: + raise DocTestError(f"{path}: unknown excluded snippets: {sorted(unknown)}") + result, covered = [], set() + for scenario in scenarios: + + def fail(message): + raise DocTestError(f"{path}: scenario {scenario.id}: {message}") + + missing = set(scenario.snippets) - snippets.keys() + if missing: + fail(f"unknown snippets: {sorted(missing)}") + overlap = set(scenario.snippets) & excluded.keys() + if overlap: + fail(f"snippets are both tested and excluded: {sorted(overlap)}") + steps = [snippets[key] for key in scenario.snippets] + kinds = {step.kind for step in steps} + if not kinds <= {"run", "python"} or len(kinds) != 1: + fail("select only shell fences or only Python fences in each scenario") + if scenario.background: + if kinds != {"run"} or scenario.handler is None: + fail("background requires shell snippets and a verification handler") + steps = [*steps[:-1], replace(steps[-1], kind="background")] + for step in steps: + if step.kind == "python": + try: + compile("\n" * (step.line - 1) + step.code, str(path), "exec") + except SyntaxError as error: + fail(f"invalid Python in {step.id} at line {error.lineno}: {error.msg}") + covered.update(scenario.snippets) + result.append(replace(scenario, path=path, steps=steps)) + unaccounted = snippets.keys() - covered - excluded.keys() + if unaccounted: + raise DocTestError(f"{path}: unaccounted snippets: {sorted(unaccounted)}") + return result diff --git a/tests/_test_utils/doc_tests/pytest_utils.py b/tests/_test_utils/doc_tests/pytest_utils.py new file mode 100644 index 00000000000..036b6cc12ce --- /dev/null +++ b/tests/_test_utils/doc_tests/pytest_utils.py @@ -0,0 +1,64 @@ +# 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. + +"""Common pytest collection and resource checks for executable documentation.""" + +import importlib.util + +import pytest + +from .parser import resolve_scenarios +from .runner import run_scenario + +__all__ = ["execute_scenario", "scenario_parameters"] + + +def scenario_parameters(path, scenarios, *, excluded=None): + """Collect every scenario, enforcing a documented disposition for every fence.""" + result = [] + for scenario in resolve_scenarios(path, scenarios, excluded=excluded): + marks = [pytest.mark.timeout(scenario.timeout_seconds + 30)] + if scenario.manual: + marks.append(pytest.mark.manual) + result.append(pytest.param(scenario, id=scenario.id, marks=marks)) + return result + + +def execute_scenario(scenario, repo, tmp_path, monkeypatch): + """Check declared resources and execute the same runner for any example.""" + for dependency in scenario.requires: + try: + available = importlib.util.find_spec(dependency) + except ModuleNotFoundError: + available = None + if available is None: + pytest.skip(f"README scenario requires {dependency}") + if scenario.profile == "gpu" or scenario.min_gpus: + # GPU discovery is deliberately deferred until execution, not collection. + import torch + + needed = max(1, scenario.min_gpus) + if torch.cuda.device_count() < needed: + pytest.skip(f"README scenario requires {needed} CUDA devices") + for key in ( + "RANK", + "LOCAL_RANK", + "WORLD_SIZE", + "LOCAL_WORLD_SIZE", + "MASTER_ADDR", + "MASTER_PORT", + ): + monkeypatch.delenv(key, raising=False) + print(run_scenario(scenario, repo, tmp_path)) diff --git a/tests/_test_utils/doc_tests/runner.py b/tests/_test_utils/doc_tests/runner.py new file mode 100644 index 00000000000..92a1bbb72ec --- /dev/null +++ b/tests/_test_utils/doc_tests/runner.py @@ -0,0 +1,189 @@ +# 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. + +"""Run trusted documentation in an isolated process group with a whole-scenario deadline.""" + +import contextlib +import importlib +import json +import os +import shlex +import signal +import subprocess +import sys +import tempfile +import time +from dataclasses import asdict +from pathlib import Path +from types import SimpleNamespace + +from .parser import DocTestError, Scenario, Snippet + +__all__ = ["run_scenario"] + + +def run_scenario(scenario, repo: Path, tmp: Path) -> str: + """Execute one scenario and return captured output; terminate its group on every exit.""" + if scenario.path is None or not scenario.steps: + raise DocTestError("Resolve snippet references before running a scenario") + repo = Path(repo).resolve() + tmp = Path(tmp).resolve() + tmp.mkdir(parents=True, exist_ok=True) + environment = dict(os.environ) + environment["PYTHONPATH"] = os.pathsep.join( + [ + str(Path(repo).resolve()), + str(Path(__file__).resolve().parents[2]), + environment.get("PYTHONPATH", ""), + ] + ) + with ( + tempfile.TemporaryDirectory(prefix="doc-test-", dir=tmp) as control_dir, + tempfile.TemporaryFile() as output, + ): + manifest = Path(control_dir) / "scenario.json" + manifest.write_text(json.dumps(asdict(scenario) | {"path": str(scenario.path)})) + process = subprocess.Popen( + [ + sys.executable, + "-u", + "-m", + "_test_utils.doc_tests.runner", + str(manifest), + str(Path(repo).resolve()), + str(tmp), + control_dir, + ], + env=environment, + stdout=output, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + timed_out = False + try: + process.wait(timeout=scenario.timeout_seconds) + except subprocess.TimeoutExpired: + timed_out = True + finally: + # Give launchers time to forward termination to their own worker groups. + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGTERM) + deadline = time.monotonic() + 3 + while time.monotonic() < deadline: + process.poll() + os.killpg(process.pid, 0) + time.sleep(0.05) + os.killpg(process.pid, signal.SIGKILL) + process.wait() + completed = (Path(control_dir) / "completed").is_file() + output.seek(0) + text = output.read().decode(errors="replace") + for index, step in enumerate(scenario.steps): + if step.kind in {"run", "background"}: + text = text.replace( + str(Path(control_dir) / f"fence-{index}.sh"), str(scenario.path) + ) + if timed_out: + reason = f"timed out after {scenario.timeout_seconds}s" + elif process.returncode: + reason = f"exit {process.returncode}" + elif not completed: + reason = "worker exited before completing the scenario" + else: + return text + raise DocTestError(f"{scenario.path}:{scenario.line}: {scenario.id}: {reason}\n{text}") + + +def _execute(scenario, repo, tmp, control): + ctx = SimpleNamespace(repo=repo, tmp=tmp, cwd=repo, env=dict(os.environ)) + called = completed = False + + def run(namespace=None): + nonlocal called, completed + if called: + raise DocTestError("A scenario handler must call ctx.run() exactly once") + called = True + namespace = {} if namespace is None else namespace + namespace.update(ctx=ctx, __name__="__main__", __file__=str(scenario.path)) + if scenario.steps[0].kind == "python": + os.chdir(ctx.cwd) + os.environ.clear() + os.environ.update(ctx.env) + for step in scenario.steps: + print(f"{scenario.path}:{step.line}: {step.id}", flush=True) + exec( + compile("\n" * (step.line - 1) + step.code, str(scenario.path), "exec"), + namespace, + ) + ctx.cwd, ctx.env = Path.cwd(), dict(os.environ) + else: + _execute_shell(scenario, ctx, control) + completed = True + return namespace + + ctx.run = run + if scenario.handler: + module, name = scenario.handler.split(":") + handler = getattr(importlib.import_module(module), name) + handler(ctx) + else: + ctx.run() + if not completed: + raise DocTestError("Scenario handler did not complete ctx.run()") + + +def _execute_shell(scenario, ctx, control): + script = ["set -Eeuo pipefail", f"DOC_TEST_SOURCE={shlex.quote(str(scenario.path))}"] + for index, step in enumerate(scenario.steps): + if step.kind not in {"run", "background"}: + continue + source = control / f"fence-{index}.sh" + source.write_text("\n" * (step.line - 1) + step.code) + script.append(f"printf '%s\\n' {shlex.quote(f'{scenario.path}:{step.line}: {step.id}')}") + script.append(f"DOC_TEST_SNIPPET={shlex.quote(step.id)}") + script.append("set -Eeuo pipefail") + script.append( + 'trap \'printf "%s:%s: %s: shell command failed\\n" ' + '"$DOC_TEST_SOURCE" "$LINENO" "$DOC_TEST_SNIPPET" >&2\' ERR' + ) + command = f"source {shlex.quote(str(source))}" + if step.kind == "background": + command += f" &\necho $! > {shlex.quote(str(control / 'background_pid'))}" + script.append(command) + script.extend( + [ + f"pwd -P > {shlex.quote(str(control / 'cwd'))}", + f"env -0 > {shlex.quote(str(control / 'env'))}", + ] + ) + shell = control / "scenario.sh" + shell.write_text("\n".join(script) + "\n") + # Inherit stdout so the parent retains partial output even if the scenario times out. + result = subprocess.run(["bash", "--noprofile", "--norc", str(shell)], cwd=ctx.cwd, env=ctx.env) + if result.returncode: + raise DocTestError(f"shell exited with {result.returncode}") + ctx.cwd = Path((control / "cwd").read_text().rstrip("\n")) + ctx.env = dict(item.split("=", 1) for item in (control / "env").read_text().split("\0") if item) + if (control / "background_pid").exists(): + ctx.background_pid = int((control / "background_pid").read_text()) + + +if __name__ == "__main__": + manifest, root, workspace, control_dir = sys.argv[1:] + payload = json.loads(Path(manifest).read_text()) + payload["path"] = Path(payload["path"]) + payload["steps"] = [Snippet(**step) for step in payload["steps"]] + _execute(Scenario(**payload), Path(root), Path(workspace), Path(control_dir)) + (Path(control_dir) / "completed").touch() diff --git a/tests/_test_utils/examples/readme.py b/tests/_test_utils/examples/readme.py new file mode 100644 index 00000000000..222edfa5181 --- /dev/null +++ b/tests/_test_utils/examples/readme.py @@ -0,0 +1,77 @@ +# 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. + +"""Read documented Python commands for the existing example test runners.""" + +import re +from pathlib import Path + +from _test_utils.doc_tests.cli import python_commands +from _test_utils.doc_tests.parser import DocTestError + +__all__ = ["readme_command"] + + +def readme_command(path, snippet, script, *, overrides=None, cwd_overrides=None): + """Extract one invocation, replacing explicit CI inputs without executing shell text.""" + path = Path(path) + commands = [ + command + for command in python_commands(path, cwd=path.parent, cwd_overrides=cwd_overrides) + if command.snippet == snippet and command.script.name == script + ] + if len(commands) != 1: + raise DocTestError( + f"{path}: {snippet}: expected one {script} invocation, found {len(commands)}" + ) + command = commands[0] + options = { + key if key.startswith("-") else "--" + key: value + for key, value in (overrides or {}).items() + } + result = [] + tokens = list(command.argv) + index = 0 + while index < len(tokens): + token = tokens[index] + option = token.split("=", 1)[0] + if option not in options: + result.append(token) + index += 1 + continue + value = options.pop(option) + result.extend([option, str(value)]) + index += 1 + if "=" not in token: + if index < len(tokens) and not re.match(r"--?[A-Za-z]", tokens[index]): + index += 1 + for option, value in options.items(): + result.extend([option, str(value)]) + + def expand_default(match): + default = match.group(1) + if default is None: + raise DocTestError( + f"{path}:{command.line}: {snippet}: unresolved shell variable {match.group(0)}" + ) + return default + + expanded = [ + re.sub(r"\$\{\w+(?::-([^{}]*))?\}|\$[A-Za-z_]\w*", expand_default, token) + for token in result + ] + if any("$" in token for token in expanded): + raise DocTestError(f"{path}:{command.line}: {snippet}: unresolved shell expression") + return expanded diff --git a/tests/examples/README.md b/tests/examples/README.md index 184f1d4344d..a7e70a172b2 100644 --- a/tests/examples/README.md +++ b/tests/examples/README.md @@ -33,3 +33,10 @@ pytest tests/examples/diffusers/test_diffusers_onnx_trt.py The following environment variables can be set to control the behavior of the tests: - `MODELOPT_LOCAL_MODEL_ROOT`: If set, the tests will use the local model directory instead of downloading the model from the internet. Default is not set, which means the model will be downloaded. + +## Executable README scenarios + +The LLM QAT and Megatron Bridge examples reference ID-marked Markdown shell +and Python fences from Python scenario definitions, with separate opt-in integrations for costly workflows. See +[the snippet contract and CPU/GPU commands](../_test_utils/doc_tests/README.md) +for local development, collection, reduced-scale inputs and acceptance checks. diff --git a/tests/examples/llm_qat/test_llm_qat.py b/tests/examples/llm_qat/test_llm_qat.py index b9c7ac09624..314ecfb586e 100644 --- a/tests/examples/llm_qat/test_llm_qat.py +++ b/tests/examples/llm_qat/test_llm_qat.py @@ -15,14 +15,26 @@ import json +import math +from pathlib import Path import pytest import torch +from _test_utils.examples.readme import readme_command from _test_utils.examples.run_command import run_example_command from _test_utils.fs_utils import assert_unmodified_tree from _test_utils.torch.transformers_models import create_tiny_qwen3_dir +from datasets import load_from_disk from safetensors.torch import load_file +REPO = Path(__file__).resolve().parents[3] +README = REPO / "examples/llm_qat/README.md" +README_SCENARIOS = { + "configs/train/qat_nvfp4.yaml": "llm-qat-quickstart", + "configs/train/qad_nvfp4.yaml": "llm-qad-quickstart", + "configs/train/qlora_nvfp4.yaml": "llm-qlora-quickstart", +} + # Mapping from backend name to accelerate config file BACKEND_CONFIGS = { "fsdp2": "configs/accelerate/fsdp2.yaml", @@ -38,72 +50,96 @@ # marked nightly. QLoRA still trains on ddp on PRs. # Fast training overrides (short runs with frequent eval) -FAST_TRAIN_ARGS = [ - "--model_max_length", - "128", - "--num_train_epochs", - "1.0", - "--save_steps", - "5", - "--eval_steps", - "5", -] - - -# fmt: off -def _fast_data_args(cache_dir: str) -> list[str]: - """Fast dataset overrides for all tests (small samples, no shuffle, temp cache).""" - return [ - "--dataset_config", "configs/dataset/blend_test.yaml", - "--train_samples", "64", - "--eval_samples", "16", - "--shuffle", "False", - "--dataset_cache_dir", cache_dir, +FAST_TRAIN_ARGS = { + "model_max_length": "128", + "num_train_epochs": "1.0", + "max_steps": "2", + "save_steps": "1", + "eval_steps": "1", +} + + +def _assert_training_progress(output): + state = json.loads((output / "trainer_state.json").read_text()) + assert state["global_step"] == 2 + losses = [ + entry[key] + for entry in state["log_history"] + for key in ("loss", "eval_loss", "train_loss") + if key in entry ] + assert losses and all(math.isfinite(value) for value in losses) + assert any("train_loss" in entry for entry in state["log_history"]) + + +def _assert_qat_export(ptq, trained, exported): + _assert_training_progress(trained) + assert (trained / "modelopt_state.pth").stat().st_size > 0 + before = load_file(str(Path(ptq) / "model.safetensors")) + after = load_file(str(trained / "model.safetensors")) + keys = [key for key in before.keys() & after.keys() if key.endswith("weight")] + assert keys and any(not torch.equal(before[key], after[key]) for key in keys) + config = json.loads((exported / "hf_quant_config.json").read_text()) + assert config["quantization"]["quant_algo"] == "NVFP4" + assert any( + value.dtype == torch.uint8 + for value in load_file(str(exported / "model.safetensors")).values() + ) -def _run_quantize(config: str, extra_cmd_args: list[str], cache_dir: str = ""): +def _fast_data_args(cache_dir: str): + """Fast dataset overrides for all tests (small samples, no shuffle, temp cache).""" + return { + "dataset_config": "configs/dataset/blend_test.yaml", + "train_samples": "64", + "eval_samples": "16", + "shuffle": "False", + "dataset_cache_dir": cache_dir, + } + + +def _readme_command(config, script, overrides): + return readme_command( + README, + README_SCENARIOS[config], + script, + overrides=overrides, + cwd_overrides={"llm-qat-quickstart": REPO}, + ) + + +def _run_quantize(config: str, overrides, cache_dir: str = ""): run_example_command( - [ - "python", "quantize.py", - "--config", config, - *_fast_data_args(cache_dir), - *extra_cmd_args, - ], + _readme_command(config, "quantize.py", {**_fast_data_args(cache_dir), **overrides}), "llm_qat", ) -def _run_train(config: str, extra_cmd_args: list[str], backend: str = "fsdp2", cache_dir: str = ""): - config_file = BACKEND_CONFIGS[backend] - gradient_args = ( - ["--gradient_checkpointing", "True"] - if backend in GRADIENT_CHECKPOINTING_BACKENDS - else [] - ) +def _run_train(config: str, overrides, backend: str = "fsdp2", cache_dir: str = ""): + args = { + "config-file": BACKEND_CONFIGS[backend], + **_fast_data_args(cache_dir), + **FAST_TRAIN_ARGS, + } + if backend in GRADIENT_CHECKPOINTING_BACKENDS: + args["gradient_checkpointing"] = "True" run_example_command( - [ - "accelerate", "launch", - "--config-file", config_file, - "train.py", - "--config", config, - *_fast_data_args(cache_dir), - *FAST_TRAIN_ARGS, - *gradient_args, - *extra_cmd_args, - ], + _readme_command(config, "train.py", {**args, **overrides}), "llm_qat", setup_free_port=True, ) -def _run_export(ckpt_dir: str, export_dir: str): +def _run_export(ckpt_dir: str, export_dir: str, config="configs/train/qat_nvfp4.yaml"): run_example_command( - [ - "python", "export.py", - "--pyt_ckpt_path", ckpt_dir, - "--export_path", export_dir, - ], + _readme_command( + config, + "export.py", + { + "pyt_ckpt_path": ckpt_dir, + "export_path": export_dir, + }, + ), "llm_qat", ) @@ -136,18 +172,12 @@ def tied_qwen3_path(tmp_path_factory): def nvfp4_ptq_dir(tied_qwen3_path, tmp_path_factory, dataset_cache_dir): """NVFP4 PTQ checkpoint shared by the QAT and QAD tests, which only read it. - quantize.py ignores the QAD-only config keys, so qat_nvfp4.yaml and qad_nvfp4.yaml give the - same arguments and one run serves both. + The QAT and QAD README quickstarts use the same PTQ command, so one run serves both. """ ptq_dir = tmp_path_factory.mktemp("ptq") _run_quantize( "configs/train/qat_nvfp4.yaml", - [ - "--model_name_or_path", tied_qwen3_path, - "--recipe", "general/ptq/nvfp4_default-kv_fp8", - "--calib_size", "64", - "--output_dir", str(ptq_dir), - ], + {"model_name_or_path": tied_qwen3_path, "calib_size": "64", "output_dir": str(ptq_dir)}, cache_dir=dataset_cache_dir, ) with assert_unmodified_tree(ptq_dir) as path: @@ -158,35 +188,44 @@ def test_dataset_utils_pretokenize(tiny_qwen3_path, tmp_path): """Test dataset_utils.py standalone CLI pre-tokenization.""" cache_dir = tmp_path / "dataset_cache" run_example_command( - [ - "python", "dataset_utils.py", - *_fast_data_args(str(cache_dir)), - "--model_name_or_path", tiny_qwen3_path, - ], + readme_command( + README, + "llm-qat-dataset", + "dataset_utils.py", + overrides={**_fast_data_args(str(cache_dir)), "model_name_or_path": tiny_qwen3_path}, + ), "llm_qat", ) assert cache_dir.exists(), "Cache directory should be created" assert any(cache_dir.iterdir()), "Cache directory should contain tokenized data" - - -@pytest.mark.parametrize("backend", [ - "fsdp2", - pytest.param("deepspeed", marks=pytest.mark.nightly), - pytest.param("ddp", marks=pytest.mark.nightly), -]) + caches = list(cache_dir.glob("*/dataset_dict.json")) + assert len(caches) == 1 + cached = load_from_disk(str(caches[0].parent)) + assert len(cached["train"]) > 0 and len(cached["eval"]) > 0 + assert {"input_ids", "attention_mask", "labels"} <= set(cached["train"].column_names) + + +@pytest.mark.parametrize( + "backend", + [ + "fsdp2", + pytest.param("deepspeed", marks=pytest.mark.nightly), + pytest.param("ddp", marks=pytest.mark.nightly), + ], +) def test_qwen3_qat_nvfp4(nvfp4_ptq_dir, dataset_cache_dir, tmp_path, backend): qat_output_dir = tmp_path / "qat" _run_train( "configs/train/qat_nvfp4.yaml", - [ - "--model_name_or_path", nvfp4_ptq_dir, - "--do_train", "True", - "--output_dir", str(qat_output_dir), - ], + {"model_name_or_path": nvfp4_ptq_dir, "output_dir": str(qat_output_dir)}, backend=backend, cache_dir=dataset_cache_dir, ) + export_dir = tmp_path / "qat_export" + _run_export(str(qat_output_dir), str(export_dir)) + _assert_qat_export(nvfp4_ptq_dir, qat_output_dir, export_dir) + @pytest.mark.skip(reason="FSDP2 LoRA checkpoint save omits adapter_model.safetensors") def test_qwen3_lora_qat_nvfp4(tiny_qwen3_path, tmp_path): @@ -196,12 +235,11 @@ def test_qwen3_lora_qat_nvfp4(tiny_qwen3_path, tmp_path): # Step 1: Quantize _run_quantize( "configs/train/qat_nvfp4.yaml", - [ - "--model_name_or_path", tiny_qwen3_path, - "--recipe", "general/ptq/nvfp4_default-kv_fp8", - "--calib_size", "64", - "--output_dir", str(ptq_output_dir), - ], + { + "model_name_or_path": tiny_qwen3_path, + "calib_size": "64", + "output_dir": str(ptq_output_dir), + }, cache_dir=cache_dir, ) @@ -209,12 +247,12 @@ def test_qwen3_lora_qat_nvfp4(tiny_qwen3_path, tmp_path): lora_qat_output_dir = tmp_path / "lora_qat" _run_train( "configs/train/qat_nvfp4.yaml", - [ - "--model_name_or_path", str(ptq_output_dir), - "--do_train", "True", - "--lora", "True", - "--output_dir", str(lora_qat_output_dir), - ], + { + "model_name_or_path": str(ptq_output_dir), + "do_train": "True", + "lora": "True", + "output_dir": str(lora_qat_output_dir), + }, backend="fsdp2", cache_dir=cache_dir, ) @@ -244,26 +282,30 @@ def test_qwen3_lora_qat_nvfp4(tiny_qwen3_path, tmp_path): assert torch.equal(base_weights[key], reference[key]), f"{key} does not match PTQ export" -@pytest.mark.parametrize("backend", [ - "fsdp2", - pytest.param("deepspeed", marks=pytest.mark.nightly), -]) +@pytest.mark.parametrize( + "backend", + [ + "fsdp2", + pytest.param("deepspeed", marks=pytest.mark.nightly), + ], +) def test_qwen3_qad_nvfp4(tied_qwen3_path, nvfp4_ptq_dir, dataset_cache_dir, tmp_path, backend): qad_output_dir = tmp_path / "qad" # QAD (quantization-aware distillation) of the shared quantized student _run_train( "configs/train/qad_nvfp4.yaml", - [ - "--model_name_or_path", nvfp4_ptq_dir, - "--do_train", "True", - "--output_dir", str(qad_output_dir), - "--distill", "True", - "--teacher_model", tied_qwen3_path, - ], + { + "model_name_or_path": nvfp4_ptq_dir, + "output_dir": str(qad_output_dir), + "teacher_model": tied_qwen3_path, + }, backend=backend, cache_dir=dataset_cache_dir, ) + export_dir = tmp_path / "qad_export" + _run_export(str(qad_output_dir), str(export_dir), config="configs/train/qad_nvfp4.yaml") + _assert_qat_export(nvfp4_ptq_dir, qad_output_dir, export_dir) def test_qwen3_qlora_nvfp4(tiny_qwen3_path, dataset_cache_dir, tmp_path): @@ -272,13 +314,11 @@ def test_qwen3_qlora_nvfp4(tiny_qwen3_path, dataset_cache_dir, tmp_path): # Step 1: Quantize with compression for QLoRA (its own checkpoint: the others are not compressed) _run_quantize( "configs/train/qlora_nvfp4.yaml", - [ - "--model_name_or_path", tiny_qwen3_path, - "--recipe", "general/ptq/nvfp4_default-kv_fp8", - "--calib_size", "64", - "--compress", "True", - "--output_dir", str(ptq_output_dir), - ], + { + "model_name_or_path": tiny_qwen3_path, + "calib_size": "64", + "output_dir": str(ptq_output_dir), + }, cache_dir=dataset_cache_dir, ) @@ -286,19 +326,16 @@ def test_qwen3_qlora_nvfp4(tiny_qwen3_path, dataset_cache_dir, tmp_path): qlora_output_dir = tmp_path / "qlora" _run_train( "configs/train/qlora_nvfp4.yaml", - [ - "--model_name_or_path", str(ptq_output_dir), - "--do_train", "True", - "--lora", "True", - "--output_dir", str(qlora_output_dir), - ], + {"model_name_or_path": str(ptq_output_dir), "output_dir": str(qlora_output_dir)}, backend="ddp", cache_dir=dataset_cache_dir, ) + _assert_training_progress(qlora_output_dir) + # Step 3: Export the QLoRA checkpoint for deployment export_dir = tmp_path / "qlora_export" - _run_export(str(qlora_output_dir), str(export_dir)) + _run_export(str(qlora_output_dir), str(export_dir), config="configs/train/qlora_nvfp4.yaml") # The base model is exported compressed; the adapters stay at the top level. base_model_dir = export_dir / "base_model" @@ -310,7 +347,9 @@ def test_qwen3_qlora_nvfp4(tiny_qwen3_path, dataset_cache_dir, tmp_path): # NVFP4 needs the packed weight and *both* scales to be dequantizable downstream. base_weights = load_file(base_model_dir / "model.safetensors") - packed_weights = [k for k, v in base_weights.items() if k.endswith(".weight") and v.dtype == torch.uint8] + packed_weights = [ + k for k, v in base_weights.items() if k.endswith(".weight") and v.dtype == torch.uint8 + ] assert packed_weights, "no NVFP4-packed weights found in the exported base model" for key in packed_weights: prefix = key.removesuffix(".weight") diff --git a/tests/examples/llm_qat/test_llm_qat_readme.py b/tests/examples/llm_qat/test_llm_qat_readme.py new file mode 100644 index 00000000000..134b85f3b06 --- /dev/null +++ b/tests/examples/llm_qat/test_llm_qat_readme.py @@ -0,0 +1,120 @@ +# 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. + +"""Execute named README snippets through the shared documentation runner.""" + +import os +from pathlib import Path + +import pytest +from _test_utils.doc_tests.parser import Scenario +from _test_utils.doc_tests.pytest_utils import execute_scenario, scenario_parameters + +REPO = Path(__file__).resolve().parents[3] +README = REPO / "examples/llm_qat/README.md" + +SCENARIOS = [ + Scenario( + "llm-qat-install", + ["llm-qat-install"], + timeout_seconds=7200, + handler="_test_utils.doc_tests.fixtures.llm_qat_readme:llm_qat_install", + manual=True, + ), + Scenario("llm-qat-recipes", ["llm-qat-recipes"], timeout_seconds=30), + Scenario( + "llm-qat-simple", + ["llm-qat-simple"], + profile="gpu", + timeout_seconds=900, + handler="_test_utils.doc_tests.fixtures.llm_qat_readme:llm_qat_simple", + ), + Scenario( + "llm-qlora-serve", + ["llm-qlora-serve"], + profile="gpu", + timeout_seconds=900, + handler="_test_utils.doc_tests.fixtures.llm_qat_readme:llm_qlora_serve", + manual=True, + requires=["vllm"], + background=True, + ), + Scenario( + "llm-qat-python", + ["llm-qat-python"], + profile="gpu", + timeout_seconds=900, + handler="_test_utils.doc_tests.fixtures.llm_qat_readme:llm_qat_python", + ), + Scenario( + "llm-qat-trainer", + ["llm-qat-trainer"], + profile="gpu", + timeout_seconds=900, + handler="_test_utils.doc_tests.fixtures.llm_qat_readme:llm_qat_trainer", + ), + Scenario( + "llm-qad-trainer", + ["llm-qad-trainer"], + profile="gpu", + timeout_seconds=900, + handler="_test_utils.doc_tests.fixtures.llm_qat_readme:llm_qad_trainer", + ), + Scenario( + "llm-qat-llama-wrap", + ["llm-qat-llama-wrap"], + profile="gpu", + timeout_seconds=900, + handler="_test_utils.doc_tests.fixtures.llm_qat_readme:llm_qat_llama_wrap", + ), + Scenario( + "llm-qat-copy-config", + ["llm-qat-copy-config"], + timeout_seconds=900, + handler="_test_utils.doc_tests.fixtures.llm_qat_readme:llm_qat_copy_config", + ), + Scenario( + "llm-qat-cli-override", + ["llm-qat-cli-override"], + profile="gpu", + timeout_seconds=900, + handler="_test_utils.doc_tests.fixtures.llm_qat_readme:llm_qat_cli_override", + ), + Scenario( + "llm-qat-evaluation", + ["llm-qat-evaluation"], + profile="gpu", + timeout_seconds=7200, + handler="_test_utils.doc_tests.fixtures.llm_qat_readme:llm_qat_evaluation", + manual=True, + requires=["lm_eval"], + ), +] + +EXCLUDED = { + "llm-qat-dataset": "Executed by test_llm_qat.py::test_dataset_utils_pretokenize.", + "llm-qat-quickstart": "Executed by test_llm_qat.py::test_qwen3_qat_nvfp4 through run_example_command.", + "llm-qad-quickstart": "Executed by test_llm_qat.py::test_qwen3_qad_nvfp4 through run_example_command.", + "llm-qlora-quickstart": "Executed by test_llm_qat.py::test_qwen3_qlora_nvfp4 through run_example_command.", +} + + +@pytest.mark.parametrize("scenario", scenario_parameters(README, SCENARIOS, excluded=EXCLUDED)) +def test_readme(scenario, tmp_path, monkeypatch): + if scenario.id in {"llm-qat-python", "llm-qat-trainer", "llm-qad-trainer"}: + # Keep API smokes single-GPU: Trainer's threaded DataParallel races attention patching. + visible = os.environ.get("CUDA_VISIBLE_DEVICES", "0") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", visible.split(",", 1)[0]) + execute_scenario(scenario, REPO, tmp_path, monkeypatch) diff --git a/tests/examples/llm_qat/test_llm_qat_readme_cli.py b/tests/examples/llm_qat/test_llm_qat_readme_cli.py new file mode 100644 index 00000000000..1c3d3c47167 --- /dev/null +++ b/tests/examples/llm_qat/test_llm_qat_readme_cli.py @@ -0,0 +1,42 @@ +# 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. + +"""Check README Python flag names without executing example scenarios.""" + +from pathlib import Path + +import pytest +from _test_utils.doc_tests.cli import check_cli_options, python_commands + +REPO = Path(__file__).resolve().parents[3] +README = REPO / "examples/llm_qat/README.md" + + +@pytest.mark.timeout(90) +@pytest.mark.parametrize( + "command", + [ + pytest.param( + command, + id=f"{command.snippet}:{command.line}", + marks=[pytest.mark.manual] if command.snippet == "llm-qat-evaluation" else [], + ) + for command in python_commands( + README, cwd=README.parent, cwd_overrides={"llm-qat-quickstart": REPO} + ) + ], +) +def test_readme_cli_options(command): + check_cli_options(README, command) diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py index eae2e14df84..d9d87c71930 100644 --- a/tests/examples/megatron_bridge/test_distill.py +++ b/tests/examples/megatron_bridge/test_distill.py @@ -22,6 +22,7 @@ import pytest import torch +from _test_utils.examples.readme import readme_command from _test_utils.examples.run_command import extend_cmd_parts, run_example_command from _test_utils.torch.puzzletron.utils import create_and_save_small_hf_model from _test_utils.torch.transformers_models import ( @@ -62,32 +63,31 @@ def test_distill_llm_hf_export(tmp_path, num_gpus, fake_hub_checkpoint): distill_output_dir = tmp_path / "distill_output" distilled_hf_path = tmp_path / "distilled_hf" all_exports = tmp_path / "all_exports" - distill_cmd_parts = extend_cmd_parts( - [ - "torchrun", - f"--nproc_per_node={num_gpus}", - "distill.py", - "--use_mock_data", - "--document_boundaries", - ], - student_hf_path=teacher_hf_path, - teacher_hf_path=teacher_hf_path, - output_dir=distill_output_dir, - tp_size=num_gpus, - pp_size=1, - # Mock token ids cycle through the 128-token vocab, so the EOD id recurs every 128 tokens - # and each window trains as ~17 THD segments. - seq_length=2048, - mbs=1, - gbs=4, - train_iters=train_iters, - lr_warmup_iters=1, - eval_interval=1, # save_interval defaults to it, so every iteration is checkpointed - eval_iters=1, - log_interval=1, - checkpoint_keep_last=-1, - hf_export_path=distilled_hf_path, + distill_cmd_parts = readme_command( + Path(__file__).resolve().parents[3] / "examples/megatron_bridge/README.md", + "bridge-distill", + "distill.py", + overrides={ + "nproc_per_node": num_gpus, + "tp_size": num_gpus, + "pp_size": 1, + "teacher_hf_path": teacher_hf_path, + "student_hf_path": teacher_hf_path, + "output_dir": distill_output_dir, + # Preserve this test's THD/document-boundary coverage. + "seq_length": 2048, + "mbs": 1, + "gbs": 4, + "train_iters": train_iters, + "lr_warmup_iters": 1, + "eval_interval": 1, + "eval_iters": 1, + "log_interval": 1, + "checkpoint_keep_last": -1, + "hf_export_path": distilled_hf_path, + }, ) + distill_cmd_parts.append("--document_boundaries") run_example_command(distill_cmd_parts, example_path="megatron_bridge") checkpoint_root = distill_output_dir / "checkpoints" @@ -97,12 +97,16 @@ def test_distill_llm_hf_export(tmp_path, num_gpus, fake_hub_checkpoint): assert (distilled_hf_path / fake_hub_checkpoint.marker).exists() # export_distilled_megatron_to_hf.py converts any saved iteration, not just the final one. - all_export_cmd_parts = extend_cmd_parts( - ["torchrun", "--nproc_per_node=1", "export_distilled_megatron_to_hf.py"], - student_hf_path=student_export_arg, - megatron_path=checkpoint_root, - hf_export_path=all_exports, - export_iterations="all", + all_export_cmd_parts = readme_command( + Path(__file__).resolve().parents[3] / "examples/megatron_bridge/README.md", + "bridge-distill-export", + "export_distilled_megatron_to_hf.py", + overrides={ + "nproc_per_node": 1, + "student_hf_path": student_export_arg, + "megatron_path": checkpoint_root, + "hf_export_path": all_exports, + }, ) run_example_command(all_export_cmd_parts, example_path="megatron_bridge") diff --git a/tests/examples/megatron_bridge/test_megatron_bridge_readme.py b/tests/examples/megatron_bridge/test_megatron_bridge_readme.py new file mode 100644 index 00000000000..8a14955d1de --- /dev/null +++ b/tests/examples/megatron_bridge/test_megatron_bridge_readme.py @@ -0,0 +1,138 @@ +# 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. + +"""Execute named README snippets through the shared documentation runner.""" + +from pathlib import Path + +import pytest +from _test_utils.doc_tests.parser import Scenario +from _test_utils.doc_tests.pytest_utils import execute_scenario, scenario_parameters + +REPO = Path(__file__).resolve().parents[3] +README = REPO / "examples/megatron_bridge/README.md" + +SCENARIOS = [ + Scenario( + "bridge-login", + ["bridge-login"], + timeout_seconds=120, + handler="_test_utils.doc_tests.fixtures.megatron_bridge_readme:bridge_login", + manual=True, + requires=["huggingface_hub"], + ), + Scenario( + "bridge-import", + ["bridge-import"], + profile="gpu", + timeout_seconds=7200, + handler="_test_utils.doc_tests.fixtures.megatron_bridge_readme:bridge_import", + manual=True, + min_gpus=8, + requires=["megatron.bridge"], + ), + Scenario( + "bridge-mlflow", + ["bridge-mlflow"], + profile="gpu", + timeout_seconds=7200, + handler="_test_utils.doc_tests.fixtures.megatron_bridge_readme:bridge_mlflow", + manual=True, + min_gpus=2, + requires=["megatron.bridge", "mlflow"], + ), + Scenario( + "bridge-real-data", + ["bridge-real-data"], + profile="gpu", + timeout_seconds=86400, + handler="_test_utils.doc_tests.fixtures.megatron_bridge_readme:bridge_real_data", + manual=True, + min_gpus=8, + requires=["megatron.bridge"], + ), + Scenario( + "bridge-qad", + ["bridge-qad"], + profile="gpu", + timeout_seconds=86400, + handler="_test_utils.doc_tests.fixtures.megatron_bridge_readme:bridge_qad", + manual=True, + min_gpus=8, + requires=["megatron.bridge"], + ), + Scenario( + "bridge-prune-active", + ["bridge-prune-active"], + profile="gpu", + timeout_seconds=7200, + handler="_test_utils.doc_tests.fixtures.megatron_bridge_readme:bridge_prune_active", + manual=True, + min_gpus=2, + requires=["megatron.bridge"], + ), + Scenario( + "bridge-prune-memory", + ["bridge-prune-memory"], + profile="gpu", + timeout_seconds=7200, + handler="_test_utils.doc_tests.fixtures.megatron_bridge_readme:bridge_prune_memory", + manual=True, + min_gpus=2, + requires=["megatron.bridge"], + ), + Scenario( + "bridge-prune-vlm", + ["bridge-prune-vlm"], + profile="gpu", + timeout_seconds=7200, + handler="_test_utils.doc_tests.fixtures.megatron_bridge_readme:bridge_prune_vlm", + manual=True, + min_gpus=2, + requires=["megatron.bridge"], + ), + Scenario( + "bridge-generation", + ["bridge-generation"], + profile="gpu", + timeout_seconds=7200, + handler="_test_utils.doc_tests.fixtures.megatron_bridge_readme:bridge_generation", + manual=True, + min_gpus=1, + requires=["vllm"], + ), +] + +EXCLUDED = { + "bridge-quantize": "Executed by test_quantize_export.py::test_quantize_and_export[readme].", + "bridge-export": "Executed by test_quantize_export.py::test_quantize_and_export[readme].", + "bridge-distill": "Executed by test_distill.py::test_distill_llm_hf_export.", + "bridge-distill-export": "Executed by test_distill.py::test_distill_llm_hf_export.", + "bridge-prune-params": "Executed by test_prune_minitron.py::test_prune_minitron[qwen3].", + "bridge-prune-manual": "Executed by test_prune_minitron.py::test_prune_minitron[readme-manual].", + "bridge-distill-help": "Validated by the README CLI preflight.", + "bridge-prune-minitron-help": "Validated by the README CLI preflight.", + "bridge-lora-ptq": "Separate LoRA PTQ integration; full-model adapter calibration is outside this pilot.", + "bridge-docker": "Interactive container provisioning runs outside pytest; CI supplies the container.", + "bridge-vlm-template": "Incomplete template; concrete training/export snippets are tested.", + "bridge-inline-export-template": "Incomplete template; concrete training/export snippets are tested.", + "bridge-single-export-template": "User-selected model/iteration; multi-iteration conversion is tested.", + "bridge-output-tree": "Illustrative output tree, not executable code.", +} + + +@pytest.mark.parametrize("scenario", scenario_parameters(README, SCENARIOS, excluded=EXCLUDED)) +def test_readme(scenario, tmp_path, monkeypatch): + execute_scenario(scenario, REPO, tmp_path, monkeypatch) diff --git a/tests/examples/megatron_bridge/test_megatron_bridge_readme_cli.py b/tests/examples/megatron_bridge/test_megatron_bridge_readme_cli.py new file mode 100644 index 00000000000..5efc620871e --- /dev/null +++ b/tests/examples/megatron_bridge/test_megatron_bridge_readme_cli.py @@ -0,0 +1,40 @@ +# 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. + +"""Check README Python flag names without executing example scenarios.""" + +from pathlib import Path + +import pytest +from _test_utils.doc_tests.cli import check_cli_options, python_commands + +REPO = Path(__file__).resolve().parents[3] +README = REPO / "examples/megatron_bridge/README.md" + + +@pytest.mark.timeout(90) +@pytest.mark.parametrize( + "command", + [ + pytest.param( + command, + id=f"{command.snippet}:{command.line}", + marks=[pytest.mark.manual] if command.snippet == "bridge-generation" else [], + ) + for command in python_commands(README, cwd=README.parent) + ], +) +def test_readme_cli_options(command): + check_cli_options(README, command) diff --git a/tests/examples/megatron_bridge/test_prune_minitron.py b/tests/examples/megatron_bridge/test_prune_minitron.py index 2cdee235d06..3177d19fdbc 100644 --- a/tests/examples/megatron_bridge/test_prune_minitron.py +++ b/tests/examples/megatron_bridge/test_prune_minitron.py @@ -14,8 +14,12 @@ # limitations under the License. """Tests for prune_minitron.py and distill.py scripts.""" +import json +from pathlib import Path + import pytest import torch +from _test_utils.examples.readme import readme_command from _test_utils.examples.run_command import extend_cmd_parts, run_example_command from _test_utils.torch.transformers_models import ( create_tiny_deepseek_v3_dir, @@ -28,7 +32,7 @@ @pytest.mark.parametrize( - ("create_teacher", "expected_pruned_config"), + ("create_teacher", "expected_pruned_config", "readme_snippet"), [ # Dense Qwen3 LM. pytest.param( @@ -36,8 +40,21 @@ tmp_path, with_tokenizer=True, return_model=True, num_hidden_layers=num_gpus ), {}, + "bridge-prune-params", id="qwen3", ), + pytest.param( + lambda tmp_path, num_gpus: create_tiny_qwen3_dir( + tmp_path, + with_tokenizer=True, + return_model=True, + num_hidden_layers=num_gpus, + intermediate_size=128, + ), + {"intermediate_size": 64}, + "bridge-prune-manual", + id="readme-manual", + ), # NemotronH (nemotron-3-nano): Mamba + attention + MoE hybrid. # MTP heads are enabled so the run covers dropping them during calibration. pytest.param( @@ -49,6 +66,7 @@ mtp_hybrid_override_pattern="*E", ), {"n_shared_experts": 1}, + None, id="nemotron_h", ), # DeepSeek-V3: MLA (Q-LoRA) + MoE sizing the shared expert as @@ -66,12 +84,13 @@ topk_group=1, ), {"n_shared_experts": 1}, + None, id="deepseek_v3", ), ], ) def test_prune_minitron( - tmp_path, num_gpus, create_teacher, expected_pruned_config, fake_hub_checkpoint + tmp_path, num_gpus, create_teacher, expected_pruned_config, readme_snippet, fake_hub_checkpoint ): """The teacher is passed as a Hub ID (see ``fake_hub_checkpoint``): it is pruned from its local copy, and the HF output carries that copy's non-model files.""" @@ -81,28 +100,45 @@ def test_prune_minitron( prune_target_params = int(teacher_params * 0.8) pruned_path = tmp_path / "pruned" - prune_command_parts = extend_cmd_parts( - ["torchrun", f"--nproc_per_node={num_gpus}", "prune_minitron.py"], - hf_model_name_or_path=teacher_arg, - output_hf_path=pruned_path, - pp_size=num_gpus, - calib_dataset_name="cnn_dailymail", - calib_num_samples=8, - seq_length=16, - prune_target_params=prune_target_params, - prune_score_func="mmlu_1pct_bs32", - score_lower_bound=0.0, # exercise the score-gate path without coupling to the model's acc - ss_channel_divisor=4, - hparams_to_skip="num_attention_heads", - top_k=1, - ) + options = { + "hf_model_name_or_path": teacher_arg, + "output_hf_path": pruned_path, + "pp_size": num_gpus, + "calib_dataset_name": "cnn_dailymail", + "calib_num_samples": 8, + "seq_length": 16, + } + if readme_snippet == "bridge-prune-manual": + options["prune_export_config"] = json.dumps({"ffn_hidden_size": 64}) + else: + options.update( + prune_target_params=prune_target_params, + prune_score_func="mmlu_1pct_bs32", + score_lower_bound=0.0, # exercise the score gate without coupling to model accuracy + ss_channel_divisor=4, + top_k=1, + ) + if readme_snippet is not None: + prune_command_parts = readme_command( + Path(__file__).resolve().parents[3] / "examples/megatron_bridge/README.md", + readme_snippet, + "prune_minitron.py", + overrides={"nproc_per_node": num_gpus, **options}, + ) + else: + prune_command_parts = extend_cmd_parts( + ["torchrun", f"--nproc_per_node={num_gpus}", "prune_minitron.py"], + **options, + hparams_to_skip="num_attention_heads", + ) run_example_command(prune_command_parts, example_path="megatron_bridge") assert (pruned_path / "config.json").exists() if fake_hub_checkpoint.active: assert (pruned_path / fake_hub_checkpoint.marker).exists() pruned_model = AutoModelForCausalLM.from_pretrained(pruned_path) - assert sum(p.numel() for p in pruned_model.parameters()) <= prune_target_params + if readme_snippet != "bridge-prune-manual": + assert sum(p.numel() for p in pruned_model.parameters()) <= prune_target_params for field, expected in expected_pruned_config.items(): assert getattr(pruned_model.config, field) == expected diff --git a/tests/examples/megatron_bridge/test_quantize_export.py b/tests/examples/megatron_bridge/test_quantize_export.py index d031990e76a..c9fc6f5ae7b 100644 --- a/tests/examples/megatron_bridge/test_quantize_export.py +++ b/tests/examples/megatron_bridge/test_quantize_export.py @@ -19,6 +19,7 @@ import pytest import torch import yaml +from _test_utils.examples.readme import readme_command from _test_utils.examples.run_command import extend_cmd_parts, run_example_command from _test_utils.torch.export.unified_checkpoint import ( assert_exported_checkpoint_matches, @@ -27,6 +28,7 @@ from _test_utils.torch.megatron.modelopt_state import assert_has_modelopt_state from _test_utils.torch.transformers_models import ( create_tiny_nemotron_h_dir, + create_tiny_qwen3_dir, create_tiny_qwen3_moe_dir, create_tiny_qwen3vl_dir, ) @@ -45,22 +47,46 @@ } +def _command(script, num_gpus, options, *, snippet=None): + if snippet is not None: + return readme_command( + Path(__file__).resolve().parents[3] / "examples/megatron_bridge/README.md", + snippet, + script, + overrides={"nproc_per_node": num_gpus, **options}, + ) + return extend_cmd_parts(["torchrun", f"--nproc_per_node={num_gpus}", script], **options) + + @pytest.mark.parametrize( ( "create_model", "model_kwargs", - "recipe", + "quantization_args", + "use_readme", "quantize_parallelism", "export_parallelism", "fakequant_export_parallelism", ), [ + # The former README smoke reuses this test's checkpoint/export assertions. + pytest.param( + create_tiny_qwen3_dir, + dict(_DENSE_KWARGS, head_dim=64), + {}, + True, + "tp_size", + "pp_size", + None, + id="readme", + ), # MoE: routed experts used to be dropped silently from the export. Calibrates with context # parallelism (not on the VLM: M-RoPE CP needs nemo:26.10). ( create_tiny_qwen3_moe_dir, _DENSE_KWARGS, - "general/ptq/nvfp4_default-kv_fp8", + {"recipe": "general/ptq/nvfp4_default-kv_fp8"}, + False, "cp_size", "pp_size", # Per-expert HF names, each EP rank owning some of them. @@ -71,7 +97,8 @@ ( create_tiny_qwen3vl_dir, {}, - "general/ptq/fp8_default-kv_fp8", + {"recipe": "general/ptq/fp8_default-kv_fp8"}, + False, "tp_size", "pp_size", "tp_size", @@ -81,13 +108,14 @@ ( create_tiny_nemotron_h_dir, {}, - "general/ptq/nvfp4_default-kv_fp8", + {"recipe": "general/ptq/nvfp4_default-kv_fp8"}, + False, "tp_size", "ep_size", "pp_size", ), ], - ids=["qwen3_moe", "qwen3vl", "nemotron_h"], + ids=["readme", "qwen3_moe", "qwen3vl", "nemotron_h"], ) @pytest.mark.timeout(420) # quantize + 2 exports in one test; 1-gpu CI exceeds the default 300s def test_quantize_and_export( @@ -95,13 +123,14 @@ def test_quantize_and_export( num_gpus, create_model, model_kwargs, - recipe, + quantization_args, + use_readme, quantize_parallelism, export_parallelism, fakequant_export_parallelism, fake_hub_checkpoint, ): - """Quantize a tiny model via a YAML recipe and export it to a unified HF checkpoint. + """Quantize a tiny model and export it to a unified HF checkpoint. The model is passed as a Hub ID (see ``fake_hub_checkpoint``): both steps load and export from its local copy, and the Megatron checkpoint records the Hub ID as its tokenizer. @@ -114,10 +143,9 @@ def test_quantize_and_export( hf_export_path = tmp_path / "quantized_hf" # Step 1: quantize and save a Megatron checkpoint - quantize_cmd = extend_cmd_parts( - ["torchrun", f"--nproc_per_node={num_gpus}", "quantize.py", "--skip_generate"], + quantize_args = dict( hf_model_name_or_path=model_arg, - recipe=recipe, + **quantization_args, **{quantize_parallelism: num_gpus}, calib_dataset_name="cnn_dailymail", calib_num_samples=4, @@ -125,6 +153,13 @@ def test_quantize_and_export( seq_length=16, export_megatron_path=megatron_path, ) + quantize_cmd = _command( + "quantize.py", + num_gpus, + quantize_args, + snippet="bridge-quantize" if use_readme else None, + ) + quantize_cmd.append("--skip_generate") run_example_command(quantize_cmd, example_path="megatron_bridge", setup_free_port=True) assert (megatron_path / "latest_checkpointed_iteration.txt").exists() assert_has_modelopt_state(megatron_path) @@ -135,13 +170,18 @@ def test_quantize_and_export( ) # Step 2: export to HF - export_cmd = extend_cmd_parts( - ["torchrun", f"--nproc_per_node={num_gpus}", "export_quantized_megatron_to_hf.py"], + export_args = dict( hf_model_name_or_path=model_arg, megatron_path=megatron_path, export_unified_hf_path=hf_export_path, **{export_parallelism: num_gpus}, ) + export_cmd = _command( + "export_quantized_megatron_to_hf.py", + num_gpus, + export_args, + snippet="bridge-export" if use_readme else None, + ) run_example_command(export_cmd, example_path="megatron_bridge", setup_free_port=True) assert (hf_export_path / "config.json").exists() assert (hf_export_path / "hf_quant_config.json").exists() @@ -164,6 +204,10 @@ def test_quantize_and_export( # outputs = llm.generate(["Hello!"], vllm.SamplingParams(max_tokens=4)) # assert outputs and outputs[0].outputs and outputs[0].outputs[0].text + if fakequant_export_parallelism is None: + # This case ends at the unified export; the other cases also check fake-quant sidecars. + return + # Step 3: export folded weights + quantizer sidecars through Megatron-Bridge for vLLM fake quant fakequant_path = tmp_path / "vllm_fakequant" fakequant_cmd = extend_cmd_parts( diff --git a/tests/unit/doc_tests/test_cli.py b/tests/unit/doc_tests/test_cli.py new file mode 100644 index 00000000000..4bdae991a2c --- /dev/null +++ b/tests/unit/doc_tests/test_cli.py @@ -0,0 +1,87 @@ +# 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. + +"""Check real script help without executing README commands or launchers.""" + +import pytest +from _test_utils.doc_tests.cli import check_cli_options, python_commands +from _test_utils.doc_tests.parser import DocTestError + + +@pytest.mark.parametrize( + "launcher", + [ + "python", + "python3", + "torchrun --nproc_per_node 8", + "accelerate launch --config_file unused.yaml", + ], +) +def test_readme_flags_against_cached_script_help(tmp_path, launcher): + nested = tmp_path / "nested" + nested.mkdir() + script = nested / "cli.py" + script.write_text("""import argparse +import os +from pathlib import Path +p = argparse.ArgumentParser(description="An obsolete --removed flag is mentioned here.") +p.add_argument('--alpha', '-a', required=True) +p.add_argument('--beta') +with Path('help-calls').open('a') as calls: + calls.write('help\\n') +assert os.environ['CUDA_VISIBLE_DEVICES'] == '' +p.parse_args() +raise AssertionError('Only --help should execute') +""") + readme = tmp_path / "README.md" + readme.write_text(f""" +```bash +cd nested +{launcher} cli.py --alpha "$(touch forbidden)" \\ + --beta=value +python3 cli.py -a two +``` +""") + commands = python_commands(readme, cwd=tmp_path) + assert len(commands) == 2 + assert commands[0].line == 4 + for command in commands: + check_cli_options(readme, command) + assert (nested / "help-calls").read_text() == "help\n" + assert not (nested / "forbidden").exists() + readme.write_text(readme.read_text().replace("--alpha", "--removed")) + with pytest.raises( + DocTestError, match=r"README.md:4: cli: cli.py: unknown options:.*--removed" + ): + check_cli_options(readme, python_commands(readme, cwd=tmp_path)[0]) + + +@pytest.mark.parametrize( + ("body", "message"), + [ + (None, "does not exist"), + ("raise SystemExit('help unavailable')", "help unavailable"), + ("import time; time.sleep(10)", "timed out"), + ("print('No option definitions')", "did not list any options"), + ], +) +def test_unusable_script_help_reports_readme_location(tmp_path, body, message): + if body is not None: + (tmp_path / "cli.py").write_text(body) + readme = tmp_path / "README.md" + readme.write_text("\n```sh\npython cli.py --alpha=one\n```\n") + command = python_commands(readme, cwd=tmp_path)[0] + with pytest.raises(DocTestError, match=rf"(?s)README.md:3: cli:.*{message}"): + check_cli_options(readme, command, timeout=0.2) diff --git a/tests/unit/doc_tests/test_markdown.py b/tests/unit/doc_tests/test_markdown.py new file mode 100644 index 00000000000..c85e2f9e806 --- /dev/null +++ b/tests/unit/doc_tests/test_markdown.py @@ -0,0 +1,86 @@ +# 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. + +"""Focused checks for README parsing, failure locations, and subprocess cleanup.""" + +import os +import time +from pathlib import Path + +import pytest +from _test_utils.doc_tests.parser import DocTestError, Scenario, parse_markdown, resolve_scenarios +from _test_utils.doc_tests.runner import run_scenario + + +def _scenario(tmp_path, code, *, language="sh", timeout_seconds=60): + path = tmp_path / "README.md" + path.write_text(f"\n```{language}\n{code}\n```\n") + return resolve_scenarios( + path, [Scenario("smoke", ["command"], timeout_seconds=timeout_seconds)] + )[0] + + +def test_named_fence_parsing_and_selection(tmp_path): + path = tmp_path / "README.md" + path.write_text(""" +~~~bash +printf 'hello' +~~~ +> +> ```python +> answer = 42 +> ``` +""") + snippets = parse_markdown(path) + assert [(s.id, s.kind, s.line, s.code) for s in snippets.values()] == [ + ("shell", "run", 3, "printf 'hello'\n"), + ("python", "python", 7, "answer = 42\n"), + ] + (scenario,) = resolve_scenarios( + path, [Scenario("smoke", ["python"])], excluded={"shell": "Illustrative command"} + ) + assert scenario.path == path + assert scenario.steps == [snippets["python"]] + + +@pytest.mark.skipif(os.name != "posix", reason="Requires bash/POSIX") +@pytest.mark.parametrize( + ("language", "code"), + [("sh", "false | cat\ntouch forbidden"), ("python", "raise RuntimeError('snippet failed')")], +) +def test_failure_reports_snippet_and_markdown_line(tmp_path, language, code): + scenario = _scenario(tmp_path, code, language=language) + with pytest.raises(DocTestError) as error: + run_scenario(scenario, tmp_path, tmp_path) + assert f"{scenario.path}:3:" in str(error.value) + assert "command" in str(error.value) + assert not (tmp_path / "forbidden").exists() + + +@pytest.mark.skipif(os.name != "posix", reason="Requires bash/POSIX") +def test_timeout_cleans_up_descendants(tmp_path): + if not Path("/proc/self/stat").exists(): + pytest.skip("Requires /proc to verify descendant exit") + scenario = _scenario(tmp_path, "sleep 30 &\necho $! > child.pid\nsleep 30", timeout_seconds=1) + with pytest.raises(DocTestError, match="timed out"): + run_scenario(scenario, tmp_path, tmp_path) + assert not list(tmp_path.glob("doc-test-*")) + pid = int((tmp_path / "child.pid").read_text()) + for _ in range(100): + status = Path(f"/proc/{pid}/stat") + if not status.exists() or status.read_text().split()[2] == "Z": + return + time.sleep(0.01) + pytest.fail(f"Child {pid} survived scenario cleanup") diff --git a/tests/unit/examples/test_readme.py b/tests/unit/examples/test_readme.py new file mode 100644 index 00000000000..e156276d3bb --- /dev/null +++ b/tests/unit/examples/test_readme.py @@ -0,0 +1,68 @@ +# 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. + +"""Check README command extraction and explicit end-to-end input overrides.""" + +import pytest +from _test_utils.doc_tests.parser import DocTestError +from _test_utils.examples.readme import readme_command + + +def test_readme_command_preserves_flags_and_substitutes_inputs(tmp_path): + path = tmp_path / "README.md" + path.write_text(""" +```bash +torchrun --nproc_per_node 8 train.py --model_path pretrained \ + --max_steps 100 --seq_length "${SEQ_LENGTH:-512}" --new-flag +python export.py --checkpoint trained +``` +""") + command = readme_command( + path, + "quickstart", + "train.py", + overrides={ + "nproc_per_node": 1, + "model_path": "local model", + "max_steps": 2, + "save_steps": 1, + }, + ) + assert command == [ + "torchrun", + "--nproc_per_node", + "1", + "train.py", + "--model_path", + "local model", + "--max_steps", + "2", + "--seq_length", + "512", + "--new-flag", + "--save_steps", + "1", + ] + assert readme_command(path, "quickstart", "export.py") == [ + "python", + "export.py", + "--checkpoint", + "trained", + ] + with pytest.raises(DocTestError): + readme_command(path, "renamed", "train.py") + path.write_text(path.read_text().replace("${SEQ_LENGTH:-512}", "$UNSET_LENGTH")) + with pytest.raises(DocTestError): + readme_command(path, "quickstart", "train.py") diff --git a/tests/unit/test_example_tests_workflow.py b/tests/unit/test_example_tests_workflow.py new file mode 100644 index 00000000000..20020e96962 --- /dev/null +++ b/tests/unit/test_example_tests_workflow.py @@ -0,0 +1,73 @@ +# 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. + +"""Exercise example CI routing without allocating GPU runners.""" + +import os +import re +import subprocess +from pathlib import Path + +import pytest +import yaml + +REPO = Path(__file__).resolve().parents[2] +ALL_LANES = {"torch", "vllm", "trtllm", "megatron", "diffusers", "onnx"} +QAT = "examples/llm_qat/README.md" +BRIDGE = "examples/megatron_bridge/README.md" + + +@pytest.mark.skipif(os.name != "posix", reason="Workflow uses Bash") +@pytest.mark.parametrize( + ("paths", "common", "stacked", "non_pr", "expected"), + [ + ([QAT], False, False, False, {"torch"}), + ([BRIDGE], False, False, False, {"megatron"}), + ([QAT, BRIDGE], False, False, False, {"torch", "megatron"}), + (["examples/hf_ptq/README.md"], False, False, False, set()), + ([], True, False, False, ALL_LANES), + ([QAT, BRIDGE], False, True, False, set()), + ([], False, False, True, ALL_LANES), + ], +) +def test_example_readme_lane_routing(tmp_path, *, paths, common, stacked, non_pr, expected): + workflow = yaml.safe_load((REPO / ".github/workflows/example_tests.yml").read_text()) + steps = workflow["jobs"]["pr-gate"]["steps"] + groups = yaml.safe_load( + next(step["with"]["files_yaml"] for step in steps if step.get("id") == "changed") + ) + flags = { + f"{group}_any_modified": any(path in groups[group] for path in paths) + for group in ("torch_docs", "megatron_docs") + } + flags["common_any_modified"] = common + lanes = next(step for step in steps if step.get("id") == "lanes") + env = dict(os.environ, GITHUB_OUTPUT=str(tmp_path / "outputs")) + for key, expression in lanes["env"].items(): + if key in {"NON_PR", "STACKED"}: + enabled = non_pr if key == "NON_PR" else stacked + else: + terms = expression.removeprefix("${{ ").removesuffix(" }}").split(" || ") + enabled = False + for term in terms: + match = re.fullmatch(r"steps\.changed\.outputs\.(\w+)(?: == 'true')?", term) + assert match, f"Unsupported lane expression: {expression}" + enabled |= flags.get(match[1], False) + env[key] = str(enabled).lower() + # Run the workflow's actual gate; no model subprocesses or GPU jobs are started. + subprocess.run(["bash", "-e", "-c", lanes["run"]], env=env, check=True, timeout=5) + outputs = dict(line.split("=", 1) for line in (tmp_path / "outputs").read_text().splitlines()) + assert {lane for lane in ALL_LANES if outputs[lane] == "true"} == expected + assert outputs["any_changed"] == str(bool(expected)).lower() diff --git a/tests/unit/torch/opt/plugins/test_modelopt_arg_parser.py b/tests/unit/torch/opt/plugins/test_modelopt_arg_parser.py index 20f789dc078..79de8625e22 100644 --- a/tests/unit/torch/opt/plugins/test_modelopt_arg_parser.py +++ b/tests/unit/torch/opt/plugins/test_modelopt_arg_parser.py @@ -70,6 +70,15 @@ def test_yaml_warmup_spelling_follows_transformers(tmp_path, arg_type, yaml, exp class TestModelOptArgParser: """Tests for ModelOptArgParser --config and --generate_docs features.""" + def test_help_includes_wrapper_options(self, capsys): + parser = ModelOptArgParser((_ModelArgs, _TrainArgs)) + with pytest.raises(SystemExit) as exc_info: + parser.parse_args_into_dataclasses(args=["--help"]) + assert exc_info.value.code == 0 + help_text = capsys.readouterr().out + assert "--config PATH" in help_text + assert "--generate_docs [PATH]" in help_text + def test_cli_args_only(self): parser = ModelOptArgParser((_ModelArgs, _TrainArgs)) model_args, train_args = parser.parse_args_into_dataclasses(