From aa572bca20d5153ba80d77fc27208fd045354293 Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Tue, 22 Sep 2026 20:24:31 -0700 Subject: [PATCH 1/4] Add Alpamayo FP8 Edge export guide --- .../user_guide/alpamayo_fp8_edge_exporter.rst | 598 ++++++++++++++++++ docsrc/user_guide/edge_exporter.rst | 11 +- docsrc/user_guide/index.rst | 1 + 3 files changed, 608 insertions(+), 2 deletions(-) create mode 100644 docsrc/user_guide/alpamayo_fp8_edge_exporter.rst diff --git a/docsrc/user_guide/alpamayo_fp8_edge_exporter.rst b/docsrc/user_guide/alpamayo_fp8_edge_exporter.rst new file mode 100644 index 0000000000..98cc3b46ad --- /dev/null +++ b/docsrc/user_guide/alpamayo_fp8_edge_exporter.rst @@ -0,0 +1,598 @@ +.. _alpamayo_fp8_edge_exporter: + +Alpamayo 1.5 ModelOpt FP8 Edge Export +===================================== + +This guide walks through the complete workflow for: + +1. quantizing ``nvidia/Alpamayo-1.5-10B`` to FP8 with NVIDIA ModelOpt; +2. checking the quantized checkpoint with Alpamayo evaluation; +3. building the TensorRT-Edge-LLM plugin; +4. exporting Alpamayo's vision, language, and action components to TensorRT; +5. checking eager-versus-TensorRT parity and the generated engine artifacts. + +The Edge exporter produces three engines and a small ``torch.export`` program +that chains them: + +.. code-block:: text + + camera patches + │ + ▼ + vision.engine ──► visual token insertion + DeepStack packing + │ + ▼ + language.engine ──► prefix K/V cache + │ + ▼ + action.engine ──► one diffusion velocity step [batch, 64, 2] + +The export smoke test compiles and checks one action-denoising step. Alpamayo's +autoregressive reasoning loop and Python-controlled diffusion sampler remain +host-side orchestration; they are not represented as one static graph. + +.. warning:: + + Alpamayo and the PhysicalAI dataset are gated. Request access before + starting, and use ``hf auth login``. Do not put Hugging Face tokens directly + in shell commands, documentation, or logs. + +Validated component requirements +-------------------------------- + +Use Python 3.12. The quantization recipe and Edge exporter intentionally use +different environments because their dependency requirements differ. + +.. list-table:: + :header-rows: 1 + :widths: 24 36 40 + + * - Component + - Quantization environment + - Edge export environment + * - PyTorch + - ``2.8.0`` + - Version required by the installed Torch-TensorRT wheel + * - Transformers + - ``4.57.1`` + - ``>=5.4.0`` + * - ModelOpt + - ``0.43.0`` + - ``>=0.44.0`` + * - TensorRT + - Not used to create the checkpoint + - Must match Torch-TensorRT and the Edge-LLM plugin + * - CUDA + - CUDA 12 or 13, matching PyTorch + - Must match the plugin build + +The ModelOpt checkpoint format carries ``modelopt_state.pth`` so the newer +export environment can reconstruct quantized modules when +``mto.enable_huggingface_checkpointing()`` is enabled before +``from_pretrained``. + +Suggested source layout +----------------------- + +The commands below use this layout: + +.. code-block:: text + + /workspace/ + ├── alpamayo-recipes/ + │ └── recipes/alpamayo1_5_quant/ + ├── TensorRT/ + ├── TensorRT-Edge-LLM/ + └── TensorRT-Torch/ + +Set paths once: + +.. code-block:: bash + + export WORKSPACE=/workspace + export ALPAMAYO_RECIPES="$WORKSPACE/alpamayo-recipes" + export TORCH_TRT_ROOT="$WORKSPACE/TensorRT-Torch" + export EDGE_LLM_ROOT="$WORKSPACE/TensorRT-Edge-LLM" + export QUANT_DIR="$ALPAMAYO_RECIPES/recipes/alpamayo1_5_quant" + export QUANT_OUTPUT="$QUANT_DIR/outputs" + export ENGINE_DIR="$WORKSPACE/alpamayo_fp8_edge" + export HF_HOME="$WORKSPACE/.cache/huggingface" + +Use storage with enough space for the base checkpoint, calibration data, +quantized checkpoint, and TensorRT build artifacts. The base model is +approximately 22 GB in BF16; the FP8 parameter payload is approximately 11 GB. +TensorRT compilation also needs temporary host and GPU memory. + +0. Obtain the sources and gated assets +-------------------------------------- + +Clone the required repositories: + +.. code-block:: bash + + git clone https://github.com/NVlabs/alpamayo.git "$ALPAMAYO_RECIPES" + git clone https://github.com/pytorch/TensorRT.git "$TORCH_TRT_ROOT" + git clone https://github.com/NVIDIA/TensorRT-Edge-LLM.git "$EDGE_LLM_ROOT" + +Check out the Torch-TensorRT revision containing Alpamayo Edge exporter support +when it has not yet landed on your default branch. + +Request access to: + +* `Alpamayo-1.5-10B `_ +* `PhysicalAI-Autonomous-Vehicles + `_ + +Authenticate interactively: + +.. code-block:: bash + + hf auth login + hf auth whoami + +Configure the PhysicalAI dataset paths required by the Alpamayo recipe: + +.. code-block:: bash + + export ALPAMAYO_WORKSPACE="$ALPAMAYO_RECIPES" + export ALPAMAYO_MODEL_DIR="$WORKSPACE/alpamayo_model_converted_from_hf" + export ALPAMAYO_PAI_LOCAL_DIR="$WORKSPACE/PAI_mini" + export ALPAMAYO_LOG_DIR="$WORKSPACE/alpamayo_logs" + +Follow the Alpamayo repository's dataset setup instructions to populate +``ALPAMAYO_PAI_LOCAL_DIR`` and obtain the calibration parquet file referenced by +the recipe. + +1. Create the ModelOpt quantization environment +----------------------------------------------- + +Create the recipe environment with ``uv``: + +.. code-block:: bash + + export UV_CACHE_DIR="$WORKSPACE/.cache/uv" + cd "$QUANT_DIR" + + uv venv am15_quant + source am15_quant/bin/activate + + # Install torch first, then build flash-attn against it. + uv sync --active --no-install-package flash-attn + MAX_JOBS=4 uv sync --active + +Verify the important versions: + +.. code-block:: bash + + python - <<'PY' + import modelopt + import torch + import transformers + + print("torch", torch.__version__) + print("transformers", transformers.__version__) + print("modelopt", modelopt.__version__) + print("cuda available", torch.cuda.is_available()) + PY + +The recipe is defined for Python 3.12, PyTorch 2.8, Transformers 4.57.1, and +ModelOpt 0.43. + +2. Quantize Alpamayo with ModelOpt FP8 +------------------------------------- + +The exporter consumes a Hugging Face checkpoint restored through ModelOpt's +checkpoint integration. For the first TensorRT integration run, use +``--fake_quant``. It preserves FP16 checkpoint weights while retaining the FP8 +Q/DQ graph that TensorRT lowers. This is the lowest-risk export format. + +.. code-block:: bash + + cd "$QUANT_DIR" + source am15_quant/bin/activate + + uv run --active quantize.py \ + --ckpt nvidia/Alpamayo-1.5-10B \ + --quant_format fp8 \ + --num_of_calib_clips 100 \ + --save_model_dir "$QUANT_OUTPUT" \ + --fake_quant + +The expected output directory is: + +.. code-block:: text + + outputs/alpamayo1.5_fp8_calib100_fakequant/ + ├── config.json + ├── modelopt_state.pth + ├── model*.safetensors + └── tokenizer and processor assets + +Confirm that ModelOpt state was saved: + +.. code-block:: bash + + export ALPAMAYO_FP8_CKPT="$QUANT_OUTPUT/alpamayo1.5_fp8_calib100_fakequant" + + test -f "$ALPAMAYO_FP8_CKPT/config.json" + test -f "$ALPAMAYO_FP8_CKPT/modelopt_state.pth" + ls -lh "$ALPAMAYO_FP8_CKPT" + +Real compressed FP8 checkpoint +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Omit ``--fake_quant`` to call ``mtq.compress`` and save a real FP8 parameter +payload: + +.. code-block:: bash + + uv run --active quantize.py \ + --quant_format fp8 \ + --num_of_calib_clips 100 \ + --save_model_dir "$QUANT_OUTPUT" + +This produces ``alpamayo1.5_fp8_calib100`` and reduces checkpoint and runtime +weight memory. Real-quant restore and downstream compilation are still +experimental in ModelOpt. Establish the fake-quant export first, then validate +the compressed checkpoint separately. + +3. Evaluate the quantized checkpoint +------------------------------------ + +Evaluate a small subset before compiling TensorRT engines: + +.. code-block:: bash + + cd "$QUANT_DIR" + source am15_quant/bin/activate + + uv run --active eval.py \ + --ckpt "$ALPAMAYO_FP8_CKPT" \ + --limit 10 \ + --num_traj_samples 6 \ + --seed 42 \ + --print_every 1 + +The evaluation should: + +* restore ``modelopt_state.pth`` without missing-module errors; +* print the quantization summary for a non-base checkpoint; +* report finite per-clip minADE values; +* finish with average minADE and average evaluation time. + +Increase ``--limit`` after the smoke test to obtain the customer acceptance +metric. + +4. Prepare the Edge export environment +-------------------------------------- + +Deactivate the quantization environment and activate or create an environment +containing the Torch-TensorRT build used for export: + +.. code-block:: bash + + deactivate 2>/dev/null || true + source /path/to/edge-export/bin/activate + +Install the model and exporter dependencies: + +.. code-block:: bash + + python -m pip install \ + "transformers>=5.4.0" \ + "nvidia-modelopt[hf]>=0.44.0" \ + physical-ai-av \ + einops hydra-core pillow + + python -m pip install \ + "git+https://github.com/NVlabs/alpamayo1.5.git" + +Install a Torch-TensorRT wheel built for this environment, or build/install +Torch-TensorRT from the selected source revision. Confirm that importing +``torch_tensorrt`` resolves to the installed package: + +.. code-block:: bash + + cd "$TORCH_TRT_ROOT" + + # tools/hf contains the experimental exporter. Do not add $TORCH_TRT_ROOT/py + # unless the source tree has been built and generated _version.py. + export PYTHONPATH="$TORCH_TRT_ROOT/tools/hf" + + python - <<'PY' + import modelopt + import torch + import torch_tensorrt + import transformers + + print("torch", torch.__version__) + print("torch_tensorrt", torch_tensorrt.__version__) + print("torch_tensorrt path", torch_tensorrt.__file__) + print("transformers", transformers.__version__) + print("modelopt", modelopt.__version__) + PY + +``torch_tensorrt.__file__`` should point to ``site-packages`` unless you built +the source tree in place. + +5. Build the TensorRT-Edge-LLM plugin +------------------------------------ + +The language engine contains TensorRT-Edge-LLM attention plugins. Build the +plugin against the same CUDA and TensorRT major versions used by +Torch-TensorRT. + +Initialize dependencies: + +.. code-block:: bash + + cd "$EDGE_LLM_ROOT" + git submodule update --init --recursive + +Configure an x86 CUDA 13 / TensorRT 11 build: + +.. code-block:: bash + + export CUDA_HOME=/usr/local/cuda-13.0 + export PATH="$CUDA_HOME/bin:$PATH" + + cmake -S . -B build-alpamayo-export \ + -DCMAKE_BUILD_TYPE=Release \ + -DTRT_PACKAGE_DIR=/usr \ + -DCUDA_CTK_VERSION=13.0 \ + -DENABLE_CUTE_DSL=OFF \ + -DBUILD_UNIT_TESTS=OFF \ + -DBUILD_PYTHON_BINDINGS=OFF + + cmake --build build-alpamayo-export \ + --target NvInfer_edgellm_plugin \ + -j"$(nproc)" + +If TensorRT headers and libraries are installed under a separate SDK root, use +that path for ``TRT_PACKAGE_DIR``. It must contain ``include/NvInfer.h`` and a +``lib`` or ``lib64`` directory containing ``libnvinfer.so``. + +Set and verify the plugin: + +.. code-block:: bash + + export EDGE_LLM_PLUGIN_SO="$EDGE_LLM_ROOT/build-alpamayo-export/libNvInfer_edgellm_plugin.so.1.0" + export LD_LIBRARY_PATH="$(dirname "$EDGE_LLM_PLUGIN_SO"):${LD_LIBRARY_PATH:-}" + + test -f "$EDGE_LLM_PLUGIN_SO" + + python - <<'PY' + import ctypes + import os + + path = os.environ["EDGE_LLM_PLUGIN_SO"] + ctypes.CDLL(path) + print("loaded", path) + PY + +Use the CUDA and TensorRT versions available on your target system rather than +copying the example values blindly. + +6. Choose an export sample +-------------------------- + +The exporter uses one PhysicalAI clip to construct realistic image, language, +trajectory-history, mRoPE, DeepStack, and action inputs. Select a clip from the +calibration parquet: + +.. code-block:: bash + + cd "$QUANT_DIR" + source am15_quant/bin/activate + + export ALPAMAYO_CLIP_ID="$( + python - <<'PY' + from alpamayo1_5_quant.utils import read_clip_ids_from_parquet + + clips = read_clip_ids_from_parquet( + "0417_5k_train_set_for_calibration_25.10.parquet" + ) + print(clips[0]) + PY + )" + + echo "$ALPAMAYO_CLIP_ID" + +Return to the Edge export environment before running the exporter. + +7. Export Alpamayo to TensorRT +------------------------------ + +Load the plugin and run the unified exporter: + +.. code-block:: bash + + cd "$TORCH_TRT_ROOT" + source /path/to/edge-export/bin/activate + + export PYTHONPATH="$TORCH_TRT_ROOT/tools/hf" + export EDGE_LLM_PLUGIN_SO="$EDGE_LLM_ROOT/build-alpamayo-export/libNvInfer_edgellm_plugin.so.1.0" + export LD_LIBRARY_PATH="$(dirname "$EDGE_LLM_PLUGIN_SO"):${LD_LIBRARY_PATH:-}" + + rm -rf "$ENGINE_DIR" + + python tools/hf/run_export.py alpamayo \ + --checkpoint "$ALPAMAYO_FP8_CKPT" \ + --clip-id "$ALPAMAYO_CLIP_ID" \ + --t0-us 5100000 \ + --engine-dir "$ENGINE_DIR" \ + --max-seq-len 4096 \ + --dtype float16 \ + --device cuda:0 + +During a successful run, the exporter: + +1. enables ModelOpt Hugging Face checkpoint restoration; +2. loads the Alpamayo checkpoint and sample clip; +3. captures unpatched eager outputs; +4. temporarily patches the original Qwen3-VL and Alpamayo ``forward`` methods; +5. compiles the vision, language, and action engines; +6. restores the original class methods; +7. compares eager and TensorRT component outputs; +8. exports and executes the outer engine-chaining graph. + +TensorRT autotuning can remain silent for several minutes. High CPU or GPU +utilization during this period usually means the build is still progressing. + +8. Validate the result +---------------------- + +The runner prints: + +* the engine mapping; +* runtime input keys; +* eager-versus-TensorRT parity for each component; +* eager and TensorRT timings; +* the final action velocity shape and mean. + +The final tensor should have shape: + +.. code-block:: text + + velocity (1, 64, 2) + +Inspect the output tree: + +.. code-block:: bash + + ls -lh "$ENGINE_DIR"/vision + ls -lh "$ENGINE_DIR"/language + ls -lh "$ENGINE_DIR"/action + +Expected layout: + +.. code-block:: text + + $ENGINE_DIR/ + ├── vision/ + │ ├── visual.engine + │ └── config.json + ├── language/ + │ ├── language.engine + │ └── config.json + └── action/ + ├── action.engine + └── config.json + +The ``config.json`` files record logical input/output names, output shapes, +dtypes, and engine filenames. The outer ``ExportedProgram`` calls these engines +through ``torch.ops.edge_llm.execute_engine``. + +Parity interpretation +^^^^^^^^^^^^^^^^^^^^^ + +Review component parity independently: + +* vision compares merged Qwen3-VL image features; +* language compares the final hidden states from multimodal prefill; +* action compares one velocity prediction with identical noisy actions, + timestep, position IDs, mask, and prefix K/V. + +Do not accept NaN/Inf values. Investigate large errors before measuring +performance. FP8 tolerances are necessarily looser than FP16, but a result that +is effectively uncorrelated with eager output indicates a packing, mRoPE, +DeepStack, mask, or checkpoint-restore mismatch. + +9. Run focused tests +-------------------- + +From the Torch-TensorRT checkout: + +.. code-block:: bash + + cd "$TORCH_TRT_ROOT" + export PYTHONPATH="$TORCH_TRT_ROOT/tools/hf" + + python -m pytest \ + -o addopts='' \ + tools/hf/exporters/tests/test_edge_exporter.py \ + -k alpamayo \ + -q + +These tests cover registration, fixed-grid and action patch wiring, visual-token +packing, and multimodal generation-position extension. They do not replace the +GPU export and parity run. + +End-to-end checklist +-------------------- + +1. [ ] Access granted for the Alpamayo model and PhysicalAI dataset. +2. [ ] Quantization environment uses Python 3.12 and the recipe-pinned packages. +3. [ ] FP8 checkpoint contains ``config.json`` and ``modelopt_state.pth``. +4. [ ] Quantized checkpoint produces finite minADE on a small evaluation subset. +5. [ ] Edge environment imports the intended Torch-TensorRT installation. +6. [ ] Edge environment uses Transformers 5.4 or newer and ModelOpt 0.44 or newer. +7. [ ] Edge-LLM plugin is built against matching CUDA/TensorRT versions. +8. [ ] ``ctypes.CDLL(EDGE_LLM_PLUGIN_SO)`` succeeds. +9. [ ] ``PYTHONPATH`` contains ``tools/hf`` but does not accidentally shadow + the installed Torch-TensorRT package with an unbuilt ``py/`` directory. +10. [ ] A valid PhysicalAI clip ID and ``t0_us`` are selected. +11. [ ] Vision, language, and action engines are written. +12. [ ] Component parity is finite and within the acceptance threshold. +13. [ ] Final velocity output has shape ``[1, 64, 2]``. + +Troubleshooting +--------------- + +.. list-table:: + :header-rows: 1 + :widths: 38 62 + + * - Symptom + - Resolution + * - ``No module named torch_tensorrt._version`` + - An unbuilt source ``py/`` directory is shadowing the installed package. + Set ``PYTHONPATH=$TORCH_TRT_ROOT/tools/hf`` and verify + ``torch_tensorrt.__file__`` points to the intended installation. + * - ``Set EDGE_LLM_PLUGIN_SO ...`` + - Build the Edge-LLM plugin, set the environment variable to the real + ``.so`` file, and add its directory to ``LD_LIBRARY_PATH``. + * - ``Plugin missing`` + - A documentation placeholder was copied literally. Use the actual build + artifact path and check it with ``test -f``. + * - ``undefined symbol`` while loading the plugin + - The plugin and Python environment use different CUDA or TensorRT + versions. Rebuild the plugin against the same TensorRT SDK. + * - ``No module named transformers.exporters`` + - Upgrade the Edge environment to Transformers 5.4 or newer. Do not use + the recipe's Transformers 4.57 environment for Edge export. + * - ModelOpt state is ignored or quantizers are absent + - Ensure ``modelopt_state.pth`` is beside the checkpoint and ModelOpt is + installed. The Alpamayo loader enables ModelOpt checkpointing before + ``from_pretrained``. + * - Real FP8 checkpoint fails to restore or export + - Retry with the ``--fake_quant`` checkpoint first. Real compressed + checkpoint restore is experimental. + * - Dataset access, 401, or missing clip + - Confirm gated access with ``hf auth whoami`` and verify + ``ALPAMAYO_PAI_LOCAL_DIR`` plus the requested clip ID. + * - Image token count does not match visual features + - Use the same processor, image set, ``image_grid_thw``, and checkpoint + tokenizer. Do not mix processor assets from another Alpamayo run. + * - Language parity is poor while vision is good + - Check multimodal position IDs, RoPE delta extension, image-token masks, + and dense DeepStack insertion. + * - Action parity is poor + - Check prefix K/V layer/head layout, non-causal expert mask, timestep + dtype, and action dimensions ``[64, 2]``. + * - CUDA out of memory during compile + - Stop unrelated GPU workloads, use a larger-memory GPU, and keep + checkpoint/cache/build directories off the root filesystem. + * - TensorRT build appears hung + - TensorRT may be autotuning without console output. Check CPU/GPU + utilization before interrupting it. + +Security and cleanup +-------------------- + +Do not export tokens directly in commands that may be captured in shell history. +If a token appears in logs, revoke it immediately and create a replacement. + +Engine files are specific to the TensorRT version, plugin build, GPU +architecture, shapes, and compilation settings used to produce them. Rebuild +the engines after changing any of those inputs. diff --git a/docsrc/user_guide/edge_exporter.rst b/docsrc/user_guide/edge_exporter.rst index 74009efe52..a63f93480c 100644 --- a/docsrc/user_guide/edge_exporter.rst +++ b/docsrc/user_guide/edge_exporter.rst @@ -20,8 +20,8 @@ difference is what happens inside. Instead of tracing the whole policy as one graph, Edge compiles one TensorRT engine per component, then records a small outer graph that only *calls* those engines. -The code is in ``tools/hf``. The entry points are ``tools/hf/run_pi05_export.py``, -``run_groot_export.py``, and ``run_nemotron_export.py``. +The code is in ``tools/hf``. Use ``tools/hf/run_export.py`` and select a model +family such as ``pi05``, ``groot``, ``nemotron``, or ``alpamayo``. .. note:: @@ -49,6 +49,13 @@ The code is in ``tools/hf``. The entry points are ``tools/hf/run_pi05_export.py` - ``nemotron_h`` - ``language`` - ``nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16`` + * - Alpamayo 1.5 + - ``alpamayo`` + - ``vision``, ``language``, ``action`` + - ``nvidia/Alpamayo-1.5-10B`` or a ModelOpt FP8 checkpoint + +For the complete ModelOpt quantization and export workflow, see +:ref:`alpamayo_fp8_edge_exporter`. Installation ------------ diff --git a/docsrc/user_guide/index.rst b/docsrc/user_guide/index.rst index 09be207dcf..3f6c10f245 100644 --- a/docsrc/user_guide/index.rst +++ b/docsrc/user_guide/index.rst @@ -9,6 +9,7 @@ Conceptual guides and how-tos for Torch-TensorRT. torch_tensorrt_explained compilation/index edge_exporter + alpamayo_fp8_edge_exporter shapes_precision/index runtime_performance/index performance_tuning From 72d15a50a9f3570824106d731e2ab02d743141bc Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Tue, 22 Sep 2026 20:30:25 -0700 Subject: [PATCH 2/4] Add Alpamayo ModelOpt FP8 example --- .../user_guide/alpamayo_fp8_edge_exporter.rst | 19 +- docsrc/user_guide/shapes_precision/index.rst | 1 + .../shapes_precision/quantization.rst | 5 + .../torch_export_alpamayo_modelopt_fp8.py | 285 ++++++++++++++++++ 4 files changed, 307 insertions(+), 3 deletions(-) create mode 100644 examples/dynamo/torch_export_alpamayo_modelopt_fp8.py diff --git a/docsrc/user_guide/alpamayo_fp8_edge_exporter.rst b/docsrc/user_guide/alpamayo_fp8_edge_exporter.rst index 98cc3b46ad..a6b3c40e1b 100644 --- a/docsrc/user_guide/alpamayo_fp8_edge_exporter.rst +++ b/docsrc/user_guide/alpamayo_fp8_edge_exporter.rst @@ -110,7 +110,7 @@ Clone the required repositories: .. code-block:: bash - git clone https://github.com/NVlabs/alpamayo.git "$ALPAMAYO_RECIPES" + git clone https://github.com/NVlabs/alpamayo-recipes.git "$ALPAMAYO_RECIPES" git clone https://github.com/pytorch/TensorRT.git "$TORCH_TRT_ROOT" git clone https://github.com/NVIDIA/TensorRT-Edge-LLM.git "$EDGE_LLM_ROOT" @@ -178,8 +178,21 @@ Verify the important versions: The recipe is defined for Python 3.12, PyTorch 2.8, Transformers 4.57.1, and ModelOpt 0.43. -2. Quantize Alpamayo with ModelOpt FP8 -------------------------------------- +2. Run the ModelOpt Alpamayo FP8 example +---------------------------------------- + +This workflow uses the public +`Alpamayo 1.5 ModelOpt quantization example +`_. +Its ``quantize.py``, ``eval.py``, calibration parquet, and pinned ``uv`` +environment are the source of truth for checkpoint creation and evaluation. +The steps below connect that example to Torch-TensorRT Edge export. + +For a compact code example that calibrates the isolated Alpamayo diffusion +expert with synthetic tensors and compiles it directly with Torch-TensorRT, see +:ref:`torch_export_alpamayo_modelopt_fp8`. Use that example to validate the +ModelOpt Q/DQ conversion path; use this guide's PhysicalAI recipe for +customer-quality calibration and checkpoint evaluation. The exporter consumes a Hugging Face checkpoint restored through ModelOpt's checkpoint integration. For the first TensorRT integration run, use diff --git a/docsrc/user_guide/shapes_precision/index.rst b/docsrc/user_guide/shapes_precision/index.rst index de0883e6e8..6ef3c62d6c 100644 --- a/docsrc/user_guide/shapes_precision/index.rst +++ b/docsrc/user_guide/shapes_precision/index.rst @@ -12,3 +12,4 @@ and reduce model size with INT8/FP8/FP4 quantization via ModelOpt. quantization ../../tutorials/_rendered_examples/dynamo/vgg16_ptq Example: ViT FP8 Quantization <../../tutorials/_rendered_examples/dynamo/quantize_vit_fp8> + Example: Alpamayo ModelOpt FP8 <../../tutorials/_rendered_examples/dynamo/torch_export_alpamayo_modelopt_fp8> diff --git a/docsrc/user_guide/shapes_precision/quantization.rst b/docsrc/user_guide/shapes_precision/quantization.rst index e703e6f372..78c777f9fc 100644 --- a/docsrc/user_guide/shapes_precision/quantization.rst +++ b/docsrc/user_guide/shapes_precision/quantization.rst @@ -26,6 +26,11 @@ Hardware requirements: * **FP8**: NVIDIA Hopper (H100) or newer. * **FP4 (NVFP4)**: NVIDIA Blackwell (B100/B200) or newer; requires TensorRT ≥ 10.8. +For a large vision-language-action example, see +:ref:`torch_export_alpamayo_modelopt_fp8`. For representative PhysicalAI +calibration, checkpoint evaluation, and complete Edge export, see +:ref:`alpamayo_fp8_edge_exporter`. + ---- INT8 / FP8 PTQ Workflow diff --git a/examples/dynamo/torch_export_alpamayo_modelopt_fp8.py b/examples/dynamo/torch_export_alpamayo_modelopt_fp8.py new file mode 100644 index 0000000000..0dcf84122d --- /dev/null +++ b/examples/dynamo/torch_export_alpamayo_modelopt_fp8.py @@ -0,0 +1,285 @@ +""" +.. _torch_export_alpamayo_modelopt_fp8: + +Compiling Alpamayo 1.5 with ModelOpt FP8 quantization +===================================================== + +This example loads +`NVIDIA Alpamayo 1.5 10B `_, +extracts its diffusion expert path, applies NVIDIA ModelOpt FP8 post-training +quantization, and compiles the result with the Torch-TensorRT Dynamo backend. + +Alpamayo's complete trajectory rollout contains Hugging Face autoregressive +generation, mutable KV caches, and a Python-controlled diffusion loop. Those +parts cannot be represented by one static ``torch.export`` graph. This example +therefore compiles the cache-independent expert graph: +``action_in_proj -> expert -> action_out_proj``. + +The calibration loop below uses synthetic tensors so the example is +self-contained. It validates the ModelOpt-to-Torch-TensorRT Q/DQ workflow, but +it is not an accuracy recipe. For representative PhysicalAI calibration, +checkpoint evaluation, and full vision/language/action Edge export, see +:ref:`alpamayo_fp8_edge_exporter` and the +`official Alpamayo ModelOpt recipe +`_. + +Requirements: + +* NVIDIA GPU with FP8 support +* ``nvidia-modelopt[hf]>=0.44.0`` +* ``transformers`` and the Alpamayo 1.5 package +* a compatible Torch-TensorRT installation +* access to the gated Alpamayo checkpoint + +.. code-block:: bash + + pip install "nvidia-modelopt[hf]>=0.44.0" transformers accelerate einops hydra-core + pip install git+https://github.com/NVlabs/alpamayo1.5.git + + export ALPAMAYO_MODEL_ID=nvidia/Alpamayo-1.5-10B + export ALPAMAYO_MODELOPT_CALIB_STEPS=8 + python torch_export_alpamayo_modelopt_fp8.py + +""" + +# %% +# Imports +# ------- + +from __future__ import annotations + +import gc +import os + +import modelopt.torch.quantization as mtq +import torch +import torch_tensorrt +from alpamayo1_5.models.alpamayo1_5 import Alpamayo1_5 +from modelopt.torch.quantization.utils import export_torch_mode + +DEVICE = "cuda:0" +MODEL_ID = os.environ.get("ALPAMAYO_MODEL_ID", "nvidia/Alpamayo-1.5-10B") +BATCH_SIZE = int(os.environ.get("ALPAMAYO_BATCH_SIZE", "1")) +CALIBRATION_STEPS = int(os.environ.get("ALPAMAYO_MODELOPT_CALIB_STEPS", "8")) + + +# %% +# Define the expert graph +# ----------------------- + + +class AlpamayoExpert(torch.nn.Module): + """Cache-independent Alpamayo diffusion expert.""" + + def __init__( + self, + action_in_proj: torch.nn.Module, + expert: torch.nn.Module, + action_out_proj: torch.nn.Module, + action_dims: tuple[int, ...], + non_causal_attention: bool, + ) -> None: + super().__init__() + self.action_in_proj = action_in_proj + self.expert = expert + self.action_out_proj = action_out_proj + self.action_dims = action_dims + self.non_causal_attention = non_causal_attention + + def forward( + self, + noisy_action: torch.Tensor, + timestep: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + batch_size = noisy_action.shape[0] + num_action_tokens = self.action_dims[0] + inputs_embeds = self.action_in_proj(noisy_action, timestep) + if inputs_embeds.dim() == 2: + inputs_embeds = inputs_embeds.view( + batch_size, + num_action_tokens, + -1, + ) + + forward_kwargs = {} + if self.non_causal_attention: + forward_kwargs["is_causal"] = False + hidden_states = self.expert( + inputs_embeds=inputs_embeds, + position_ids=position_ids, + use_cache=False, + return_dict=False, + **forward_kwargs, + )[0] + return self.action_out_proj(hidden_states).reshape( + batch_size, + *self.action_dims, + ) + + +# %% +# Load Alpamayo +# ------------- +# Retain only the action projections and diffusion expert so the VLM does not +# consume memory during calibration and TensorRT compilation. + +print(f"Loading {MODEL_ID} ...") +model = ( + Alpamayo1_5.from_pretrained( + MODEL_ID, + dtype=torch.float16, + attn_implementation="eager", + ) + .to(DEVICE) + .eval() +) + +action_dims = tuple(model.action_space.get_action_space_dims()) +expert_model = ( + AlpamayoExpert( + model.action_in_proj, + model.expert, + model.action_out_proj, + action_dims, + model.config.expert_non_causal_attention, + ) + .to(DEVICE) + .eval() +) + +model.action_in_proj = None +model.expert = None +model.action_out_proj = None +del model +gc.collect() +torch.cuda.empty_cache() + + +# %% +# Create example inputs +# --------------------- + +num_action_tokens, action_width = action_dims +noisy_action = torch.randn( + BATCH_SIZE, + num_action_tokens, + action_width, + device=DEVICE, + dtype=torch.float16, +) +timestep = torch.rand( + BATCH_SIZE, + 1, + 1, + device=DEVICE, + dtype=torch.float16, +) +position_ids = ( + torch.arange(num_action_tokens, device=DEVICE) + .view(1, 1, num_action_tokens) + .expand(3, BATCH_SIZE, num_action_tokens) + .clone() +) +example_inputs = (noisy_action, timestep, position_ids) + + +# %% +# Calibrate ModelOpt FP8 +# ---------------------- +# ``FP8_DEFAULT_CFG`` quantizes eligible Linear inputs and weights. The +# synthetic loop exercises the complete isolated expert graph. Production +# calibration should use representative PhysicalAI clips as described in the +# user guide. + + +def calibration_loop(module: torch.nn.Module) -> None: + generator = torch.Generator(device=DEVICE).manual_seed(1234) + with torch.no_grad(): + for _ in range(CALIBRATION_STEPS): + calib_action = torch.randn( + noisy_action.shape, + generator=generator, + device=DEVICE, + dtype=noisy_action.dtype, + ) + calib_timestep = torch.rand( + timestep.shape, + generator=generator, + device=DEVICE, + dtype=timestep.dtype, + ) + module(calib_action, calib_timestep, position_ids) + + +mtq.quantize( + expert_model, + mtq.FP8_DEFAULT_CFG, + forward_loop=calibration_loop, +) +mtq.print_quant_summary(expert_model) + + +# %% +# Capture the quantized eager reference +# ------------------------------------- + +with torch.no_grad(): + eager_output = expert_model(*example_inputs) + + +# %% +# Export ModelOpt Q/DQ operations +# ------------------------------- +# ``export_torch_mode`` lowers ModelOpt TensorQuantizers to explicit +# ``torch.ops.tensorrt`` Q/DQ operations that Torch-TensorRT converters +# recognize. Keep weights in fake-quant form until after export. + +with export_torch_mode(): + exported_program = torch.export.export( + expert_model, + example_inputs, + strict=False, + ) + + +# %% +# Compile with Torch-TensorRT +# --------------------------- + +compiled_expert = torch_tensorrt.dynamo.compile( + exported_program, + arg_inputs=list(example_inputs), + truncate_double=True, + min_block_size=1, + require_full_compilation=True, + immutable_weights=False, + offload_module_to_cpu=True, +) + +del exported_program, expert_model +gc.collect() +torch.cuda.empty_cache() + + +# %% +# Compare eager and TensorRT outputs +# ---------------------------------- + +with torch.no_grad(): + trt_output = compiled_expert(*example_inputs) +if isinstance(trt_output, (tuple, list)): + trt_output = trt_output[0] +trt_output = trt_output.to(device=eager_output.device, dtype=eager_output.dtype) + +absolute_difference = (trt_output.float() - eager_output.float()).abs() +close = torch.isclose( + trt_output.float(), + eager_output.float(), + rtol=1e-2, + atol=1e-2, +) +print(f"output shape: {tuple(trt_output.shape)}") +print(f"max |Δoutput|: {absolute_difference.max().item():.6g}") +print(f"mean |Δoutput|: {absolute_difference.mean().item():.6g}") +print(f"close elements: {close.float().mean().item() * 100:.2f}%") From 152511b0fbe0d77c598e9e460a1c9d0cf4350a8d Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Tue, 22 Sep 2026 20:35:10 -0700 Subject: [PATCH 3/4] Match official Alpamayo ModelOpt recipe --- .../user_guide/alpamayo_fp8_edge_exporter.rst | 62 ++-- docsrc/user_guide/alpamayo_modelopt_fp8.rst | 289 ++++++++++++++++++ docsrc/user_guide/index.rst | 1 + docsrc/user_guide/shapes_precision/index.rst | 1 - .../shapes_precision/quantization.rst | 6 +- .../torch_export_alpamayo_modelopt_fp8.py | 285 ----------------- 6 files changed, 321 insertions(+), 323 deletions(-) create mode 100644 docsrc/user_guide/alpamayo_modelopt_fp8.rst delete mode 100644 examples/dynamo/torch_export_alpamayo_modelopt_fp8.py diff --git a/docsrc/user_guide/alpamayo_fp8_edge_exporter.rst b/docsrc/user_guide/alpamayo_fp8_edge_exporter.rst index a6b3c40e1b..135d783db7 100644 --- a/docsrc/user_guide/alpamayo_fp8_edge_exporter.rst +++ b/docsrc/user_guide/alpamayo_fp8_edge_exporter.rst @@ -186,18 +186,10 @@ This workflow uses the public `_. Its ``quantize.py``, ``eval.py``, calibration parquet, and pinned ``uv`` environment are the source of truth for checkpoint creation and evaluation. -The steps below connect that example to Torch-TensorRT Edge export. - -For a compact code example that calibrates the isolated Alpamayo diffusion -expert with synthetic tensors and compiles it directly with Torch-TensorRT, see -:ref:`torch_export_alpamayo_modelopt_fp8`. Use that example to validate the -ModelOpt Q/DQ conversion path; use this guide's PhysicalAI recipe for -customer-quality calibration and checkpoint evaluation. - -The exporter consumes a Hugging Face checkpoint restored through ModelOpt's -checkpoint integration. For the first TensorRT integration run, use -``--fake_quant``. It preserves FP16 checkpoint weights while retaining the FP8 -Q/DQ graph that TensorRT lowers. This is the lowest-risk export format. +The quantization procedure is also documented separately in +:ref:`alpamayo_modelopt_fp8`. The steps below use the same full-model +PhysicalAI calibration and compressed FP8 output, then continue into +Torch-TensorRT Edge export. .. code-block:: bash @@ -205,60 +197,62 @@ Q/DQ graph that TensorRT lowers. This is the lowest-risk export format. source am15_quant/bin/activate uv run --active quantize.py \ - --ckpt nvidia/Alpamayo-1.5-10B \ - --quant_format fp8 \ - --num_of_calib_clips 100 \ - --save_model_dir "$QUANT_OUTPUT" \ - --fake_quant + --quant_format=fp8 \ + --num_of_calib_clips=100 \ + --save_model_dir=./outputs The expected output directory is: .. code-block:: text - outputs/alpamayo1.5_fp8_calib100_fakequant/ + outputs/alpamayo1.5_fp8_calib100/ ├── config.json ├── modelopt_state.pth ├── model*.safetensors └── tokenizer and processor assets +The calibration loop exercises the full VLM rollout and diffusion path with +the selected PhysicalAI clips. By default, ``quantize.py`` calls +``mtq.compress(model)`` before saving, so this checkpoint contains real FP8 +weights. + Confirm that ModelOpt state was saved: .. code-block:: bash - export ALPAMAYO_FP8_CKPT="$QUANT_OUTPUT/alpamayo1.5_fp8_calib100_fakequant" + export ALPAMAYO_FP8_CKPT="$QUANT_DIR/outputs/alpamayo1.5_fp8_calib100" test -f "$ALPAMAYO_FP8_CKPT/config.json" test -f "$ALPAMAYO_FP8_CKPT/modelopt_state.pth" ls -lh "$ALPAMAYO_FP8_CKPT" -Real compressed FP8 checkpoint -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -Omit ``--fake_quant`` to call ``mtq.compress`` and save a real FP8 parameter -payload: +For a long-running calibration, use the upstream background command: .. code-block:: bash - uv run --active quantize.py \ - --quant_format fp8 \ - --num_of_calib_clips 100 \ - --save_model_dir "$QUANT_OUTPUT" - -This produces ``alpamayo1.5_fp8_calib100`` and reduces checkpoint and runtime -weight memory. Real-quant restore and downstream compilation are still -experimental in ModelOpt. Establish the fake-quant export first, then validate -the compressed checkpoint separately. + nohup uv run --active quantize.py \ + --quant_format=fp8 \ + --num_of_calib_clips=100 \ + --save_model_dir=./outputs \ + > quantize_fp8.log 2>&1 & 3. Evaluate the quantized checkpoint ------------------------------------ -Evaluate a small subset before compiling TensorRT engines: +Run the official evaluation command: .. code-block:: bash cd "$QUANT_DIR" source am15_quant/bin/activate + uv run --active eval.py \ + --ckpt ./outputs/alpamayo1.5_fp8_calib100 + +For a shorter smoke test before the complete evaluation: + +.. code-block:: bash + uv run --active eval.py \ --ckpt "$ALPAMAYO_FP8_CKPT" \ --limit 10 \ diff --git a/docsrc/user_guide/alpamayo_modelopt_fp8.rst b/docsrc/user_guide/alpamayo_modelopt_fp8.rst new file mode 100644 index 0000000000..45f36c9bee --- /dev/null +++ b/docsrc/user_guide/alpamayo_modelopt_fp8.rst @@ -0,0 +1,289 @@ +.. _alpamayo_modelopt_fp8: + +Alpamayo 1.5 ModelOpt FP8 Quantization +====================================== + +This guide follows the public +`Alpamayo 1.5 quantization recipe +`_ +to post-training quantize the complete Alpamayo model with NVIDIA ModelOpt and +evaluate the resulting checkpoint with minADE. + +Unlike an isolated expert smoke test, this workflow: + +* loads the full Alpamayo vision-language-action model; +* calibrates with PhysicalAI driving clips; +* exercises both VLM rollout and diffusion paths; +* compresses fake-quant weights to real FP8 by default; +* saves ``modelopt_state.pth`` with the Hugging Face checkpoint; +* reloads the checkpoint and evaluates trajectory accuracy. + +The commands below intentionally match the upstream recipe. After completing +this guide, continue with :ref:`alpamayo_fp8_edge_exporter` to build TensorRT +engines. + +Prerequisites +------------- + +The upstream recipe is tested with: + +* NVIDIA RTX 5090 with CUDA 12; +* NVIDIA B300 with CUDA 13; +* Python 3.12; +* PyTorch 2.8.0; +* torchvision 0.23.0; +* NVIDIA ModelOpt 0.43.0. + +Request access to: + +* `Alpamayo-1.5-10B `_ +* `PhysicalAI-Autonomous-Vehicles + `_ + +Clone the official recipe: + +.. code-block:: bash + + export YOUR_HOME=/path/to/your/workspace + git clone https://github.com/NVlabs/alpamayo-recipes.git \ + "$YOUR_HOME/alpamayo-recipes" + +1. Create the recipe environment +-------------------------------- + +.. code-block:: bash + + export UV_CACHE_DIR="$YOUR_HOME/.cache/uv" + + cd "$YOUR_HOME/alpamayo-recipes/recipes/alpamayo1_5_quant" + uv venv am15_quant + source am15_quant/bin/activate + + # Install all dependencies except flash-attn first so torch is available. + uv sync --active --no-install-package flash-attn + + # Build flash-attn against the installed torch. + MAX_JOBS=4 uv sync --active + +Verify the pinned environment: + +.. code-block:: bash + + python - <<'PY' + import modelopt + import torch + import torchvision + import transformers + + print("torch", torch.__version__) + print("torchvision", torchvision.__version__) + print("transformers", transformers.__version__) + print("modelopt", modelopt.__version__) + print("cuda available", torch.cuda.is_available()) + PY + +2. Configure model, dataset, and cache paths +------------------------------------------- + +.. code-block:: bash + + export ALPAMAYO_WORKSPACE="$YOUR_HOME/alpamayo-recipes" + export ALPAMAYO_MODEL_DIR="$YOUR_HOME/alpamayo_model_converted_from_hf" + export ALPAMAYO_PAI_LOCAL_DIR="$YOUR_HOME/PAI_mini" + export ALPAMAYO_LOG_DIR="$YOUR_HOME/alpamayo_logs" + export HF_HOME="$YOUR_HOME/.cache/huggingface" + +``ALPAMAYO_PAI_LOCAL_DIR`` must contain the PhysicalAI data expected by +``load_physical_aiavdataset``. The recipe directory already contains: + +* ``0417_5k_train_set_for_calibration_25.10.parquet`` for calibration; +* ``1005_7cam_gold_eval_metadb_public.parquet`` for evaluation. + +For offline execution after the assets are cached: + +.. code-block:: bash + + export HF_HUB_OFFLINE=1 + export TRANSFORMERS_OFFLINE=1 + +3. Authenticate with Hugging Face +--------------------------------- + +Authenticate interactively so the token is not written into shell history: + +.. code-block:: bash + + hf auth login + hf auth whoami + +4. Quantize the full model to FP8 +-------------------------------- + +Run the upstream FP8 command: + +.. code-block:: bash + + cd "$YOUR_HOME/alpamayo-recipes/recipes/alpamayo1_5_quant" + source am15_quant/bin/activate + + uv run --active quantize.py \ + --quant_format=fp8 \ + --num_of_calib_clips=100 \ + --save_model_dir=./outputs + +The calibration loop loads the requested PhysicalAI clips and calls +``sample_trajectories_from_data_with_vlm_rollout``. This observes activation +ranges through the VLM, expert, action projections, and diffusion-related +paths. + +By default, ``quantize.py`` calls ``mtq.compress(model)`` before saving. The +checkpoint therefore contains real FP8 weights and is expected at: + +.. code-block:: text + + outputs/alpamayo1.5_fp8_calib100/ + ├── config.json + ├── modelopt_state.pth + ├── model*.safetensors + └── tokenizer and processor assets + +Capture logs for a long-running calibration: + +.. code-block:: bash + + nohup uv run --active quantize.py \ + --quant_format=fp8 \ + --num_of_calib_clips=100 \ + --save_model_dir=./outputs \ + > quantize_fp8.log 2>&1 & + + tail -f quantize_fp8.log + +Check the result: + +.. code-block:: bash + + export ALPAMAYO_FP8_CKPT="$PWD/outputs/alpamayo1.5_fp8_calib100" + + test -f "$ALPAMAYO_FP8_CKPT/config.json" + test -f "$ALPAMAYO_FP8_CKPT/modelopt_state.pth" + ls -lh "$ALPAMAYO_FP8_CKPT" + +Fake-quant checkpoint option +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +The upstream recipe also supports ``--fake_quant`` for downstream SDK +debugging. It preserves the original weights while saving Q/DQ state: + +.. code-block:: bash + + uv run --active quantize.py \ + --quant_format=fp8 \ + --num_of_calib_clips=100 \ + --save_model_dir=./outputs \ + --fake_quant + +This writes ``alpamayo1.5_fp8_calib100_fakequant``. It is not the default +workflow and does not provide the real-weight memory reduction of the +compressed checkpoint. + +5. Evaluate the FP8 checkpoint +------------------------------ + +Run the official evaluation command: + +.. code-block:: bash + + uv run --active eval.py \ + --ckpt ./outputs/alpamayo1.5_fp8_calib100 + +For a short smoke test before the complete evaluation: + +.. code-block:: bash + + uv run --active eval.py \ + --ckpt ./outputs/alpamayo1.5_fp8_calib100 \ + --limit 10 \ + --num_traj_samples 6 \ + --seed 42 \ + --print_every 1 + +Capture a full evaluation in the background: + +.. code-block:: bash + + nohup uv run --active eval.py \ + --ckpt ./outputs/alpamayo1.5_fp8_calib100 \ + > eval_fp8.log 2>&1 & + + tail -f eval_fp8.log + +A correct run reports: + +* the number of clip IDs loaded from the evaluation parquet; +* per-clip minADE and inference time; +* failed clips, if any; +* average minADE over successful clips; +* average evaluation time per clip. + +Compare the FP8 result with a base-model run using the same settings: + +.. code-block:: bash + + uv run --active eval.py \ + --ckpt nvidia/Alpamayo-1.5-10B \ + --limit 10 \ + --num_traj_samples 6 \ + --seed 42 \ + --print_every 1 + +Keep ``--limit``, ``--num_traj_samples``, and ``--seed`` identical when +comparing checkpoints. + +6. Continue to TensorRT Edge export +----------------------------------- + +The checkpoint consumed by the Edge exporter is: + +.. code-block:: bash + + export ALPAMAYO_FP8_CKPT="$YOUR_HOME/alpamayo-recipes/recipes/alpamayo1_5_quant/outputs/alpamayo1.5_fp8_calib100" + +Continue with :ref:`alpamayo_fp8_edge_exporter`. That guide creates a separate +export environment, restores ``modelopt_state.pth``, builds the Edge-LLM +plugin, compiles the vision/language/action engines, and checks component +parity. + +Troubleshooting +--------------- + +``401`` or gated-repository error +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Confirm access and run ``hf auth login`` again. Do not paste tokens into +commands or logs. + +``No space left on device`` +^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Move ``HF_HOME``, ``UV_CACHE_DIR``, and ``outputs`` to storage with enough +capacity for the base and quantized checkpoints. + +FlashAttention build failure +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Install all other dependencies first, confirm PyTorch imports, then rerun +``MAX_JOBS=4 uv sync --active``. + +ModelOpt state is not restored +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Confirm ``modelopt_state.pth`` is in the checkpoint directory and that the +loading process calls ``mto.enable_huggingface_checkpointing()`` before +``Alpamayo1_5.from_pretrained``. + +Evaluation produces no successful clips +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Verify ``ALPAMAYO_PAI_LOCAL_DIR``, the parquet path, dataset access, and +``t0_us``. Test one known clip directly before starting the full evaluation. diff --git a/docsrc/user_guide/index.rst b/docsrc/user_guide/index.rst index 3f6c10f245..450e471c6c 100644 --- a/docsrc/user_guide/index.rst +++ b/docsrc/user_guide/index.rst @@ -9,6 +9,7 @@ Conceptual guides and how-tos for Torch-TensorRT. torch_tensorrt_explained compilation/index edge_exporter + alpamayo_modelopt_fp8 alpamayo_fp8_edge_exporter shapes_precision/index runtime_performance/index diff --git a/docsrc/user_guide/shapes_precision/index.rst b/docsrc/user_guide/shapes_precision/index.rst index 6ef3c62d6c..de0883e6e8 100644 --- a/docsrc/user_guide/shapes_precision/index.rst +++ b/docsrc/user_guide/shapes_precision/index.rst @@ -12,4 +12,3 @@ and reduce model size with INT8/FP8/FP4 quantization via ModelOpt. quantization ../../tutorials/_rendered_examples/dynamo/vgg16_ptq Example: ViT FP8 Quantization <../../tutorials/_rendered_examples/dynamo/quantize_vit_fp8> - Example: Alpamayo ModelOpt FP8 <../../tutorials/_rendered_examples/dynamo/torch_export_alpamayo_modelopt_fp8> diff --git a/docsrc/user_guide/shapes_precision/quantization.rst b/docsrc/user_guide/shapes_precision/quantization.rst index 78c777f9fc..151d2a8e96 100644 --- a/docsrc/user_guide/shapes_precision/quantization.rst +++ b/docsrc/user_guide/shapes_precision/quantization.rst @@ -26,9 +26,9 @@ Hardware requirements: * **FP8**: NVIDIA Hopper (H100) or newer. * **FP4 (NVFP4)**: NVIDIA Blackwell (B100/B200) or newer; requires TensorRT ≥ 10.8. -For a large vision-language-action example, see -:ref:`torch_export_alpamayo_modelopt_fp8`. For representative PhysicalAI -calibration, checkpoint evaluation, and complete Edge export, see +For full-model Alpamayo 1.5 quantization with representative PhysicalAI +calibration and minADE evaluation, see :ref:`alpamayo_modelopt_fp8`. To +continue from that checkpoint through TensorRT Edge export, see :ref:`alpamayo_fp8_edge_exporter`. ---- diff --git a/examples/dynamo/torch_export_alpamayo_modelopt_fp8.py b/examples/dynamo/torch_export_alpamayo_modelopt_fp8.py deleted file mode 100644 index 0dcf84122d..0000000000 --- a/examples/dynamo/torch_export_alpamayo_modelopt_fp8.py +++ /dev/null @@ -1,285 +0,0 @@ -""" -.. _torch_export_alpamayo_modelopt_fp8: - -Compiling Alpamayo 1.5 with ModelOpt FP8 quantization -===================================================== - -This example loads -`NVIDIA Alpamayo 1.5 10B `_, -extracts its diffusion expert path, applies NVIDIA ModelOpt FP8 post-training -quantization, and compiles the result with the Torch-TensorRT Dynamo backend. - -Alpamayo's complete trajectory rollout contains Hugging Face autoregressive -generation, mutable KV caches, and a Python-controlled diffusion loop. Those -parts cannot be represented by one static ``torch.export`` graph. This example -therefore compiles the cache-independent expert graph: -``action_in_proj -> expert -> action_out_proj``. - -The calibration loop below uses synthetic tensors so the example is -self-contained. It validates the ModelOpt-to-Torch-TensorRT Q/DQ workflow, but -it is not an accuracy recipe. For representative PhysicalAI calibration, -checkpoint evaluation, and full vision/language/action Edge export, see -:ref:`alpamayo_fp8_edge_exporter` and the -`official Alpamayo ModelOpt recipe -`_. - -Requirements: - -* NVIDIA GPU with FP8 support -* ``nvidia-modelopt[hf]>=0.44.0`` -* ``transformers`` and the Alpamayo 1.5 package -* a compatible Torch-TensorRT installation -* access to the gated Alpamayo checkpoint - -.. code-block:: bash - - pip install "nvidia-modelopt[hf]>=0.44.0" transformers accelerate einops hydra-core - pip install git+https://github.com/NVlabs/alpamayo1.5.git - - export ALPAMAYO_MODEL_ID=nvidia/Alpamayo-1.5-10B - export ALPAMAYO_MODELOPT_CALIB_STEPS=8 - python torch_export_alpamayo_modelopt_fp8.py - -""" - -# %% -# Imports -# ------- - -from __future__ import annotations - -import gc -import os - -import modelopt.torch.quantization as mtq -import torch -import torch_tensorrt -from alpamayo1_5.models.alpamayo1_5 import Alpamayo1_5 -from modelopt.torch.quantization.utils import export_torch_mode - -DEVICE = "cuda:0" -MODEL_ID = os.environ.get("ALPAMAYO_MODEL_ID", "nvidia/Alpamayo-1.5-10B") -BATCH_SIZE = int(os.environ.get("ALPAMAYO_BATCH_SIZE", "1")) -CALIBRATION_STEPS = int(os.environ.get("ALPAMAYO_MODELOPT_CALIB_STEPS", "8")) - - -# %% -# Define the expert graph -# ----------------------- - - -class AlpamayoExpert(torch.nn.Module): - """Cache-independent Alpamayo diffusion expert.""" - - def __init__( - self, - action_in_proj: torch.nn.Module, - expert: torch.nn.Module, - action_out_proj: torch.nn.Module, - action_dims: tuple[int, ...], - non_causal_attention: bool, - ) -> None: - super().__init__() - self.action_in_proj = action_in_proj - self.expert = expert - self.action_out_proj = action_out_proj - self.action_dims = action_dims - self.non_causal_attention = non_causal_attention - - def forward( - self, - noisy_action: torch.Tensor, - timestep: torch.Tensor, - position_ids: torch.Tensor, - ) -> torch.Tensor: - batch_size = noisy_action.shape[0] - num_action_tokens = self.action_dims[0] - inputs_embeds = self.action_in_proj(noisy_action, timestep) - if inputs_embeds.dim() == 2: - inputs_embeds = inputs_embeds.view( - batch_size, - num_action_tokens, - -1, - ) - - forward_kwargs = {} - if self.non_causal_attention: - forward_kwargs["is_causal"] = False - hidden_states = self.expert( - inputs_embeds=inputs_embeds, - position_ids=position_ids, - use_cache=False, - return_dict=False, - **forward_kwargs, - )[0] - return self.action_out_proj(hidden_states).reshape( - batch_size, - *self.action_dims, - ) - - -# %% -# Load Alpamayo -# ------------- -# Retain only the action projections and diffusion expert so the VLM does not -# consume memory during calibration and TensorRT compilation. - -print(f"Loading {MODEL_ID} ...") -model = ( - Alpamayo1_5.from_pretrained( - MODEL_ID, - dtype=torch.float16, - attn_implementation="eager", - ) - .to(DEVICE) - .eval() -) - -action_dims = tuple(model.action_space.get_action_space_dims()) -expert_model = ( - AlpamayoExpert( - model.action_in_proj, - model.expert, - model.action_out_proj, - action_dims, - model.config.expert_non_causal_attention, - ) - .to(DEVICE) - .eval() -) - -model.action_in_proj = None -model.expert = None -model.action_out_proj = None -del model -gc.collect() -torch.cuda.empty_cache() - - -# %% -# Create example inputs -# --------------------- - -num_action_tokens, action_width = action_dims -noisy_action = torch.randn( - BATCH_SIZE, - num_action_tokens, - action_width, - device=DEVICE, - dtype=torch.float16, -) -timestep = torch.rand( - BATCH_SIZE, - 1, - 1, - device=DEVICE, - dtype=torch.float16, -) -position_ids = ( - torch.arange(num_action_tokens, device=DEVICE) - .view(1, 1, num_action_tokens) - .expand(3, BATCH_SIZE, num_action_tokens) - .clone() -) -example_inputs = (noisy_action, timestep, position_ids) - - -# %% -# Calibrate ModelOpt FP8 -# ---------------------- -# ``FP8_DEFAULT_CFG`` quantizes eligible Linear inputs and weights. The -# synthetic loop exercises the complete isolated expert graph. Production -# calibration should use representative PhysicalAI clips as described in the -# user guide. - - -def calibration_loop(module: torch.nn.Module) -> None: - generator = torch.Generator(device=DEVICE).manual_seed(1234) - with torch.no_grad(): - for _ in range(CALIBRATION_STEPS): - calib_action = torch.randn( - noisy_action.shape, - generator=generator, - device=DEVICE, - dtype=noisy_action.dtype, - ) - calib_timestep = torch.rand( - timestep.shape, - generator=generator, - device=DEVICE, - dtype=timestep.dtype, - ) - module(calib_action, calib_timestep, position_ids) - - -mtq.quantize( - expert_model, - mtq.FP8_DEFAULT_CFG, - forward_loop=calibration_loop, -) -mtq.print_quant_summary(expert_model) - - -# %% -# Capture the quantized eager reference -# ------------------------------------- - -with torch.no_grad(): - eager_output = expert_model(*example_inputs) - - -# %% -# Export ModelOpt Q/DQ operations -# ------------------------------- -# ``export_torch_mode`` lowers ModelOpt TensorQuantizers to explicit -# ``torch.ops.tensorrt`` Q/DQ operations that Torch-TensorRT converters -# recognize. Keep weights in fake-quant form until after export. - -with export_torch_mode(): - exported_program = torch.export.export( - expert_model, - example_inputs, - strict=False, - ) - - -# %% -# Compile with Torch-TensorRT -# --------------------------- - -compiled_expert = torch_tensorrt.dynamo.compile( - exported_program, - arg_inputs=list(example_inputs), - truncate_double=True, - min_block_size=1, - require_full_compilation=True, - immutable_weights=False, - offload_module_to_cpu=True, -) - -del exported_program, expert_model -gc.collect() -torch.cuda.empty_cache() - - -# %% -# Compare eager and TensorRT outputs -# ---------------------------------- - -with torch.no_grad(): - trt_output = compiled_expert(*example_inputs) -if isinstance(trt_output, (tuple, list)): - trt_output = trt_output[0] -trt_output = trt_output.to(device=eager_output.device, dtype=eager_output.dtype) - -absolute_difference = (trt_output.float() - eager_output.float()).abs() -close = torch.isclose( - trt_output.float(), - eager_output.float(), - rtol=1e-2, - atol=1e-2, -) -print(f"output shape: {tuple(trt_output.shape)}") -print(f"max |Δoutput|: {absolute_difference.max().item():.6g}") -print(f"mean |Δoutput|: {absolute_difference.mean().item():.6g}") -print(f"close elements: {close.float().mean().item() * 100:.2f}%") From e7d526bfb9cc806ce0eda93f41a66bf78173327b Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Tue, 22 Sep 2026 20:39:01 -0700 Subject: [PATCH 4/4] Fix Alpamayo guide headings --- docsrc/user_guide/alpamayo_fp8_edge_exporter.rst | 2 +- docsrc/user_guide/alpamayo_modelopt_fp8.rst | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docsrc/user_guide/alpamayo_fp8_edge_exporter.rst b/docsrc/user_guide/alpamayo_fp8_edge_exporter.rst index 135d783db7..e1179ec07a 100644 --- a/docsrc/user_guide/alpamayo_fp8_edge_exporter.rst +++ b/docsrc/user_guide/alpamayo_fp8_edge_exporter.rst @@ -323,7 +323,7 @@ Torch-TensorRT from the selected source revision. Confirm that importing the source tree in place. 5. Build the TensorRT-Edge-LLM plugin ------------------------------------- +--------------------------------------------- The language engine contains TensorRT-Edge-LLM attention plugins. Build the plugin against the same CUDA and TensorRT major versions used by diff --git a/docsrc/user_guide/alpamayo_modelopt_fp8.rst b/docsrc/user_guide/alpamayo_modelopt_fp8.rst index 45f36c9bee..e50ff7ff54 100644 --- a/docsrc/user_guide/alpamayo_modelopt_fp8.rst +++ b/docsrc/user_guide/alpamayo_modelopt_fp8.rst @@ -83,7 +83,7 @@ Verify the pinned environment: PY 2. Configure model, dataset, and cache paths -------------------------------------------- +----------------------------------------------- .. code-block:: bash @@ -107,7 +107,7 @@ For offline execution after the assets are cached: export TRANSFORMERS_OFFLINE=1 3. Authenticate with Hugging Face ---------------------------------- +------------------------------------- Authenticate interactively so the token is not written into shell history: @@ -117,7 +117,7 @@ Authenticate interactively so the token is not written into shell history: hf auth whoami 4. Quantize the full model to FP8 --------------------------------- +------------------------------------- Run the upstream FP8 command: