From 70ca7f52979fd378948481023cceaa637fb3d38d Mon Sep 17 00:00:00 2001 From: hungho77 Date: Thu, 13 Aug 2026 15:47:12 +0700 Subject: [PATCH 01/30] feat(internvla-n1-dualvln): add recipe scaffolding, docs and scheme matrix Adds recipes/internvla-n1-dualvln/, the front matter for porting a working FP8 quantization and TensorRT-Edge-LLM deployment pipeline for InternVLA-N1-DualVLN on Jetson Thor. This commit is structure and documentation only; the Python and shell entrypoints land in follow-ups. The recipe is laid out along its dependency boundary rather than by file type. InternVLA-N1-DualVLN declares model_type internvla_n1 and ships no modeling code, so it cannot be loaded with trust_remote_code -- the class has to come from the InternNav repository. A repackaging step strips the eight System-1 tensor prefixes and rewrites config.json to a stock Qwen2.5-VL, using nothing but safetensors surgery, which takes InternNav off the hot path for quantization, export, engine build and latent verification. Only System-1 export and the agent-level checks need INTERNNAV_PATH, and the layout says so per step. configs/schemes.yaml encodes the scheme x strategy validity matrix so impossible combinations are rejected before the model loads rather than forty minutes into a run. Two entries are hardware facts: NVFP4 cannot touch the vision tower because the ViT MLP intermediate_size is 3420 and 3420/16 is not an integer, and the KV cache stays FP8 under NVFP4 weights because NVFP4 KV needs sm100f while Thor is sm110. NVFP4 weights are marked experimental rather than supported: they quantize, export and generate fluent text, but z_latents cosine falls to 0.647, which breaks the System2 -> System1 bridge. For a navigation model the bridge is the acceptance metric, not text fluency, so the README leads with that distinction. The README also retracts an earlier claim carried in the source project that FP16 engines are broken on Thor and FP8 is therefore mandatory. The garbage came from a Myelin fc_h_fusion miscompile at TensorRT 10.13; with the workaround applied the unquantized FP16 engine is in fact the highest-fidelity variant measured. FP8 is recommended on size and latency, not correctness. numpy and scipy are deliberately absent from the pyproject extra. This recipe needs numpy 1.x while three other extras pin numpy==2.2.6, so rather than assert a resolution outcome that has not been tested here, they are installed as a documented post-sync step. Note the repository already carries three mutually exclusive transformers pins across extras, so this is not a new class of problem; run uv lock and check before concluding anything about it. Directory names avoid build/ and lib/, which the root .gitignore silently swallows. A recipe-level .gitignore covers the multi-gigabyte engine and ONNX artifacts, which nothing else in the repository excludes. (cherry picked from commit 7b62c501871b408cfca3d5a0ae235cafdc92dcbd) --- README.md | 24 ++ pyproject.toml | 23 ++ recipes/internvla-n1-dualvln/.gitignore | 8 + recipes/internvla-n1-dualvln/Makefile | 137 +++++++++ recipes/internvla-n1-dualvln/README.md | 259 ++++++++++++++++++ .../internvla-n1-dualvln/configs/schemes.yaml | 109 ++++++++ .../requirements-torch.txt | 31 +++ recipes/internvla-n1-dualvln/requirements.txt | 50 ++++ 8 files changed, 641 insertions(+) create mode 100644 recipes/internvla-n1-dualvln/.gitignore create mode 100644 recipes/internvla-n1-dualvln/Makefile create mode 100644 recipes/internvla-n1-dualvln/README.md create mode 100644 recipes/internvla-n1-dualvln/configs/schemes.yaml create mode 100644 recipes/internvla-n1-dualvln/requirements-torch.txt create mode 100644 recipes/internvla-n1-dualvln/requirements.txt diff --git a/README.md b/README.md index e391dfd..40df38b 100644 --- a/README.md +++ b/README.md @@ -28,6 +28,7 @@ The repository currently includes: | `recipes/qwen36-27b/` | Universal causal LLM quantization with llmcompressor | `uv sync --extra qwen36-27b` | | `recipes/qwen36-moe-35b-nvfp4/` | INT8, FP8, and NVFP4 quantization for hybrid MoE models | `uv sync --extra qwen36-moe-35b-nvfp4` | | `recipes/cosmos-reason2/` | NVFP4 quantization for Cosmos Reason2 (2B, 8B) with llmcompressor and Hugging Face export | `uv sync --extra cosmos-reason2` | +| `recipes/internvla-n1-dualvln/` | FP8 quantization and TensorRT-Edge-LLM deployment for InternVLA-N1-DualVLN, a dual-system vision-language navigation model, on NVIDIA Jetson Thor | `uv sync --extra internvla-n1-dualvln` | ## Architecture @@ -51,6 +52,7 @@ uv sync --extra qwen3-asr uv sync --extra qwen36-27b uv sync --extra qwen36-moe-35b-nvfp4 uv sync --extra cosmos-reason2 +uv sync --extra internvla-n1-dualvln ``` ## Quick Start @@ -124,6 +126,27 @@ MODEL_PATH=/path/to/Cosmos-Reason2-2B OUTPUT_PATH=/path/to/output ./quantize.sh MODEL_PATH=/path/to/Cosmos-Reason2-8B OUTPUT_PATH=/path/to/output ./quantize.sh ``` +### InternVLA-N1 DualVLN + +```bash +cd recipes/internvla-n1-dualvln + +export INTERNVLA_CKPT=/path/to/InternVLA-N1-DualVLN +make repackage # strip System 1 -> stock Qwen2.5-VL System 2 checkpoint +make quantize-fp8 # FP8 W8A8, LLM backbone only +make export-build # ONNX export + TensorRT-Edge-LLM engines +make verify-latents # acceptance gate: System2 -> System1 bridge fidelity +``` + +The recipe is split by dependency boundary: + +- `quantize/` turns the checkpoint into a quantized stock Qwen2.5-VL — needs no InternNav +- `trt-edgellm/` builds and verifies the engines; only System 1 export and the agent-level + checks require `INTERNNAV_PATH` + +Acceptance is measured on `z_latents` cosine (the System 2 → System 1 bridge), not on text +fluency — a checkpoint can caption correctly and still navigate wrongly. + ## Benchmark and Results Representative results from `recipes/qwen3-asr/`: @@ -155,6 +178,7 @@ For the edge results, Jetson measurements use unified memory while RTX measureme │ ├── _template/ │ ├── cosmos-reason2/ │ ├── gemma4/ +│ ├── internvla-n1-dualvln/ │ ├── qwen3-asr/ │ ├── qwen36-27b/ │ └── qwen36-moe-35b-nvfp4/ diff --git a/pyproject.toml b/pyproject.toml index 7add770..8cd8e28 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -83,6 +83,29 @@ cosmos-reason2 = [ "tqdm==4.67.3", ] +# numpy and scipy are deliberately omitted here. This recipe requires numpy 1.x +# (OpenCV and diffusers break under numpy 2 on Jetson) while three other extras pin +# numpy==2.2.6, so rather than assert a resolution outcome we install them as a +# documented post-sync step; see recipes/internvla-n1-dualvln/requirements-torch.txt. +internvla-n1-dualvln = [ + "torch>=2.9.0", + "torchvision>=0.24.0", + "transformers==4.51.3", + "diffusers==0.33.1", + "accelerate==1.13.0", + "safetensors>=0.8.0", + "huggingface-hub>=0.36.0", + "nvidia-modelopt>=0.44.0", + "onnx==1.22.0", + "onnxscript==0.7.1", + "onnx-graphsurgeon==0.6.1", + "pandas>=2.0.0", + "pillow>=10.0.0", + "numpy-quaternion>=2023.0.0", + "pyyaml>=6.0.0", + "tqdm>=4.66.0", +] + [tool.ruff] target-version = "py312" diff --git a/recipes/internvla-n1-dualvln/.gitignore b/recipes/internvla-n1-dualvln/.gitignore new file mode 100644 index 0000000..292337d --- /dev/null +++ b/recipes/internvla-n1-dualvln/.gitignore @@ -0,0 +1,8 @@ +# Recipe artifacts. These are multi-gigabyte and reproducible; nothing in the root +# .gitignore covers them, so a stray WORK_DIR=. would otherwise stage ~40 GB. +work/ +*.engine +*.onnx +*.onnx.data +bench_imgs/ +reports/ diff --git a/recipes/internvla-n1-dualvln/Makefile b/recipes/internvla-n1-dualvln/Makefile new file mode 100644 index 0000000..98ffa98 --- /dev/null +++ b/recipes/internvla-n1-dualvln/Makefile @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +# Makefile — Top-level command panel for internvla-n1-dualvln +# +# Override paths via environment variables, e.g.: +# make quantize-fp8 INTERNVLA_CKPT=/data/InternVLA-N1-DualVLN WORK_DIR=/mnt/scratch +# +# Targets that run Python/bash assume the virtualenv is already activated. +# +# Targets are grouped by dependency boundary. Everything under "System 2" runs +# without InternNav; System 1 and the agent-level checks require INTERNNAV_PATH. + +# ── Configurable paths ───────────────────────────────────────────────────────── +INTERNVLA_CKPT ?= $(HOME)/InternNav/checkpoints/InternVLA-N1-DualVLN +INTERNNAV_PATH ?= $(HOME)/InternNav +TRT_EDGELLM_DIR ?= $(HOME)/modelopt/TensorRT-Edge-LLM +WORK_DIR ?= $(HOME)/vln-opt-work +ENGINE_DIR ?= $(WORK_DIR)/engines +CALIB_DATA_ROOT ?= $(WORK_DIR)/calib_scenes +SCHEME ?= fp8_default +STRATEGY ?= s1 +DEVICE ?= cuda +REPO_ROOT := $(abspath ../..) + +REPKG_CKPT := $(WORK_DIR)/qwen25vl_system2 +QUANT_CKPT := $(WORK_DIR)/qwen25vl_$(STRATEGY)_$(SCHEME) + +# ── System 2: quantize (no InternNav required) ──────────────────────────────── +.PHONY: fetch-calib repackage quantize quantize-fp8 quantize-nvfp4 + +fetch-calib: + bash quantize/scripts/00_fetch_calib_scenes.sh \ + --output_path $(CALIB_DATA_ROOT) + +repackage: + bash quantize/scripts/01_repackage.sh \ + --model_path $(INTERNVLA_CKPT) \ + --output_path $(REPKG_CKPT) + +quantize: + bash quantize/scripts/02_quantize.sh \ + --model_path $(REPKG_CKPT) \ + --output_path $(QUANT_CKPT) \ + --scheme $(SCHEME) \ + --strategy $(STRATEGY) \ + --device $(DEVICE) + +quantize-fp8: + $(MAKE) quantize SCHEME=fp8_default STRATEGY=s1 + +quantize-nvfp4: + $(MAKE) quantize SCHEME=nvfp4_default STRATEGY=s1 + +# ── System 2: export, build, verify (no InternNav required) ──────────────────── +.PHONY: export-build export-build-base-fp16 verify-latents benchmark + +export-build: + bash trt-edgellm/scripts/03_export_build_system2.sh \ + --model_path $(QUANT_CKPT) \ + --engine_dir $(ENGINE_DIR)/$(STRATEGY)_$(SCHEME) \ + --trt_edgellm_dir $(TRT_EDGELLM_DIR) + +export-build-base-fp16: + bash trt-edgellm/scripts/03_export_build_system2.sh \ + --model_path $(REPKG_CKPT) \ + --engine_dir $(ENGINE_DIR)/base_fp16 \ + --trt_edgellm_dir $(TRT_EDGELLM_DIR) \ + --no_quantization + +verify-latents: + bash trt-edgellm/scripts/05_verify.sh \ + --engine_dir $(ENGINE_DIR)/$(STRATEGY)_$(SCHEME) \ + --repkg_ckpt $(REPKG_CKPT) \ + --calib_data_root $(CALIB_DATA_ROOT) + +benchmark: + bash trt-edgellm/scripts/06_benchmark.sh \ + --engine_dir $(ENGINE_DIR)/$(STRATEGY)_$(SCHEME) \ + --work_dir $(WORK_DIR) + +# ── System 1: BF16 engines (requires INTERNNAV_PATH) ────────────────────────── +.PHONY: export-system1 verify-system1 + +export-system1: + bash trt-edgellm/scripts/04_export_system1.sh \ + --internvla_ckpt $(INTERNVLA_CKPT) \ + --internnav_path $(INTERNNAV_PATH) \ + --engine_dir $(ENGINE_DIR)/system1 + +verify-system1: + bash trt-edgellm/scripts/05_verify.sh \ + --system1 \ + --internvla_ckpt $(INTERNVLA_CKPT) \ + --internnav_path $(INTERNNAV_PATH) \ + --engine_dir $(ENGINE_DIR)/system1 + +# ── NVFP4 investigation ─────────────────────────────────────────────────────── +.PHONY: investigate-nvfp4 + +investigate-nvfp4: + python trt-edgellm/investigate_nvfp4.py \ + --repkg_ckpt $(REPKG_CKPT) \ + --calib_data_root $(CALIB_DATA_ROOT) \ + --work_dir $(WORK_DIR) + +# ── Housekeeping ────────────────────────────────────────────────────────────── +.PHONY: clean-onnx help + +# ONNX is a pure intermediate and the largest reclaimable artifact (9-16 GB). +clean-onnx: + rm -rf $(WORK_DIR)/onnx + @echo "Removed $(WORK_DIR)/onnx" + +help: + @echo "internvla-n1-dualvln — targets" + @echo "" + @echo " System 2 (no InternNav needed):" + @echo " fetch-calib download VLN calibration scenes" + @echo " repackage InternVLA checkpoint -> stock Qwen2.5-VL System 2" + @echo " quantize quantize with SCHEME/STRATEGY (default fp8_default/s1)" + @echo " quantize-fp8 shorthand for SCHEME=fp8_default STRATEGY=s1" + @echo " quantize-nvfp4 shorthand for SCHEME=nvfp4_default STRATEGY=s1 (experimental)" + @echo " export-build ONNX export + LLM and visual engines" + @echo " export-build-base-fp16 unquantized FP16 reference engine" + @echo " verify-latents acceptance gate: z_latents cosine vs FP32" + @echo " benchmark latency and memory" + @echo "" + @echo " System 1 (requires INTERNNAV_PATH):" + @echo " export-system1 traj_dit + memory block -> BF16 engines" + @echo " verify-system1 trajectory parity vs PyTorch" + @echo "" + @echo " Other:" + @echo " investigate-nvfp4 why NVFP4 breaks the System2 -> System1 bridge" + @echo " clean-onnx remove the ONNX intermediate (9-16 GB)" + @echo "" + @echo " Variables: INTERNVLA_CKPT INTERNNAV_PATH TRT_EDGELLM_DIR WORK_DIR" + @echo " ENGINE_DIR CALIB_DATA_ROOT SCHEME STRATEGY DEVICE" diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md new file mode 100644 index 0000000..066e948 --- /dev/null +++ b/recipes/internvla-n1-dualvln/README.md @@ -0,0 +1,259 @@ +# internvla-n1-dualvln + +FP8 quantization and TensorRT-Edge-LLM deployment for InternVLA-N1-DualVLN, a dual-system +vision-language navigation model, targeting NVIDIA Jetson Thor. + +## Status + +- Owner: unassigned +- Model family: `InternVLA-N1-DualVLN` (System 2 = Qwen2.5-VL-7B, System 1 = NextDiT trajectory head) +- Quantization presets: `fp8_default`, `fp8_per_channel` (validated) · `nvfp4_*` (experimental, see Notes) +- Runtime target: TensorRT-Edge-LLM engines on Jetson Thor (sm_110) + +## What this model is, and what the metric has to be + +InternVLA-N1 is a **dual-system** navigation policy: + + camera image ─► ViT (BF16/FP8 TRT) ─► LLM (FP16/FP8 TRT) ─► branch + ├─ discrete action (↑ ← → STOP) + └─ pixel goal (coordinate) + System 2 (planner, ~2 Hz) ▼ + ──────────────────────────── latent_queries ─► LLM hidden ─► norm ─► cond_projector + │ z_latents + System 1 (controller, ~15 Hz) ▼ + ──────────────────────────── RGB ─► memory block (BF16 TRT) ─┐ + ├─► traj_dit (BF16 TRT, x10) + z_latents ───┘ + +Only **System 2** is quantized. System 1 stays BF16. + +The two are joined by `z_latents`: the last-layer hidden states of 4 trajectory tokens, run +through a host-side `final_norm` and `cond_projector`. **That bridge is the acceptance +metric, not text quality.** A quantized checkpoint can emit perfectly fluent captions and +still be useless for navigation — NVFP4 does exactly that here (see Notes). Every claim in +this recipe is gated on `z_latents` cosine against the FP32 reference, not on generated text. + +## Results + +Measured on Jetson Thor, 12 held-out multi-image VLN steps: + +| LLM variant | z_latents vs FP32 | agrees w/ PyTorch | System 2 latency | LLM engine | +|---|---|---|---|---| +| PyTorch BF16 (baseline) | 0.99974 | — | 1631 ms · 1.00x | ~14 GB weights | +| base FP16 TensorRT (no quant) | **0.99985** | 12/12 | 770 ms · 2.12x | 14.2 GB | +| FP8 TensorRT | 0.99559 | 11/12 | **646 ms · 2.53x** | **7.6 GB** | +| NVFP4 TensorRT | **0.647** ✗ | tokens fine, bridge broken | — | 4.5 GB | + +Other engines: ViT 1.3 GB BF16 → 0.68 GB FP8; traj_dit 0.07 GB; memory block 0.11 GB. + +**Calibration data made no measurable difference.** Held-out z_latents came out at 0.99143 +with generic `cnn_dailymail` text versus 0.99146 with a domain-specific VLN set — equal +within noise. An earlier apparent gain turned out to be overlap between the calibration and +probe sets. The VLN calibration loader ships anyway (it is the honest default to offer for a +navigation model), but do not expect it to buy accuracy. + +## Files + + . + ├── README.md + ├── Makefile # command panel for both paths + ├── requirements.txt # PyPI dependencies (excluding PyTorch) + ├── requirements-torch.txt # PyTorch installation guide + ├── configs/ + │ └── schemes.yaml # scheme x strategy validity matrix + ├── quantize/ # HF checkpoint -> quantized HF checkpoint + │ ├── README.md + │ ├── repackage_system2.py # strip System 1 -> stock Qwen2.5-VL checkpoint + │ ├── quantize.py # ModelOpt driver + │ ├── configs.py # presets, strategies, validity gate + │ ├── calibration.py # text / multimodal / VLN calibration loaders + │ ├── model.py # load, calibrate, export + │ ├── prompt_builder.py # VLN prompt — single source of truth + │ └── scripts/{00_fetch_calib_scenes,01_repackage,02_quantize}.sh + └── trt-edgellm/ # quantized checkpoint -> engines -> verification + ├── README.md + ├── export_traj_dit.py # System 1 diffusion head -> ONNX -> BF16 engine + ├── export_memory_block.py # System 1 memory block -> ONNX -> BF16 engine + ├── engine_runner.py # direct-TensorRT LLM harness (hand-built 3D mRoPE) + ├── internvla_compat.py # patches needed to load System 1 + ├── investigate_nvfp4.py # why z_latents collapse under NVFP4 + ├── verify/ # 7 fidelity checks + ├── benchmark/ # 3 latency/memory benchmarks + ├── deploy/run_eval_engine.py + └── scripts/{03_export_build_system2,04_export_system1,05_verify,06_benchmark}.sh + +## The repackage step, and why it matters + +`InternVLA-N1-DualVLN` declares `model_type: internvla_n1`, ships **no** modeling code in the +checkpoint, and therefore cannot be loaded with `trust_remote_code` — the class has to come +from the InternNav repository. + +`quantize/repackage_system2.py` sidesteps that for the entire quantization flow. It is pure +safetensors surgery: it streams the checkpoint, drops the eight System-1 tensor prefixes, +and rewrites `config.json` to `model_type: qwen2_5_vl` / +`architectures: ["Qwen2_5_VLForConditionalGeneration"]`. **It imports nothing from +InternNav.** + +After that step everything downstream is a stock Qwen2.5-VL flow: + +| Step | Needs `INTERNNAV_PATH`? | +|---|---| +| `00_fetch_calib_scenes.sh` | no | +| `01_repackage.sh` | **no** — pure file surgery | +| `02_quantize.sh` | no — operates on stock Qwen2.5-VL | +| `03_export_build_system2.sh` | no | +| `04_export_system1.sh` | **yes** | +| `05_verify.sh` (latents) | no | +| `05_verify.sh` (agent-level) | **yes** | + +Cost of this approach is one ~15 GB intermediate checkpoint. Pass `--free_source` to delete +each input shard as its converted copy is written if disk is tight. + +## Setup + +**Step 1 — PyTorch.** On Jetson, use the JetPack wheel; do not install from PyPI. + +| Platform | Command | +|---|---| +| Jetson Thor (JetPack 7.1, CUDA 13.0) | use the JetPack-provided `torch==2.10.0`; see requirements-torch.txt | +| x86 CUDA 12.8 | `pip install torch==2.10.0+cu128 --extra-index-url https://download.pytorch.org/whl/cu128` | + +**Step 2 — remaining dependencies:** + + pip install -r requirements.txt + +**Step 2b — pin numpy afterwards.** This recipe needs numpy 1.x (OpenCV and diffusers break +under numpy 2 on Jetson), while the other recipes in this repository pin `numpy==2.2.6`. +Since `uv lock` resolves every extra into one universal lock, declaring both would make the +lock unsatisfiable, so numpy is deliberately absent from this recipe's extra. Run once after +`uv sync`: + + pip install "numpy==1.26.4" "scipy==1.13.1" + +**Step 3 — dependencies that are not on PyPI.** These must be present before running anything: + +| Dependency | How | +|---|---| +| TensorRT 10.13 | ships with JetPack at `/usr/lib/python3.12/dist-packages` | +| `tensorrt-edgellm` 0.8.0 | build from source, then `pip install --no-deps -e $TRT_EDGE_LLM` | +| InternNav | `git clone` it; export `INTERNNAV_PATH`. Only needed for System 1 and agent-level checks | +| OpenCV | system package | + +## Environment + + export INTERNVLA_CKPT=/path/to/InternVLA-N1-DualVLN # source checkpoint + export INTERNNAV_PATH=/path/to/InternNav # System 1 only + export TRT_EDGE_LLM=/path/to/TensorRT-Edge-LLM # build root + export VLN_OPT_WORK=$HOME/vln-opt-work # intermediate artifacts + export VLN_OPT_ENGINES=$VLN_OPT_WORK/engines # engine output + +Two flags are mandatory and set by the scripts themselves — listed here so their absence is +diagnosable, not so you set them by hand: + +- `TRITON_BACKENDS_IN_TREE=1` for every quantization run. +- `__LUNOWUD="-peep:fc_h_fusion=off"` for every engine build on TensorRT 10.13/10.14. Without + it the FP16 engine emits gibberish — a Myelin miscompile on sm_110, not a precision problem. + +## Quick Start + + make repackage # InternVLA checkpoint -> stock Qwen2.5-VL System 2 + make quantize-fp8 # s1 FP8, text calibration + make export-build # ONNX export + FP8 LLM engine + visual engine + make verify-latents # the acceptance gate: z_latents cosine > 0.99 + +`make help` lists every target. Each script is also runnable directly; see the per-path +READMEs in `quantize/` and `trt-edgellm/`. + +## CLI Reference + + quantize/quantize.py [options] + + --model_path PATH Repackaged System 2 checkpoint (required) + --output_path PATH Destination for the quantized checkpoint (required) + --strategy {s1,s2,s3,s4} s1 LLM · s2 +KV cache · s3 +ViT · s4 +ViT+KV (default: s1) + --scheme NAME Preset from configs/schemes.yaml (default: fp8_default) + --calib {auto,text,multimodal,vln} Calibration source (default: auto) + --calib_data PATH Root for VLN calibration scenes + --num_calib_samples N Calibration samples (default: 512; image paths cap at 128) + --max_seq_len N Calibration truncation length (default: 512) + --dtype {fp16,bf16} Load dtype (default: bf16) + --device DEVICE Torch device (default: cuda) + --resume DIR Layerwise checkpoint dir, for AWQ/Hessian crash recovery + --allow_experimental Permit schemes marked experimental (NVFP4) + --dry_run Load and validate, skip quantization + + quantize/repackage_system2.py [options] + + --model_path PATH Source InternVLA-N1-DualVLN checkpoint (required) + --output_path PATH Destination stock Qwen2.5-VL checkpoint (required) + --free_source Delete each source shard once converted (destructive) + +## Notes + +### What is and is not quantized + +Quantized: the System 2 LLM backbone (`model.layers.*`), and the vision tower under `s3`/`s4`. +Never quantized: System 1 (`traj_dit`, memory block), `lm_head`, and the host-side bridge ops +(`final_norm`, `cond_projector`). + +### Scheme validity + +| | s1 (LLM) | s2 (+KV) | s3 (+ViT) | s4 (+ViT+KV) | +|---|---|---|---|---| +| `fp8_default`, `fp8_per_channel` | yes | yes | yes | yes | +| `nvfp4_*` | experimental | experimental | **blocked** | **blocked** | + +- **NVFP4 with the vision tower is blocked, permanently.** The Qwen2.5-VL ViT MLP has + `intermediate_size = 3420`, and 3420 / 16 = 213.75 — not divisible by the NVFP4 block size. + The recipe rejects these combinations up front rather than letting them crash mid-run. +- **KV cache is always FP8**, even with NVFP4 weights. NVFP4 KV requires `sm100f` (datacenter + Blackwell); Thor is sm110. This is a hardware limit, not a configuration mistake — do not + "fix" it by selecting an NVFP4 KV preset. +- **NVFP4 weights are experimental and currently not usable for navigation.** They quantize, + export and generate fluent text, but `z_latents` cosine falls to 0.647, which breaks the + System 2 → System 1 bridge. Requires `--allow_experimental`. See + `trt-edgellm/investigate_nvfp4.py`. + +### FP16 on Thor is fine — an earlier claim to the contrary was wrong + +Earlier notes in the source project stated that a plain FP16 LLM engine "produces garbage on +Thor, so FP8 is mandatory". **That is incorrect and has been retracted.** The garbage came +from a Myelin `fc_h_fusion` miscompile on sm_110 at TensorRT 10.13: the horizontal fusion of +the gate/up projections is wrong at batch 1. TensorRT-Edge-LLM already disables it, but only +for TensorRT >= 10.15, so 10.13 slips through the gap. FP8 quietly dodged the same bug +because its Q/DQ nodes break the fusion pattern — which is why FP8 *looked* mandatory. + +Exporting `__LUNOWUD="-peep:fc_h_fusion=off"` fixes FP16 completely; the base FP16 engine is +the highest-fidelity variant in the results table. FP8 remains the recommended deployment +choice on size and latency, not on correctness. + +### One visual engine, sized for multi-image prompts + +The visual encoder is built once at `--minImageTokens 4 --maxImageTokens 4096 +--maxImageTokensPerImage 1024`, and every consumer reads that one engine. + +This is worth stating because the source project got it wrong in a way that is easy to +repeat: its default build used `128 / 512 / 512`, taken from a single-image demo. A VLN +prompt carries 9–10 images and roughly 1,764 image tokens, so it does not fit in 512, and the +multi-image verification scripts were quietly pointed at a hand-built engine that no script +in the repository produced. If a verification here reports a shape or capacity error, check +the visual engine's sizing before suspecting the weights. + +### Scope + +These are **conversion-fidelity** numbers plus offline planner metrics. They are not a +closed-loop navigation success rate — that needs the Habitat/InternUtopia simulator, which is +not part of this recipe. `trt-edgellm/verify/verify_engine_policy.py` checks that the +simulator adapter is wired correctly, but cannot exercise it. + +## Tested Environments + +- **OS:** Ubuntu 24.04 (JetPack 7.1) +- **Hardware:** NVIDIA Jetson Thor (Blackwell, sm_110, 128 GB unified memory) +- **Python:** 3.12.3 +- **PyTorch:** 2.10.0 +- **CUDA:** 13.0 +- **TensorRT:** 10.13.3.9 +- **tensorrt-edgellm:** 0.8.0 +- **nvidia-modelopt:** 0.44.0 +- **transformers:** 4.51.3 · **diffusers:** 0.33.1 · **onnx:** 1.22.0 · **numpy:** 1.26.4 diff --git a/recipes/internvla-n1-dualvln/configs/schemes.yaml b/recipes/internvla-n1-dualvln/configs/schemes.yaml new file mode 100644 index 0000000..42fb6c6 --- /dev/null +++ b/recipes/internvla-n1-dualvln/configs/schemes.yaml @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +# +# Quantization schemes and the scheme x strategy validity matrix for +# InternVLA-N1-DualVLN on Jetson Thor (sm_110). +# +# Two of the constraints below are hardware facts, not tuning choices. Encoding them +# here means an impossible combination is rejected before the model is loaded, instead +# of crashing forty minutes into a run: +# +# * NVFP4 cannot touch the Qwen2.5-VL vision tower. Its MLP intermediate_size is +# 3420, and 3420 / 16 = 213.75 -- not divisible by the NVFP4 block size. +# * NVFP4 KV cache requires sm100f (datacenter Blackwell). Thor is sm110, so the KV +# cache is FP8 even when the weights are NVFP4. This is deliberate; do not +# "fix" it by pointing at an NVFP4 KV preset. + +# --------------------------------------------------------------------------- # +# Strategies -- which parts of System 2 get quantized +# --------------------------------------------------------------------------- # +strategies: + s1: + description: LLM backbone only -- fastest, safest, the validated default + quantize_kv_cache: false + quantize_visual: false + s2: + description: LLM + KV cache -- lower memory for long multi-image prompts + quantize_kv_cache: true + quantize_visual: false + s3: + description: LLM + vision tower -- full VLM quantization + quantize_kv_cache: false + quantize_visual: true + s4: + description: LLM + vision tower + KV cache -- maximum compression + quantize_kv_cache: true + quantize_visual: true + +# --------------------------------------------------------------------------- # +# Schemes -- weight/activation formats, mapped to ModelOpt presets +# +# calib_batch_size applies to text calibration only; image calibration is always 1. +# --------------------------------------------------------------------------- # +schemes: + fp8_default: + modelopt_preset: FP8_DEFAULT_CFG + description: FP8 W8A8 per-tensor -- balanced accuracy and speed + calib_batch_size: 8 + status: validated + fp8_per_channel: + modelopt_preset: FP8_PER_CHANNEL_PER_TOKEN_CFG + description: FP8 W8A8 per-channel weight, per-token activation -- highest FP8 accuracy + calib_batch_size: 8 + status: validated + nvfp4_default: + modelopt_preset: NVFP4_DEFAULT_CFG + description: NVFP4 W4A4 -- smallest footprint + calib_batch_size: 4 + status: experimental + nvfp4_awq_full: + modelopt_preset: NVFP4_AWQ_FULL_CFG + description: NVFP4 W4A4 + AWQ (lite + clip) -- best NVFP4 accuracy + calib_batch_size: 1 + status: experimental + nvfp4_local_hessian: + modelopt_preset: NVFP4_W4A4_WEIGHT_LOCAL_HESSIAN_CFG + description: NVFP4 W4A4 + local Hessian -- highest-accuracy NVFP4, slow + calib_batch_size: 1 + status: experimental + +# --------------------------------------------------------------------------- # +# KV cache pairing +# --------------------------------------------------------------------------- # +kv_cache: + preset: FP8_KV_CFG + applies_to: [fp8_default, fp8_per_channel, nvfp4_default, nvfp4_awq_full, nvfp4_local_hessian] + note: >- + FP8 KV for every weight format. NVFP4 KV needs sm100f at inference time, which Thor + (sm110) does not provide; FP8 KV pairs portably with NVFP4 weights. + +# --------------------------------------------------------------------------- # +# Modules excluded when a strategy leaves the vision tower alone. +# +# ModelOpt presets exclude *lm_head* by default but NOT the vision tower, so without +# these globs s1/s2 would silently quantize the ViT. +# --------------------------------------------------------------------------- # +visual_exclude_patterns: + - "*visual*" + - "*vision_tower*" + - "*multi_modal_projector*" + +# --------------------------------------------------------------------------- # +# Validity matrix -- every combination not listed as blocked or experimental is allowed +# --------------------------------------------------------------------------- # +blocked: + - schemes: [nvfp4_default, nvfp4_awq_full, nvfp4_local_hessian] + strategies: [s3, s4] + reason: >- + NVFP4 cannot quantize the Qwen2.5-VL vision tower: its MLP intermediate_size is + 3420 and 3420 / 16 = 213.75, which is not an integer, so the tensor cannot be + split into NVFP4 blocks. Use an fp8_* scheme for s3/s4. + +experimental: + - schemes: [nvfp4_default, nvfp4_awq_full, nvfp4_local_hessian] + strategies: [s1, s2] + reason: >- + NVFP4 weights quantize, export and generate fluent text, but the System 2 -> System 1 + bridge breaks: z_latents cosine against the FP32 reference falls to 0.647 versus + 0.9956 for FP8. The model will caption an image correctly and still navigate wrongly. + Pass --allow_experimental to proceed. See trt-edgellm/investigate_nvfp4.py. diff --git a/recipes/internvla-n1-dualvln/requirements-torch.txt b/recipes/internvla-n1-dualvln/requirements-torch.txt new file mode 100644 index 0000000..aff3bb1 --- /dev/null +++ b/recipes/internvla-n1-dualvln/requirements-torch.txt @@ -0,0 +1,31 @@ +# Install PyTorch BEFORE running: pip install -r requirements.txt +# Pick the line matching your platform: +# +# Jetson Thor (JetPack 7.1, CUDA 13.0) — this is the tested configuration: +# Use the JetPack-provided wheel. Do NOT install torch from PyPI on Jetson; +# the PyPI aarch64 wheels are not built for sm_110 and will fail at runtime. +# Verify with: +# python -c "import torch; print(torch.__version__, torch.cuda.get_arch_list())" +# Expect torch 2.10.0 and 'sm_110' present in the arch list. +# +# x86 CUDA 12.8: +# pip install torch==2.10.0+cu128 torchvision==0.25.0+cu128 --extra-index-url https://download.pytorch.org/whl/cu128 +# +# x86 CUDA 13.x: +# pip install torch==2.10.0+cu130 torchvision==0.25.0+cu130 --extra-index-url https://download.pytorch.org/whl/cu130 +# +# ----------------------------------------------------------------------------- +# numpy must be pinned AFTER `uv sync`, not by it. +# +# This recipe needs numpy 1.x: OpenCV and diffusers on Jetson break under numpy 2, +# while three other extras in this repository pin numpy==2.2.6. Rather than declare a +# pin whose resolution outcome we have not verified, numpy and scipy are left out of +# this recipe's pyproject extra and installed explicitly below. +# +# Run this once after `uv sync --extra internvla-n1-dualvln`: +# +# pip install "numpy==1.26.4" "scipy==1.13.1" +# +# Installing `datasets` later can silently upgrade numpy again; re-run the line above +# if it does. +# ----------------------------------------------------------------------------- diff --git a/recipes/internvla-n1-dualvln/requirements.txt b/recipes/internvla-n1-dualvln/requirements.txt new file mode 100644 index 0000000..a77cb51 --- /dev/null +++ b/recipes/internvla-n1-dualvln/requirements.txt @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +# +# InternVLA-N1-DualVLN — Python dependencies (Jetson Thor / aarch64, Python 3.12). +# PyTorch is intentionally absent; see requirements-torch.txt. +# +# pip install -r requirements.txt + +# Hugging Face Ecosystem +transformers==4.51.3 +accelerate==1.13.0 +safetensors==0.8.0 +huggingface-hub>=0.23.0 + +# Diffusion (System 1 trajectory head) +diffusers==0.33.1 + +# Numerics — pinned below numpy 2 so OpenCV and diffusers keep working on Jetson +numpy==1.26.4 +scipy==1.13.1 +pandas>=2.0.0 +pillow>=10.0.0 + +# ONNX export +onnx==1.22.0 +onnxscript==0.7.1 +onnx-graphsurgeon==0.6.1 + +# Quantization +nvidia-modelopt==0.44.0 + +# Utilities +numpy-quaternion>=2023.0.0 +pyyaml>=6.0.0 +tqdm>=4.66.0 + +# Calibration dataset loader — installing this may pull numpy>=2; if it does, +# re-pin numpy==1.26.4 and scipy==1.13.1 afterwards. +# datasets==2.19.0 + +# ============================================================================= +# Not available on PyPI for aarch64 — install these separately: +# +# TensorRT 10.13 provided by JetPack, imported from /usr/lib/python3.12/dist-packages +# tensorrt-edgellm build from source, then: pip install --no-deps -e $TRT_EDGE_LLM +# OpenCV (cv2) system package on Jetson +# InternNav git clone; export INTERNNAV_PATH. Needed only for System 1 +# export and the agent-level verifications. +# flash-attn optional, only for the PyTorch reference baseline +# ============================================================================= From ea12e217003542c894d078c5a97091f27996cf00 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Thu, 13 Aug 2026 15:55:13 +0700 Subject: [PATCH 02/30] feat(internvla-n1-dualvln): add repackage step and scheme validity gate repackage_system2.py strips the eight System-1 tensor prefixes and rewrites config.json to a stock Qwen2.5-VL. Verified on the real 16 GB InternVLA-N1-DualVLN checkpoint: 729 System-2 keys kept, 609 System-1 keys dropped, and the result loads through AutoConfig and AutoProcessor as Qwen2_5_VLConfig / Qwen2_5_VLProcessor with no InternVLA fields left. That is what takes InternNav off the hot path for everything except System-1 export. Five bridge tensors (latent_queries, cond_projector.{0,2}.{weight,bias}) are written to a separate bridge.safetensors instead of being dropped with the rest of System 1. They are what the z_latents fidelity check consumes, so keeping them here lets quantize, export, build and verify all run from the repackaged directory without reopening the 16 GB source. quant_schemes.py loads configs/schemes.yaml and rejects impossible combinations before the model is loaded, so a bad request costs seconds instead of dying mid-quantization. The vision-tower check reads vision_config.intermediate_size from the checkpoint rather than hardcoding 3420, so it stays correct for other Qwen2.5-VL sizes; on this checkpoint it reports '3420 / 16 = 213.75' in the error. Verified against all four strategies and both scheme families: fp8 passes everywhere, nvfp4 is blocked on s3/s4 and gated behind --allow_experimental on s1/s2. A free-space preflight refuses to start below 18 GB, and --free_source deletes each source shard as its converted copy lands for machines that cannot hold both checkpoints. (cherry picked from commit 0f828a726cce5168deb864fe23160a26c8a4b97f) --- .../quantize/quant_schemes.py | 224 ++++++++++++++++++ .../quantize/repackage_system2.py | 213 +++++++++++++++++ 2 files changed, 437 insertions(+) create mode 100644 recipes/internvla-n1-dualvln/quantize/quant_schemes.py create mode 100644 recipes/internvla-n1-dualvln/quantize/repackage_system2.py diff --git a/recipes/internvla-n1-dualvln/quantize/quant_schemes.py b/recipes/internvla-n1-dualvln/quantize/quant_schemes.py new file mode 100644 index 0000000..1e04a7a --- /dev/null +++ b/recipes/internvla-n1-dualvln/quantize/quant_schemes.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Quantization scheme registry, validity gate, and ModelOpt config composer. + +Schemes and strategies live in ``configs/schemes.yaml`` rather than in Python so the +matrix is readable without tracing code, following the ``qwen36-27b`` recipe's pattern. + +The validity gate exists because two combinations are impossible on this hardware and +neither fails early on its own: + +* NVFP4 cannot quantize the Qwen2.5-VL vision tower. The ViT MLP ``intermediate_size`` + is 3420 and the NVFP4 block size is 16; 3420 / 16 = 213.75. Without a gate this only + surfaces once quantization is already under way. +* NVFP4 KV cache requires ``sm100f`` (datacenter Blackwell). Jetson Thor is sm110, so the + KV cache is FP8 even when the weights are NVFP4. + +The divisibility check reads ``vision_config.intermediate_size`` from the checkpoint being +quantized rather than hardcoding 3420, so it stays correct for other Qwen2.5-VL sizes. +""" +import copy +import json +import os +from typing import Any, Optional + +import yaml + +_DEFAULT_SCHEMES_YAML = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "configs", "schemes.yaml" +) + +# NVFP4 packs weights in blocks of this many elements along the reduction axis. +NVFP4_BLOCK_SIZE = 16 + + +def load_registry(path: Optional[str] = None) -> dict: + """Load ``configs/schemes.yaml``.""" + path = path or _DEFAULT_SCHEMES_YAML + if not os.path.isfile(path): + raise FileNotFoundError(f"scheme registry not found: {path}") + with open(path) as f: + return yaml.safe_load(f) + + +def scheme_names(registry: Optional[dict] = None) -> list[str]: + return sorted((registry or load_registry())["schemes"]) + + +def strategy_names(registry: Optional[dict] = None) -> list[str]: + return sorted((registry or load_registry())["strategies"]) + + +def is_nvfp4(scheme: str) -> bool: + return scheme.startswith("nvfp4") + + +def _vision_intermediate_size(model_path: str) -> Optional[int]: + """Read ``vision_config.intermediate_size`` from a checkpoint, if it has one.""" + cfg_path = os.path.join(model_path, "config.json") + if not os.path.isfile(cfg_path): + return None + with open(cfg_path) as f: + cfg = json.load(f) + vision = cfg.get("vision_config") + if isinstance(vision, dict): + return vision.get("intermediate_size") + return None + + +def validate(scheme: str, strategy: str, model_path: Optional[str] = None, + allow_experimental: bool = False, + registry: Optional[dict] = None) -> None: + """Raise ``ValueError`` if this combination cannot or should not run. + + Called before the model is loaded, so a rejected combination costs seconds rather + than a full checkpoint load followed by a mid-quantization crash. + """ + registry = registry or load_registry() + + if scheme not in registry["schemes"]: + raise ValueError(f"unknown scheme {scheme!r}; available: {scheme_names(registry)}") + if strategy not in registry["strategies"]: + raise ValueError(f"unknown strategy {strategy!r}; available: {strategy_names(registry)}") + + strat = registry["strategies"][strategy] + + for rule in registry.get("blocked", []): + if scheme in rule["schemes"] and strategy in rule["strategies"]: + detail = "" + # Prefer the checkpoint's own number over the one in the message. + if strat["quantize_visual"] and model_path: + size = _vision_intermediate_size(model_path) + if size is not None: + detail = (f" This checkpoint's vision_config.intermediate_size is {size}; " + f"{size} / {NVFP4_BLOCK_SIZE} = {size / NVFP4_BLOCK_SIZE}.") + raise ValueError(f"{scheme} + {strategy} is not supported. " + f"{rule['reason'].strip()}{detail}") + + for rule in registry.get("experimental", []): + if scheme in rule["schemes"] and strategy in rule["strategies"]: + if not allow_experimental: + raise ValueError(f"{scheme} + {strategy} is experimental. " + f"{rule['reason'].strip()}") + print(f"[WARN] {scheme} + {strategy} is experimental. {rule['reason'].strip()}") + + +def build_quant_config(scheme: str, strategy: str, + layerwise_checkpoint_dir: Optional[str] = None, + registry: Optional[dict] = None) -> dict: + """Compose a ModelOpt quant_cfg from a scheme preset plus a strategy. + + Mirrors NVIDIA's ``build_quant_config`` pattern: start from a base preset, then merge + in the KV-cache entries and append visual exclusions as the strategy dictates. + """ + import modelopt.torch.quantization as mtq + + registry = registry or load_registry() + scheme_cfg = registry["schemes"][scheme] + strat = registry["strategies"][strategy] + + preset_name = scheme_cfg["modelopt_preset"] + if not hasattr(mtq, preset_name): + raise ValueError(f"modelopt has no preset {preset_name!r} " + f"(scheme {scheme!r}); check the installed nvidia-modelopt version") + quant_cfg = copy.deepcopy(getattr(mtq, preset_name)) + + if strat["quantize_kv_cache"]: + kv_name = registry["kv_cache"]["preset"] + if scheme not in registry["kv_cache"]["applies_to"]: + raise ValueError(f"no KV-cache preset registered for scheme {scheme!r}") + # FP8 KV for every weight format; see the note in schemes.yaml. + kv_cfg = getattr(mtq, kv_name) + quant_cfg["quant_cfg"] = quant_cfg["quant_cfg"] + kv_cfg["quant_cfg"] + + if not strat["quantize_visual"]: + # ModelOpt presets exclude *lm_head* by default but not the vision tower, so + # without these the ViT would be quantized on s1/s2 without anyone asking. + for pattern in registry["visual_exclude_patterns"]: + quant_cfg["quant_cfg"].append({"quantizer_name": pattern, "enable": False}) + + if layerwise_checkpoint_dir is not None: + algo: Any = quant_cfg.get("algorithm") + if isinstance(algo, str): + algo = {"method": algo} + elif algo is None: + algo = {} + elif isinstance(algo, dict): + algo = dict(algo) + else: + raise TypeError(f"unexpected algorithm type: {type(algo)}") + algo["layerwise"] = {"enable": True, "checkpoint_dir": layerwise_checkpoint_dir} + quant_cfg["algorithm"] = algo + + return quant_cfg + + +def calib_batch_size(scheme: str, is_image_calib: bool, + registry: Optional[dict] = None) -> int: + """Calibration batch size. Image calibration is always 1 (GPU-memory bound).""" + if is_image_calib: + return 1 + registry = registry or load_registry() + return int(registry["schemes"][scheme].get("calib_batch_size", 1)) + + +def render_matrix(registry: Optional[dict] = None) -> str: + """Render the scheme x strategy matrix for pasting into a README.""" + registry = registry or load_registry() + strategies = strategy_names(registry) + blocked = {(s, st) for r in registry.get("blocked", []) + for s in r["schemes"] for st in r["strategies"]} + experimental = {(s, st) for r in registry.get("experimental", []) + for s in r["schemes"] for st in r["strategies"]} + + lines = ["| scheme | " + " | ".join(strategies) + " |", + "|---" * (len(strategies) + 1) + "|"] + for scheme in scheme_names(registry): + cells = [] + for strategy in strategies: + if (scheme, strategy) in blocked: + cells.append("blocked") + elif (scheme, strategy) in experimental: + cells.append("experimental") + else: + cells.append("yes") + lines.append(f"| `{scheme}` | " + " | ".join(cells) + " |") + return "\n".join(lines) + + +def describe() -> str: + """Human-readable listing of schemes and strategies, for --help epilogs.""" + registry = load_registry() + out = ["schemes:"] + for name in scheme_names(registry): + cfg = registry["schemes"][name] + out.append(f" {name:22s} {cfg['description']} [{cfg['status']}]") + out.append("strategies:") + for name in strategy_names(registry): + out.append(f" {name:22s} {registry['strategies'][name]['description']}") + return "\n".join(out) + + +def main() -> int: + import argparse + + parser = argparse.ArgumentParser(description="Inspect the quantization scheme registry") + parser.add_argument("--list_scheme_names", action="store_true") + parser.add_argument("--list_strategy_names", action="store_true") + parser.add_argument("--print_matrix", action="store_true") + parser.add_argument("--describe", action="store_true") + args = parser.parse_args() + + if args.list_scheme_names: + print("\n".join(scheme_names())) + elif args.list_strategy_names: + print("\n".join(strategy_names())) + elif args.print_matrix: + print(render_matrix()) + else: + print(describe()) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/quantize/repackage_system2.py b/recipes/internvla-n1-dualvln/quantize/repackage_system2.py new file mode 100644 index 0000000..e7e9c5a --- /dev/null +++ b/recipes/internvla-n1-dualvln/quantize/repackage_system2.py @@ -0,0 +1,213 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Repackage the InternVLA-N1 System 2 into a standalone Qwen2.5-VL checkpoint. + +InternVLA-N1-DualVLN declares ``model_type: internvla_n1`` and ships no modeling code, so +it cannot be loaded with ``trust_remote_code`` -- the class has to come from the InternNav +repository. That dependency is avoidable for everything except System 1 itself, because the +System 2 backbone (vision tower + Qwen2.5 LLM) inside the checkpoint already uses standard +Qwen2.5-VL key names. + +This is a lossless subset copy, streamed through ``safe_open`` so peak memory is one shard: + + keep ``visual.*``, ``model.embed_tokens``, ``model.norm``, ``model.layers.*``, ``lm_head`` + drop the System 1 modules (traj_dit, rgb_model, rgb_resampler, memory_encoder, + cond_projector, action_encoder, action_decoder, pos_encoding) + rewrite ``config.json`` to ``model_type=qwen2_5_vl`` / + ``architectures=[Qwen2_5_VLForConditionalGeneration]`` + +Weights are copied bit-for-bit; the source is opened read-only. + +Two System 1 tensors are *not* simply dropped. ``latent_queries`` and ``cond_projector`` form +the System 2 -> System 1 bridge, and the fidelity checks need them to compute z_latents. They +are written to a separate ``bridge.safetensors`` (~25 MB) so that quantization, export, engine +build and verification can all run from this directory alone, without reopening the 16 GB +source checkpoint and without importing InternNav. +""" +import argparse +import json +import os +import shutil + +from safetensors import safe_open +from safetensors.torch import save_file + +SYSTEM1_PREFIXES = ( + "model.traj_dit", + "model.rgb_model", + "model.rgb_resampler", + "model.memory_encoder", + "model.cond_projector", + "model.action_encoder", + "model.action_decoder", + "model.pos_encoding", +) + +# Tensors that belong to System 1 but are needed to evaluate the System 2 -> System 1 +# bridge. Kept aside rather than dropped; see the module docstring. +BRIDGE_KEYS = ("model.latent_queries", "model.cond_projector") + +# Non-weight files carried over verbatim. Matching by prefix rather than an exact list +# keeps this working when a checkpoint ships an extra tokenizer artifact. +COPY_PREFIXES = ( + "tokenizer", + "vocab", + "merges", + "preprocessor", + "chat_template", + "generation_config", + "added_tokens", + "special_tokens", +) + +# A repackaged 7B System 2 is ~15.5 GB. Refuse to start rather than die at 90 %. +MIN_FREE_GB = 18 + + +def is_system2(key: str) -> bool: + """Keep the standard vision + LLM keys; drop every System 1 module.""" + if key.startswith(SYSTEM1_PREFIXES): + return False + return ( + key.startswith("visual.") + or key == "lm_head.weight" + or key.startswith("model.embed_tokens") + or key.startswith("model.norm") + or key.startswith("model.layers.") + ) + + +def is_bridge(key: str) -> bool: + """Tensors kept aside for z_latents evaluation.""" + return key.startswith(BRIDGE_KEYS) + + +def build_config(src: str) -> dict: + """Rewrite the InternVLA config into a plain Qwen2.5-VL one.""" + with open(os.path.join(src, "config.json")) as f: + cfg = json.load(f) + for key in ("system1", "n_query", "model_cfg", "model_type", "architectures", "auto_map"): + cfg.pop(key, None) + cfg["model_type"] = "qwen2_5_vl" + cfg["architectures"] = ["Qwen2_5_VLForConditionalGeneration"] + return cfg + + +def free_gb(path: str) -> float: + stat = os.statvfs(path) + return stat.f_bavail * stat.f_frsize / 1e9 + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--model_path", required=True, + help="Source InternVLA-N1-DualVLN checkpoint directory") + parser.add_argument("--output_path", required=True, + help="Destination directory for the stock Qwen2.5-VL System 2 checkpoint") + parser.add_argument("--free_source", action="store_true", + help="Delete each source shard once its converted copy is written. " + "Peak disk becomes one shard rather than two checkpoints. " + "Destructive -- only use this if the source can be re-downloaded.") + parser.add_argument("--skip_disk_check", action="store_true", + help="Skip the free-space preflight") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + src, dst = args.model_path, args.output_path + + if not os.path.isdir(src): + print(f"[ERROR] --model_path does not exist: {src}") + return 1 + index_path = os.path.join(src, "model.safetensors.index.json") + if not os.path.isfile(index_path): + print(f"[ERROR] no model.safetensors.index.json under {src}; " + f"a sharded checkpoint is expected") + return 1 + + os.makedirs(dst, exist_ok=True) + + if not args.skip_disk_check and not args.free_source: + avail = free_gb(dst) + if avail < MIN_FREE_GB: + print(f"[ERROR] only {avail:.1f} GB free under {dst}; need >= {MIN_FREE_GB} GB. " + f"Free space, point --output_path elsewhere, or pass --free_source to " + f"delete each source shard as it is converted.") + return 1 + + with open(index_path) as f: + weight_map = json.load(f)["weight_map"] + + kept = [k for k in weight_map if is_system2(k)] + bridge = [k for k in weight_map if is_bridge(k)] + dropped = [k for k in weight_map if not is_system2(k)] + print(f"keep {len(kept)} System2 keys, drop {len(dropped)} System1 keys " + f"({len(bridge)} of them kept aside as bridge tensors)") + + # Group by source shard so each is opened once and read sequentially. + by_shard: dict[str, list[str]] = {} + for key in kept: + by_shard.setdefault(weight_map[key], []).append(key) + bridge_by_shard: dict[str, list[str]] = {} + for key in bridge: + bridge_by_shard.setdefault(weight_map[key], []).append(key) + + new_weight_map: dict[str, str] = {} + bridge_tensors: dict[str, "object"] = {} + total_bytes = 0 + out_shards = sorted(by_shard) + n_shards = len(out_shards) + + for i, shard in enumerate(out_shards, 1): + out_name = f"model-{i:05d}-of-{n_shards:05d}.safetensors" + tensors = {} + with safe_open(os.path.join(src, shard), framework="pt") as f: + for key in by_shard[shard]: + tensor = f.get_tensor(key) + tensors[key] = tensor + total_bytes += tensor.numel() * tensor.element_size() + new_weight_map[key] = out_name + for key in bridge_by_shard.get(shard, []): + bridge_tensors[key] = f.get_tensor(key) + save_file(tensors, os.path.join(dst, out_name), metadata={"format": "pt"}) + print(f" [{i}/{n_shards}] {out_name}: {len(tensors)} tensors") + del tensors + if args.free_source: + os.remove(os.path.join(src, shard)) + print(f" removed source shard {shard}") + + with open(os.path.join(dst, "model.safetensors.index.json"), "w") as f: + json.dump({"metadata": {"total_size": total_bytes}, "weight_map": new_weight_map}, + f, indent=2) + + with open(os.path.join(dst, "config.json"), "w") as f: + json.dump(build_config(src), f, indent=2) + + if bridge_tensors: + save_file(bridge_tensors, os.path.join(dst, "bridge.safetensors"), + metadata={"format": "pt"}) + print(f" bridge.safetensors: {len(bridge_tensors)} tensors " + f"({', '.join(sorted(bridge_tensors))})") + else: + print("[WARN] no bridge tensors found; z_latents verification will need the " + "original checkpoint") + + for name in os.listdir(src): + if name.endswith(".safetensors"): + continue + if any(name.startswith(p) for p in COPY_PREFIXES): + shutil.copy2(os.path.join(src, name), os.path.join(dst, name)) + + dropped_modules = sorted({k.split(".")[1] for k in dropped if "." in k}) + print(f"\nDone -> {dst}") + print(f" {n_shards} shards, {total_bytes / 1e9:.1f} GB System 2") + print(" config model_type=qwen2_5_vl") + print(f" dropped System 1 modules: {dropped_modules}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 565de9933bb1f09fd6d4730252735a823fb664a8 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Thu, 13 Aug 2026 16:04:13 +0700 Subject: [PATCH 03/30] feat(internvla-n1-dualvln): add quantize driver, calibration and model loader Completes the System-2 quantization path. Verified end to end on the real checkpoint: 16.6 GB repackaged System 2 -> 9.4 GB FP8, 1164 TensorQuantizers, 80s quantize plus 16s export, hf_quant_config reporting FP8 with lm_head and visual* excluded, and tokenizer.model_max_length still 8192 (calibration truncates at 512 as a call-time argument only; a leak there would silently cap inference sequence length). model_loader.py drops the InternVLA-N1 branch entirely rather than carrying it forward. Post-repackage the input is always a stock Qwen2.5-VL, so the branch, its two InternNav imports, the diffusers gradient-checkpointing monkeypatch and the unused _is_system2_checkpoint helper are all dead weight -- 70 lines removed, and with them the dependency that would otherwise sit on the critical path of every quantization run. quantize.py calls the validity gate before loading the model, so an impossible request costs seconds. Confirmed: nvfp4+s3 is rejected in under a second with the checkpoint's own '3420 / 16 = 213.75' and never touches the weights, while fp8+s1 loads in 14s and reports sm_110. --cfg is renamed --scheme, kebab-case flags become snake_case to match the other recipes, and the output tag no longer hardcodes qwen2.5-vl-7b. calibration.py gains an offline path. This machine sits behind a TLS-intercepting gateway, so load_dataset() fails certificate verification even though the parquet shards are already in the hub cache -- and HF's own offline mode does not help, because it still wants the dataset script. The loader now probes reachability and reads the cached shards directly, which is the difference between a working and a non-working calibration run on any air-gapped or proxied deployment machine, not just this one. (cherry picked from commit 260bf45ea3a418262ab65b0c67a8b676f48f8df6) --- .../quantize/calibration.py | 421 ++++++++++++++++++ .../quantize/model_loader.py | 291 ++++++++++++ .../quantize/prompt_builder.py | 237 ++++++++++ .../internvla-n1-dualvln/quantize/quantize.py | 323 ++++++++++++++ 4 files changed, 1272 insertions(+) create mode 100644 recipes/internvla-n1-dualvln/quantize/calibration.py create mode 100644 recipes/internvla-n1-dualvln/quantize/model_loader.py create mode 100644 recipes/internvla-n1-dualvln/quantize/prompt_builder.py create mode 100755 recipes/internvla-n1-dualvln/quantize/quantize.py diff --git a/recipes/internvla-n1-dualvln/quantize/calibration.py b/recipes/internvla-n1-dualvln/quantize/calibration.py new file mode 100644 index 0000000..f703653 --- /dev/null +++ b/recipes/internvla-n1-dualvln/quantize/calibration.py @@ -0,0 +1,421 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Calibration dataloaders for Qwen2.5-VL quantization. + +Three dataloader paths: + * ``text_calib_dataloader`` — ``cnn_dailymail`` ``input_ids`` only, used when + quantizing LLM backbone alone (no visual quantization). + * ``multimodal_calib_dataloader`` — ``lmms-lab/MMMU`` image+text pairs streamed + through the model's own ``AutoProcessor`` chat template, used when visual + tower quantization is enabled. + * ``vln_calib_dataloader`` — the in-distribution set: real InternData-N1 VLN-CE + episodes (the data System 2 was fine-tuned on) assembled through the deployed + agent's own prompt path, so ModelOpt sees the true activation ranges of the + farthest-pixel-goal navigation task instead of out-of-distribution web text/QA. + +The first two mirror NVIDIA's reference implementation. +""" + +import os +from typing import Any, Optional + +import torch +from torch.utils.data import DataLoader + + +# --------------------------------------------------------------------------- # +# Offline fallbacks +# --------------------------------------------------------------------------- # +def _hub_reachable(timeout: float = 5.0) -> bool: + """Cheap reachability probe for huggingface.co. + + Returns False on a TLS interception proxy as well as on a real outage, which is + what we want: in both cases the cached parquet is the usable path. + """ + import urllib.error + import urllib.request + + try: + urllib.request.urlopen("https://huggingface.co/api/whoami-v2", timeout=timeout) + return True + except urllib.error.HTTPError: + return True # reachable, just unauthenticated + except Exception: + return False + + +def _local_parquet_split(dataset_name: str, split: str) -> list[str]: + """Find cached parquet shards for ``dataset_name``/``split`` in the HF hub cache.""" + cache = os.environ.get("HF_HUB_CACHE") or os.path.expanduser( + "~/.cache/huggingface/hub") + repo_dir = os.path.join(cache, "datasets--" + dataset_name.replace("/", "--")) + snapshots = os.path.join(repo_dir, "snapshots") + if not os.path.isdir(snapshots): + return [] + found: list[str] = [] + for root, _dirs, files in os.walk(snapshots): + for name in sorted(files): + if name.startswith(split) and name.endswith(".parquet"): + found.append(os.path.join(root, name)) + return sorted(found) + + +# --------------------------------------------------------------------------- # +# Text calibration (LLM-only strategies S1, S2) +# --------------------------------------------------------------------------- # +def text_calib_dataloader( + tokenizer, + dataset_name: str = "abisee/cnn_dailymail", + batch_size: int = 1, + num_samples: int = 512, + max_length: int = 512, +) -> DataLoader: + """Return a DataLoader of tokenised ``input_ids`` for calibration. + + Mirrors NVIDIA's ``_text_calib_dataloader``. ``max_length`` only governs + tokenizer truncation during this call and does NOT modify + ``tokenizer.model_max_length`` — verified by an assert in the caller. + """ + from datasets import load_dataset + + # A Jetson is often behind a TLS-intercepting proxy or fully air-gapped, where + # load_dataset() fails even though the parquet shards are already in the hub cache + # (HF's own offline mode does not help: it still wants the dataset script). Fall + # back to reading those shards directly so a calibration run does not depend on + # network reachability. Set HF_DATASETS_LOCAL_ONLY=1 to skip the hub attempt. + local = _local_parquet_split(dataset_name, split="train") + if local and (os.environ.get("HF_DATASETS_LOCAL_ONLY") == "1" or not _hub_reachable()): + print(f" [calib] reading cached parquet ({len(local)} shard(s)) instead of the Hub") + ds = load_dataset("parquet", data_files=local, split="train") + col = "article" if "article" in ds.column_names else ds.column_names[0] + texts = ds[col][:num_samples] + elif "abisee/cnn_dailymail" in dataset_name: + ds = load_dataset(dataset_name, name="3.0.0", split="train") + texts = ds["article"][:num_samples] + else: + ds = load_dataset(dataset_name, split="train") + if "text" in ds.column_names: + col = "text" + elif "article" in ds.column_names: + col = "article" + else: + raise ValueError( + f"Dataset {dataset_name!r} has no 'text' or 'article' column: " + f"{ds.column_names}" + ) + texts = ds[col][:num_samples] + + enc = tokenizer( + texts, + return_tensors="pt", + padding=True, + truncation=True, + max_length=max_length, + ) + return DataLoader(enc["input_ids"], batch_size=batch_size, shuffle=False) + + +# --------------------------------------------------------------------------- # +# Multimodal calibration (visual-quantization strategies S3, S4) +# --------------------------------------------------------------------------- # +def _iter_image_question_pairs(dataset_name: str): + """Yield ``(image, question)`` pairs from a HuggingFace calibration dataset. + + Mirrors NVIDIA's ``_iter_image_question_pairs``. Tolerant of two common + schemas: + * ScienceQA-style: single ``image`` column. + * MMMU-style: numbered ``image_1`` / ``image_2`` / ... columns. + + Splits are tried in the order ``dev`` → ``validation`` → ``train``. + """ + from datasets import load_dataset + + last_err: Optional[Exception] = None + ds = None + for split in ("dev", "validation", "train"): + try: + ds = load_dataset(dataset_name, split=split, streaming=True) + break + except Exception as e: # noqa: BLE001 + last_err = e + if ds is None: + raise RuntimeError( + f"Could not load {dataset_name!r} via any of " + f"split=dev/validation/train" + ) from last_err + + for example in ds: + image = example.get("image") + if image is None: + for i in range(1, 8): + image = example.get(f"image_{i}") + if image is not None: + break + question = example.get("question") or "" + if image is not None and question: + yield image, question + + +def multimodal_calib_dataloader( + processor, + dataset_name: str = "lmms-lab/MMMU", + num_samples: int = 128, + max_length: int = 512, +) -> list[dict[str, Any]]: + """Materialise a list of ``BatchFeature`` dicts with ``input_ids`` + ``pixel_values``. + + Mirrors NVIDIA's ``_multimodal_calib_dataloader``. Streams image-question + pairs through the model's own ``AutoProcessor`` chat template so the visual + tower receives real activations. + + NVIDIA caps multimodal calibration at 128 samples because VLM calibration + is GPU-memory bound. The caller is expected to enforce this cap. + + Returns a *list* (not generator) so ModelOpt can re-iterate forward_loop + during algorithm selection (mirrors NVIDIA's design). + """ + batches: list[dict[str, Any]] = [] + for image, question in _iter_image_question_pairs(dataset_name): + messages = [{ + "role": "user", + "content": [ + {"type": "image", "image": image}, + {"type": "text", "text": question}, + ], + }] + + inputs = processor.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + ) + + batches.append({ + k: v + for k, v in inputs.items() + if v is not None + and not (isinstance(v, torch.Tensor) and v.numel() == 0) + }) + if len(batches) >= num_samples: + break + + if not batches: + raise RuntimeError( + f"No usable multimodal samples from {dataset_name!r}. " + "Check dataset access / processor chat template." + ) + return batches + + +# --------------------------------------------------------------------------- # +# VLN calibration (in-distribution — InternData-N1 VLN-CE) +# --------------------------------------------------------------------------- # +def _import_prompt_builder(): + """Import ``lib/prompt_builder`` (the single source of truth for the System 2 + prompt) regardless of how this module was launched. Kept lazy so the text/MMMU + paths never pay for it.""" + import os + import sys + + lib = os.path.join( + os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), + "lib", + ) + if lib not in sys.path: + sys.path.insert(0, lib) + import prompt_builder # noqa: E402 + + return prompt_builder + + +def _discover_vln_episodes(data_root: str, rgb_key: str): + """Walk a LeRobot dataset tree and yield one record per episode. + + ``data_root`` may be a single scene dir or any parent of several — every dir + holding ``meta/episodes.jsonl`` is picked up. Frames are stored as per-frame + JPGs under ``videos/chunk-XXX//episode_{idx:06d}_{frame}.jpg`` (no + video decoding needed). Falls back to the first available ``*.rgb.*`` stream + when ``rgb_key`` is absent in a scene. + """ + import glob + import json + import os + + records = [] + meta_files = glob.glob( + os.path.join(data_root, "**", "meta", "episodes.jsonl"), recursive=True + ) + if not os.path.isdir(os.path.join(data_root, "meta")) and not meta_files: + raise FileNotFoundError( + f"No LeRobot episodes.jsonl found under {data_root!r} " + "(expected /meta/episodes.jsonl)." + ) + # Include data_root itself if it is a scene dir. + if os.path.isfile(os.path.join(data_root, "meta", "episodes.jsonl")): + meta_files.append(os.path.join(data_root, "meta", "episodes.jsonl")) + + for meta_file in sorted(set(meta_files)): + scene_dir = os.path.dirname(os.path.dirname(meta_file)) + info_path = os.path.join(scene_dir, "meta", "info.json") + chunks_size = 1000 + if os.path.isfile(info_path): + chunks_size = json.load(open(info_path)).get("chunks_size", 1000) + + # Resolve the rgb stream dir for this scene (prefer the requested key). + video_root = os.path.join(scene_dir, "videos") + chunk_dirs = sorted(glob.glob(os.path.join(video_root, "chunk-*"))) + if not chunk_dirs: + continue + + def _rgb_dir_for_chunk(chunk_dir): + want = os.path.join(chunk_dir, rgb_key) + if os.path.isdir(want): + return want + alts = sorted(glob.glob(os.path.join(chunk_dir, "*.rgb.*"))) + return alts[0] if alts else None + + with open(meta_file) as f: + for line in f: + line = line.strip() + if not line: + continue + ep = json.loads(line) + ep_idx = ep["episode_index"] + length = ep.get("length", 0) + tasks = ep.get("tasks") or [] + if length <= 0 or not tasks: + continue + chunk_dir = os.path.join( + video_root, f"chunk-{ep_idx // chunks_size:03d}" + ) + rgb_dir = _rgb_dir_for_chunk(chunk_dir) or _rgb_dir_for_chunk( + chunk_dirs[0] + ) + if rgb_dir is None: + continue + records.append( + { + "scene_dir": scene_dir, + "ep_idx": ep_idx, + "length": length, + "instruction": tasks[0], + "rgb_dir": rgb_dir, + } + ) + return records + + +def vln_calib_dataloader( + processor, + data_root: str, + num_samples: int = 128, + rgb_key: str = "observation.images.rgb.125cm_0deg", + seed: int = 0, +) -> list[dict[str, Any]]: + """Materialise calibration ``BatchFeature`` dicts from InternData-N1 VLN-CE. + + Each sample is one navigation step: a real instruction + a sequence of + egocentric RGB frames (sub-sampled history + current), assembled through the + SAME path the deployed agent uses (``prompt_builder.build_sample_inputs``), so + the number of history frames, the prompt template and the ```` layout + all match ``InternVLAN1Net.s2_step`` exactly. + + Returns a *list* (not a generator) so ModelOpt can re-iterate the forward loop + during algorithm selection, mirroring ``multimodal_calib_dataloader``. + """ + import os + import random + import shutil + import tempfile + + import numpy as np + from PIL import Image + + pb = _import_prompt_builder() + num_history = pb.NUM_HISTORY # single source of truth — must match the prompt + + def frame_path(rgb_dir, ep_idx, frame): + return os.path.join(rgb_dir, f"episode_{ep_idx:06d}_{frame}.jpg") + + episodes = _discover_vln_episodes(data_root, rgb_key) + if not episodes: + raise RuntimeError( + f"No usable VLN episodes under {data_root!r} (rgb_key={rgb_key!r})." + ) + + # The deployed agent resizes every RGB frame to (resize_w, resize_h) before the + # processor (internvla_n1_policy.py). Match that here so the visual-token count + # and ViT activations track deployment, not the raw camera resolution. Resized + # frames are cached under a temp dir and cleaned up once every batch is built + # (the returned batches hold materialised tensors, not paths). + tmp_dir = tempfile.mkdtemp(prefix="vln_calib_") + resized_cache: dict[str, str] = {} + + def resized_frame(src_path): + cached = resized_cache.get(src_path) + if cached is not None: + return cached + img = Image.open(src_path).convert("RGB").resize( + (pb.RESIZE_W, pb.RESIZE_H) + ) + dst = os.path.join(tmp_dir, f"f{len(resized_cache):06d}.jpg") + img.save(dst, quality=95) + resized_cache[src_path] = dst + return dst + + rng = random.Random(seed) + batches: list[dict[str, Any]] = [] + attempts = 0 + max_attempts = num_samples * 20 + + try: + while len(batches) < num_samples and attempts < max_attempts: + attempts += 1 + ep = rng.choice(episodes) + length = ep["length"] + # Pick a "current" step; prefer t>=1 so the sample carries history. + t = rng.randint(1, length - 1) if length > 1 else 0 + + # History frame indices — identical rule to + # prompt_builder.build_conversation (same NUM_HISTORY constant). + if t == 0: + hist = [] + else: + hist = np.unique( + np.linspace(0, t - 1, num_history, dtype=np.int32) + ).tolist() + + src_paths = [frame_path(ep["rgb_dir"], ep["ep_idx"], h) for h in hist] + src_paths.append(frame_path(ep["rgb_dir"], ep["ep_idx"], t)) + if not all(os.path.isfile(p) for p in src_paths): + continue # frame missing on disk — skip this sample + + sample = { + "images": [resized_frame(p) for p in src_paths], + "episode_idx": t, + "instruction": ep["instruction"], + } + try: + inputs = pb.build_sample_inputs(sample, processor) + except Exception: # noqa: BLE001 — drop malformed samples, keep going + continue + + batches.append( + { + k: v + for k, v in inputs.items() + if v is not None + and not (isinstance(v, torch.Tensor) and v.numel() == 0) + } + ) + finally: + shutil.rmtree(tmp_dir, ignore_errors=True) + + if not batches: + raise RuntimeError( + f"Could not build any VLN calibration sample from {data_root!r} " + f"after {attempts} attempts (rgb_key={rgb_key!r})." + ) + return batches \ No newline at end of file diff --git a/recipes/internvla-n1-dualvln/quantize/model_loader.py b/recipes/internvla-n1-dualvln/quantize/model_loader.py new file mode 100644 index 0000000..4ef2819 --- /dev/null +++ b/recipes/internvla-n1-dualvln/quantize/model_loader.py @@ -0,0 +1,291 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Model loading, calibration forward loops, and export for Qwen2.5-VL. + +Logic mirrors NVIDIA's reference quantize_and_export implementation, with the +transformers >= 5.x workaround patches kept intact. + +There is deliberately no InternVLA-N1 branch here. ``repackage_system2.py`` runs first +and turns the checkpoint into a stock Qwen2.5-VL, so this module never sees an +``internvla_n1`` config and never needs to import InternNav. Loading the original +checkpoint directly (with ``config.system1 = "none"`` to suppress System 1) also works +and saves the intermediate copy, but it puts a gated third-party repository -- and the +three monkeypatches needed to load it -- on the critical path of every quantization run. +""" + +import json +import os +import shutil +from typing import Optional + +import torch +from tqdm import tqdm +from transformers import ( + AutoModel, + AutoModelForCausalLM, + AutoModelForImageTextToText, + AutoProcessor, + AutoTokenizer, +) + + +# --------------------------------------------------------------------------- # +# Model loading +# --------------------------------------------------------------------------- # +def load_model( + model_dir: str, + dtype: str = "bf16", + device: str = "cuda", +): + """Load model + tokenizer + optional processor via Auto* classes. + + Mirrors NVIDIA's ``_load_model`` for the standard VLM path: tries + ``AutoModelForImageTextToText`` first (so multimodal architectures are + not silently downgraded to their text-only registration), then falls + back to ``AutoModelForCausalLM`` and finally ``AutoModel``. + """ + if dtype == "fp16": + torch_dtype = torch.float16 + elif dtype == "bf16": + torch_dtype = torch.bfloat16 + else: + raise ValueError(f"Unsupported dtype: {dtype!r}") + + tokenizer = AutoTokenizer.from_pretrained( + model_dir, trust_remote_code=True + ) + + try: + processor = AutoProcessor.from_pretrained( + model_dir, + trust_remote_code=True, + min_pixels=128 * 28 * 28, + max_pixels=2048 * 32 * 32, + ) + except Exception: + processor = None + + last_err: Optional[Exception] = None + model = None + for factory in ( + AutoModelForImageTextToText, + AutoModelForCausalLM, + AutoModel, + ): + try: + model = factory.from_pretrained( + model_dir, + torch_dtype=torch_dtype, + trust_remote_code=True, + ).to(device) + break + except (ValueError, KeyError) as e: + last_err = e + if model is None: + raise RuntimeError( + f"Could not load {model_dir} via any AutoModel factory" + ) from last_err + + model.to(torch_dtype) + + # ModelOpt export_hf_checkpoint crashes when architectures is None. + if getattr(model.config, "architectures", None) is None: + model.config.architectures = [type(model).__name__] + + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + return model, tokenizer, processor + + +# --------------------------------------------------------------------------- # +# Calibration forward loops +# --------------------------------------------------------------------------- # +def calibrate_text(model, dataloader): + """Forward-loop calibration pass for text-only DataLoader. + + Mirrors NVIDIA's ``_calibrate``. + """ + for data in tqdm(dataloader, desc="Calibrating (text)"): + data = data.to(model.device) + model(data) + + +def calibrate_multimodal(model, batches): + """Forward-loop calibration pass for multimodal ``BatchFeature`` dicts. + + Mirrors NVIDIA's ``_calibrate_multimodal``. Float inputs (pixel_values) + inherit the model's dtype; int inputs (input_ids, attention_mask) are + untouched. ``use_cache=False`` matches the reference. + + NaN amax assertions from ModelOpt's internal calibrator are caught and + the offending batch is skipped so calibration can complete on the + remaining valid samples. + """ + device = model.device + model_dtype = next(model.parameters()).dtype + valid_batches = 0 + skipped_nan_batches = 0 + + for batch in tqdm(batches, desc="Calibrating (multimodal)"): + kwargs = {} + for k, v in batch.items(): + if isinstance(v, torch.Tensor): + v = v.to(device) + if v.dtype.is_floating_point: + v = v.to(model_dtype) + kwargs[k] = v + kwargs.setdefault("use_cache", False) + + with torch.no_grad(): + try: + model(**kwargs) + valid_batches += 1 + except AssertionError as exc: + if "detected nan values in amax" in str(exc): + skipped_nan_batches += 1 + continue + raise + + if valid_batches == 0: + raise RuntimeError( + "All multimodal calibration batches were skipped due to NaN amax." + ) + if skipped_nan_batches > 0: + print( + f"[WAR] Skipped {skipped_nan_batches} multimodal calibration " + f"batch(es) with NaN amax." + ) + + +# --------------------------------------------------------------------------- # +# WAR patches (transformers >= 5.x) +# --------------------------------------------------------------------------- # +def normalize_tied_weights_keys(model) -> None: + """WAR for transformers >= 5.x ``_tied_weights_keys`` format change. + + Mirrors NVIDIA's ``_normalize_tied_weights_keys``. Newer transformers + expects each submodule's ``_tied_weights_keys`` attribute to be a + dict-like (so ``modeling_utils._get_tied_weight_keys`` can call ``.keys()``). + Older custom modeling code still declares it as a list, which crashes + ``model.save_pretrained``. + + Convert list-shaped attributes to ``{key: key}`` dicts in place. The dict's + keys exactly match the original list, preserving behavior for downstream + tied-weight tracking. No-op for modules already in the dict format. + """ + for module in model.modules(): + attr = getattr(module, "_tied_weights_keys", None) + if isinstance(attr, list): + module._tied_weights_keys = {k: k for k in attr} + + +def fix_generation_config_for_strict_validate(model) -> None: + """WAR for transformers >= 5.x ``GenerationConfig.validate(strict=True)``. + + Mirrors NVIDIA's ``_fix_generation_config_for_strict_validate``. ModelOpt's + ``export_hf_checkpoint`` calls ``model.save_pretrained`` which runs + ``validate(strict=True)`` on the generation config; that validator rejects + HF checkpoints whose ``generation_config.json`` sets sampling-only kwargs + (``top_p`` / ``top_k`` / ``temperature``) without ``do_sample = True``. + + Force ``do_sample = True`` when any sampling kwarg is present. This only + changes the saved ``generation_config.json``; runtime params are read + separately and are not affected. + """ + gc = getattr(model, "generation_config", None) + if gc is None: + return + sampling_set = ( + getattr(gc, "top_p", None) not in (None, 1.0) + or getattr(gc, "top_k", None) not in (None, 0, 50) + or getattr(gc, "temperature", None) not in (None, 1.0) + ) + if sampling_set and not getattr(gc, "do_sample", False): + gc.do_sample = True + + +# --------------------------------------------------------------------------- # +# Export helpers +# --------------------------------------------------------------------------- # + +# Files copied verbatim from source model dir — preserves preprocessing config +# untouched by calibration. Matches NVIDIA's copy list. +PROCESSOR_FILES = ( + "preprocessor_config.json", + "processor_config.json", + "video_preprocessor_config.json", + "chat_template.jinja", + "chat_template.json", +) + +# Tokenizer files copied verbatim from source. Quantization only rewrites model +# weights (+ config.json / generation_config.json); it never alters the +# tokenizer, so the source tokenizer files are byte-for-byte authoritative. +# Copying them verbatim — instead of trusting tokenizer.save_pretrained() to +# round-trip cleanly — makes it impossible for calibration-time truncation state +# (max_length=512, truncation_strategy, ...) to leak into the deployed +# tokenizer_config.json. That leak silently caps inference sequence length +# (e.g. MME benchmark truncated to 512). Non-existent files are skipped, so +# listing both fast (tokenizer.json) and slow (vocab.json/merges.txt) artifacts +# is safe across tokenizer variants. +TOKENIZER_FILES = ( + "tokenizer_config.json", + "tokenizer.json", + "vocab.json", + "merges.txt", + "special_tokens_map.json", + "added_tokens.json", +) + + +# --------------------------------------------------------------------------- # +# Export +# --------------------------------------------------------------------------- # +def export_quantized_model( + model, + tokenizer, + processor, + model_dir: str, + output_dir: str, +) -> None: + """Save a quantized model with HF-compatible layout. + + Mirrors NVIDIA's export tail: + 1. Apply transformers >= 5.x WARs. + 2. ``export_hf_checkpoint`` writes weights + ``config.json`` + + ``generation_config.json`` + ``hf_quant_config.json``. + 3. Save tokenizer + processor (so the full file set always exists). + 4. Overwrite the tokenizer and preprocessor config files with verbatim + source copies. This is the authoritative step: quantization never + changes these, so the source files are correct by definition, and a + verbatim copy guarantees calibration's ``max_length=512`` truncation + state can never leak into ``tokenizer_config.json``. Only allow-listed + filenames are overwritten — the quantized weights, ``config.json``, + ``generation_config.json`` and ``hf_quant_config.json`` from step 2 are + left untouched. + """ + from modelopt.torch.export import export_hf_checkpoint + + fix_generation_config_for_strict_validate(model) + normalize_tied_weights_keys(model) + + os.makedirs(output_dir, exist_ok=True) + + with torch.inference_mode(): + export_hf_checkpoint(model, export_dir=output_dir) + + tokenizer.save_pretrained(output_dir) + + if processor is not None: + processor.save_pretrained(output_dir) + + # Overwrite tokenizer + processor configs with verbatim source copies so + # quantization/calibration cannot alter them. This is the authoritative + # final write: it runs after both save_pretrained() calls and only touches + # the allow-listed filenames, so weights, config.json, generation_config.json + # and hf_quant_config.json produced by export_hf_checkpoint are untouched. + for fname in (*TOKENIZER_FILES, *PROCESSOR_FILES): + src = os.path.join(model_dir, fname) + if os.path.isfile(src): + shutil.copy2(src, os.path.join(output_dir, fname)) diff --git a/recipes/internvla-n1-dualvln/quantize/prompt_builder.py b/recipes/internvla-n1-dualvln/quantize/prompt_builder.py new file mode 100644 index 0000000..6d7a217 --- /dev/null +++ b/recipes/internvla-n1-dualvln/quantize/prompt_builder.py @@ -0,0 +1,237 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Prompt / conversation builders for InternVLA-N1 System 2. + +Reconstructs the exact multi-image chat prompt that ``InternVLAN1Net.s2_step`` builds, for both +the normal turn and the look-down (turn 2) follow-up. Kept as a single source of truth so the +verification and calibration scripts never diverge from the agent's real prompt format. + +Read-only: only a processor/tokenizer is loaded; no model weights, no writes to any source tree. +""" +from __future__ import annotations + +import argparse +import json +import os +import sys +from typing import Any + +import numpy as np +from PIL import Image +from transformers import AutoProcessor + +# Eval config values (InternNav h1_internvla_n1_async_cfg). +RESIZE_W = 384 +RESIZE_H = 384 +NUM_HISTORY = 8 + +# Special token ids (internvla_n1.py). +IMAGE_TOKEN_INDEX = 151655 +TRAJ_TOKEN_INDEX = 151667 + +# Base prompt template (internvla_n1_policy.py), verbatim. +PROMPT_TEMPLATE = ( + "You are an autonomous navigation assistant. Your task is to . " + "Where should you go next to stay on track? Please output the next waypoint's " + "coordinates in the image. Please output STOP when you have successfully " + "completed the task." +) +CONJUNCTION = "you can see " +DEFAULT_IMAGE_TOKEN = "" + +SAMPLE_INSTRUCTION = ( + "walk out of the bathroom and turn left, then walk down the hallway and " + "stop in front of the second door on your right" +) + +# Two processor construction variants: the deploy default vs the calibration override. +PROCESSOR_VARIANTS: dict[str, dict[str, Any]] = { + "deploy": {}, + "calib": {"min_pixels": 128 * 28 * 28, "max_pixels": 2048 * 32 * 32}, +} + + +def split_and_clean(text: str) -> list[str]: + """Split around '' and drop empty parts. + + Mirrors ``internnav.model.utils.vln_utils.split_and_clean`` (re-implemented so this module + runs without the full internnav dependency tree). + """ + import re + + parts = re.split(r"()", text) + return [p for p in (s.strip() if s != DEFAULT_IMAGE_TOKEN else s for s in parts) if p] + + +def build_conversation(episode_idx: int, instruction: str) -> tuple[list[dict], int]: + """Build the normal-turn prompt (look_down=False), matching ``internvla_n1_policy.py``. + + Returns ``(conversation, n_images)``. + """ + sources_value = PROMPT_TEMPLATE.replace(".", instruction) + + if episode_idx == 0: + history_id: list[int] = [] + else: + history_id = np.unique( + np.linspace(0, episode_idx - 1, NUM_HISTORY, dtype=np.int32) + ).tolist() + placeholder = (DEFAULT_IMAGE_TOKEN + "\n") * len(history_id) + sources_value += f" These are your historical observations: {placeholder}." + + # The current frame is always appended last. + sources_value += f" {CONJUNCTION}{DEFAULT_IMAGE_TOKEN}." + + n_images = len(history_id) + 1 + + content: list[dict] = [] + for part in split_and_clean(sources_value): + if part == DEFAULT_IMAGE_TOKEN: + content.append({"type": "image", "image": None}) # placeholder + else: + content.append({"type": "text", "text": part}) + + return [{"role": "user", "content": content}], n_images + + +def build_conversation_lookdown( + episode_idx: int, instruction: str, assistant_reply: str +) -> list[dict]: + """Build the turn-2 (look-down) conversation, matching the agent. + + The real flow is two turns: turn 1 returns '↓' (action 5 = look down); the robot tilts its + camera; turn 2 sends the extra downward view **without resetting history**, and the coordinate + is produced here. The turn-2 user message repeats no instruction/history — only + " you can see ." — and only the newly appended look-down image binds to it. + """ + turn1, _ = build_conversation(episode_idx, instruction) + value = f" {CONJUNCTION}{DEFAULT_IMAGE_TOKEN}." + + content: list[dict] = [] + for part in split_and_clean(value): + if part == DEFAULT_IMAGE_TOKEN: + content.append({"type": "image", "image": None}) + else: + content.append({"type": "text", "text": part}) + + return [ + turn1[0], + {"role": "assistant", "content": [{"type": "text", "text": assistant_reply}]}, + {"role": "user", "content": content}, + ] + + +def build_sample_inputs(sample: dict, processor): + """Build processor ``inputs`` for one golden sample, handling both turn 1 and turn 2. + + Single source of truth for gate / calibration / profiling: every hand-rebuilt prompt is a + place that can drift from real behavior. ``turn`` defaults to 1. + """ + from PIL import Image + + turn = sample.get("turn", 1) + images = [Image.open(p).convert("RGB") for p in sample["images"]] + + if turn == 2: + reply = sample.get("assistant_turn1") + if not reply: + raise ValueError( + f"Turn-2 sample missing 'assistant_turn1' " + f"({sample['episode']} ep{sample['episode_idx']})." + ) + conv = build_conversation_lookdown( + sample["episode_idx"], sample["instruction"], reply) + else: + conv, _ = build_conversation(sample["episode_idx"], sample["instruction"]) + + j = 0 + for t in conv: + for it in t["content"]: + if it["type"] == "image" and it.get("image") is None: + it["image"] = images[j] + j += 1 + if j != len(images): + raise ValueError( + f" token count ({j}) != number of images ({len(images)}) — " + f"{sample['episode']} ep{sample['episode_idx']} turn {turn}" + ) + + text = processor.apply_chat_template(conv, tokenize=False, add_generation_prompt=True) + return processor(text=[text], images=images, return_tensors="pt") + + +def census_one(processor, episode_idx: int, instruction: str) -> dict[str, Any]: + """Count visual vs text tokens for one episode step (dummy images sized to the resize).""" + conversation, n_images = build_conversation(episode_idx, instruction) + + images = [ + Image.fromarray(np.zeros((RESIZE_H, RESIZE_W, 3), dtype=np.uint8)).convert("RGB") + for _ in range(n_images) + ] + img_i = 0 + for item in conversation[0]["content"]: + if item["type"] == "image": + item["image"] = images[img_i] + img_i += 1 + + text = processor.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True) + inputs = processor(text=[text], images=images, return_tensors="pt") + + ids = inputs["input_ids"][0] + total = int(ids.numel()) + n_visual = int((ids == IMAGE_TOKEN_INDEX).sum()) + n_traj = int((ids == TRAJ_TOKEN_INDEX).sum()) + grid = inputs.get("image_grid_thw") + return { + "episode_idx": episode_idx, + "n_images": n_images, + "total_tokens": total, + "visual_tokens": n_visual, + "text_tokens": total - n_visual - n_traj, + "traj_tokens": n_traj, + "visual_pct": round(100.0 * n_visual / total, 2), + "tokens_per_image": (n_visual // n_images) if n_images else 0, + "image_grid_thw": grid.tolist() if grid is not None else None, + } + + +def main() -> int: + """Token census: report the visual/text token split the LLM actually sees at runtime.""" + ap = argparse.ArgumentParser(description=main.__doc__) + ap.add_argument("--model-path", default=os.environ.get( + "INTERNVLA_CKPT", os.path.expanduser("~/InternNav/checkpoints/InternVLA-N1-DualVLN"))) + ap.add_argument("--episodes", type=int, nargs="+", default=[0, 1, 2, 4, 8, 16, 32, 64]) + ap.add_argument("--out", default="token_census.json") + args = ap.parse_args() + + if not os.path.isdir(args.model_path): + print(f"ERROR: model-path not found: {args.model_path}", file=sys.stderr) + return 1 + + report: dict[str, Any] = {"model_path": args.model_path, "variants": {}} + for variant, kwargs in PROCESSOR_VARIANTS.items(): + print(f"\n### processor variant: {variant} {kwargs or '(default)'}") + try: + processor = AutoProcessor.from_pretrained( + args.model_path, trust_remote_code=True, **kwargs) + except Exception as exc: # noqa: BLE001 + print(f" !! could not load processor: {exc}") + report["variants"][variant] = {"error": str(exc)} + continue + rows = [census_one(processor, ep, SAMPLE_INSTRUCTION) for ep in args.episodes] + for r in rows: + print(f" ep={r['episode_idx']:>3} imgs={r['n_images']:>2} " + f"total={r['total_tokens']:>6} visual={r['visual_tokens']:>6} " + f"visual%={r['visual_pct']:>6.2f}") + report["variants"][variant] = {"processor_kwargs": kwargs, "rows": rows} + + os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True) + with open(args.out, "w") as f: + json.dump(report, f, indent=2) + print(f"\nWrote: {args.out}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/quantize/quantize.py b/recipes/internvla-n1-dualvln/quantize/quantize.py new file mode 100755 index 0000000..e14abd2 --- /dev/null +++ b/recipes/internvla-n1-dualvln/quantize/quantize.py @@ -0,0 +1,323 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Quantize Qwen2.5-VL with NVIDIA ModelOpt. + +Usage: + python quantize.py --strategy s1 --cfg fp8_default + python quantize.py --strategy s3 --cfg nvfp4_awq_full --dry-run + +Supports 4 strategies x 5 schemes; see ``configs/schemes.yaml`` for the matrix +and which combinations are blocked or experimental on this hardware. +""" + +import argparse +import os +import sys +import time + +import modelopt.torch.quantization as mtq +import torch + +from calibration import ( + multimodal_calib_dataloader, + text_calib_dataloader, + vln_calib_dataloader, +) +from quant_schemes import ( + build_quant_config, + calib_batch_size, + describe, + load_registry, + scheme_names, + strategy_names, + validate, +) +from model_loader import ( + calibrate_multimodal, + calibrate_text, + export_quantized_model, + load_model, +) + + +# --------------------------------------------------------------------------- # +# Defaults +# --------------------------------------------------------------------------- # +DEFAULT_MODEL_DIR = os.path.expanduser( + os.environ.get("INTERNVLA_CKPT", "~/InternNav/checkpoints/InternVLA-N1-DualVLN") +) +DEFAULT_OUTPUT_BASE = os.path.expanduser( + os.environ.get("VLN_OPT_WORK", "~/vln-opt-work") +) +TEXT_DATASET = "abisee/cnn_dailymail" +MULTIMODAL_DATASET = "lmms-lab/MMMU" +MULTIMODAL_MAX_SAMPLES = 128 # NVIDIA's cap — VLM calibration is GPU-mem bound +# In-distribution calibration: the InternData-N1 VLN-CE subset System 2 was +# fine-tuned on. Override with --calib-data or $VLN_CALIB_DATA. +VLN_CALIB_DATA = os.path.expanduser( + os.environ.get( + # Default = the output of build/00_fetch_calib_scenes.sh. Point it at any + # InternData-N1 VLN-CE tree (a scene dir or a parent of several) via the env var. + "VLN_CALIB_DATA", + "~/vln-opt-work/calib_scenes", + ) +) +VLN_RGB_KEY = os.environ.get("VLN_RGB_KEY", "observation.images.rgb.125cm_0deg") + + +# --------------------------------------------------------------------------- # +# CLI +# --------------------------------------------------------------------------- # +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Quantize Qwen2.5-VL with NVIDIA ModelOpt", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=describe(), + ) + parser.add_argument( + "--strategy", + required=True, + choices=strategy_names(), + help="Quantization strategy (see options below).", + ) + parser.add_argument( + "--scheme", + required=True, + choices=scheme_names(), + help="Quantization scheme (see options below).", + ) + parser.add_argument( + "--model_path", + default=DEFAULT_MODEL_DIR, + help=f"Path to source HF model (default: {DEFAULT_MODEL_DIR}).", + ) + parser.add_argument( + "--output_path", + default=None, + help="Output dir for quantized model. Default: " + f"{DEFAULT_OUTPUT_BASE}/qwen2.5-vl-7b--/", + ) + parser.add_argument( + "--num_calib_samples", + type=int, + default=512, + help="Calibration sample count. Capped at 128 for multimodal " + "strategies (S3/S4) per NVIDIA's reference. Default: 512.", + ) + parser.add_argument( + "--calib", + default="auto", + choices=["auto", "text", "multimodal", "vln"], + help="Calibration data source. 'auto' (default) keeps the legacy " + "behaviour: text (cnn_dailymail) for LLM-only strategies, multimodal " + "(MMMU) when the visual tower is quantized. 'vln' uses the " + "in-distribution InternData-N1 VLN-CE set (recommended for this model).", + ) + parser.add_argument( + "--calib_data_root", + default=VLN_CALIB_DATA, + help=f"Root of the VLN-CE LeRobot data for --calib vln " + f"(default: {VLN_CALIB_DATA}).", + ) + parser.add_argument( + "--dtype", + default="bf16", + choices=["fp16", "bf16"], + help="Model load dtype. Default: bf16 (matches source config).", + ) + parser.add_argument( + "--device", + default="cuda", + help="Device for model + calibration. Default: cuda.", + ) + parser.add_argument( + "--resume", + default=None, + metavar="DIR", + help="Enable layerwise calibration resume from DIR (useful for " + "AWQ/Hessian recovery on crash).", + ) + parser.add_argument( + "--allow_experimental", + action="store_true", + help="Permit schemes marked experimental in configs/schemes.yaml. NVFP4 needs " + "this: it quantizes and generates fluent text, but the System 2 -> System 1 " + "bridge breaks (z_latents cosine 0.647 vs 0.9956 for FP8).", + ) + parser.add_argument( + "--dry_run", + action="store_true", + help="Load model + validate environment, skip quantization. " + "Useful for verifying setup before a long run.", + ) + return parser.parse_args() + + +def derive_output_path(args: argparse.Namespace) -> str: + if args.output_path is not None: + return args.output_path + tag = f"internvla-n1-system2-{args.strategy}-{args.scheme}" + return os.path.join(DEFAULT_OUTPUT_BASE, tag) + + +def print_environment() -> None: + """Quick environment summary for logging at the top of every run.""" + print("=" * 70) + print("Environment:") + print(f" torch: {torch.__version__}") + print(f" CUDA: {torch.cuda.is_available()}") + if torch.cuda.is_available(): + props = torch.cuda.get_device_properties(0) + cc = torch.cuda.get_device_capability(0) + print(f" device: {props.name}") + print(f" arch: sm_{cc[0]}{cc[1]}") + print(f" VRAM: {props.total_memory / 1e9:.1f} GB") + print(f" arch list: {torch.cuda.get_arch_list()}") + print("=" * 70) + + +# --------------------------------------------------------------------------- # +# Main +# --------------------------------------------------------------------------- # +def main() -> int: + args = parse_args() + output_path = derive_output_path(args) + + print_environment() + + registry = load_registry() + + # Reject impossible combinations before the model is loaded: a bad request should + # cost seconds, not a checkpoint load followed by a mid-quantization crash. + try: + validate(args.scheme, args.strategy, model_path=args.model_path, + allow_experimental=args.allow_experimental, registry=registry) + except ValueError as exc: + print(f"[ERROR] {exc}") + return 1 + + strat = registry["strategies"][args.strategy] + cfg_info = registry["schemes"][args.scheme] + + # Resolve calibration source. 'auto' preserves the legacy behaviour; an + # explicit --calib overrides it (e.g. VLN calib for an LLM-only strategy). + calib_mode = args.calib + if calib_mode == "auto": + calib_mode = "multimodal" if strat["quantize_visual"] else "text" + calib_desc = { + "text": "text-only (cnn_dailymail)", + "multimodal": "multimodal (MMMU image+text)", + "vln": "VLN in-distribution (InternData-N1 VLN-CE)", + }[calib_mode] + + print() + print(f"Strategy: {args.strategy} — {strat['description']}") + print(f"CFG preset: {args.scheme} — {cfg_info['description']}") + print(f"Model path: {args.model_path}") + print(f"Output path: {output_path}") + print(f"Calibration: {calib_desc}") + print() + + if not os.path.isdir(args.model_path): + print(f"ERROR: model_path not found: {args.model_path}", file=sys.stderr) + return 1 + + # -- Load model ------------------------------------------------------- # + t_load_start = time.time() + print("[1/4] Loading model...") + model, tokenizer, processor = load_model( + args.model_path, dtype=args.dtype, device=args.device + ) + original_max_length = tokenizer.model_max_length + print(f" done in {time.time() - t_load_start:.1f}s") + print(f" model class: {type(model).__name__}") + print(f" tokenizer.model_max_length (preserved): {original_max_length}") + + if args.dry_run: + print("[dry-run] Skipping quantization. Environment OK.") + return 0 + + # -- Build quant config ---------------------------------------------- # + quant_cfg = build_quant_config( + scheme=args.scheme, + strategy=args.strategy, + layerwise_checkpoint_dir=args.resume, + ) + + # -- Build dataloader & forward loop --------------------------------- # + t_calib_start = time.time() + print("[2/4] Preparing calibration data...") + + if calib_mode in ("multimodal", "vln"): + if processor is None: + raise RuntimeError( + f"{calib_mode} calibration requires an AutoProcessor but none " + "was found in the model dir." + ) + # Both image paths are GPU-memory bound (multi-image forward), so they + # share NVIDIA's multimodal sample cap. + mm_samples = min(args.num_calib_samples, MULTIMODAL_MAX_SAMPLES) + if mm_samples < args.num_calib_samples: + print(f" capping num_samples {args.num_calib_samples} → " + f"{mm_samples} (multimodal GPU-mem cap)") + if calib_mode == "vln": + batches = vln_calib_dataloader( + processor, + data_root=args.calib_data_root, + num_samples=mm_samples, + rgb_key=VLN_RGB_KEY, + ) + print(f" VLN batches prepared: {len(batches)} " + f"(data: {args.calib_data_root})") + else: + batches = multimodal_calib_dataloader( + processor, + dataset_name=MULTIMODAL_DATASET, + num_samples=mm_samples, + ) + print(f" multimodal batches prepared: {len(batches)}") + forward_loop = lambda m: calibrate_multimodal(m, batches) # noqa: E731 + else: + batch_size = calib_batch_size(args.scheme, is_image_calib=False) + loader = text_calib_dataloader( + tokenizer, + dataset_name=TEXT_DATASET, + batch_size=batch_size, + num_samples=args.num_calib_samples, + ) + forward_loop = lambda m: calibrate_text(m, loader) # noqa: E731 + print(f" text loader prepared: batch_size={batch_size}, " + f"samples={args.num_calib_samples}") + + # -- Quantize -------------------------------------------------------- # + print("[3/4] Running quantization...") + mtq.quantize(model, quant_cfg, forward_loop=forward_loop) + mtq.print_quant_summary(model) + print(f" quantization done in {time.time() - t_calib_start:.1f}s") + + # Safety: tokenizer.model_max_length must not have been touched by + # calibration. This is the user's main concern about preprocessor leakage. + assert tokenizer.model_max_length == original_max_length, ( + f"tokenizer.model_max_length changed during calibration: " + f"{original_max_length} → {tokenizer.model_max_length}" + ) + + # -- Export ---------------------------------------------------------- # + t_export_start = time.time() + print("[4/4] Exporting quantized checkpoint...") + export_quantized_model( + model=model, + tokenizer=tokenizer, + processor=processor, + model_dir=args.model_path, + output_dir=output_path, + ) + print(f" export done in {time.time() - t_export_start:.1f}s") + print() + print(f"Saved to {output_path}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) \ No newline at end of file From ddbe32e59e8d260020e7806cc70f150b8ec5329c Mon Sep 17 00:00:00 2001 From: hungho77 Date: Thu, 13 Aug 2026 16:29:19 +0700 Subject: [PATCH 04/30] feat(internvla-n1-dualvln): load ModelOpt checkpoints back with their scales applied Accuracy benchmarking a quantized checkpoint in PyTorch needs it loaded correctly, and the obvious way is silently wrong. export_hf_checkpoint writes LLM weights as real torch.float8_e4m3fn with dequantization scales in separate weight_scale/input_scale tensors. AutoModelForImageTextToText.from_pretrained on that directory appears to succeed -- it warns the scale tensors 'were not used when initializing' and continues -- but every quantized weight then sits at its raw FP8 magnitude. Measured on this checkpoint, the naive load produces weights 630x to 1799x too large depending on the layer. Any accuracy number taken from it would be reporting a broken load as quantization damage. load_quantized.py reconstructs w_bf16 = w_fp8.to(bf16) * weight_scale over a normal from_pretrained (which is what materialises the rope buffers; a meta-device build leaves them empty). Verified against the unquantized System 2 in float64: cosine 0.99965 and 2.6-2.7% relative error per projection, which is what FP8 E4M3 per-tensor should cost, against a scale ratio of exactly 1.000. Modules excluded from quantization -- the vision tower and lm_head -- are stored bf16 and pass through untouched. Note what this does and does not model: weight quantization only. Real FP8 W8A8 also quantizes activations, which the TensorRT engine does and this does not, so a number from here is a lower bound on the engine's deviation rather than a prediction of it. The docstring says so, because the distinction is easy to lose. (cherry picked from commit 63f376eb99b5f2a6ac3e05992898bc2b4b6e31bc) --- .../quantize/load_quantized.py | 128 ++++++++++++++++++ 1 file changed, 128 insertions(+) create mode 100644 recipes/internvla-n1-dualvln/quantize/load_quantized.py diff --git a/recipes/internvla-n1-dualvln/quantize/load_quantized.py b/recipes/internvla-n1-dualvln/quantize/load_quantized.py new file mode 100644 index 0000000..ff1e723 --- /dev/null +++ b/recipes/internvla-n1-dualvln/quantize/load_quantized.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Load a ModelOpt-exported checkpoint back into PyTorch with its scales applied. + +This exists because the obvious thing does not work. ``export_hf_checkpoint`` writes the +LLM weights as real ``torch.float8_e4m3fn`` with the dequantization scales in separate +``*.weight_scale`` / ``*.input_scale`` tensors. Calling +``AutoModelForImageTextToText.from_pretrained`` on that directory *appears* to succeed -- +it warns that the scale tensors "were not used when initializing" and moves on -- but the +resulting model has every quantized weight off by its scale factor. Any accuracy measured +that way is measuring a broken load, not quantization error, and it will look far worse +than the engine actually is. + +So reconstruct the weights explicitly:: + + w_bf16 = w_fp8.to(bfloat16) * weight_scale + +Two caveats worth stating plainly: + +* This reproduces **weight** quantization error only. Real FP8 W8A8 also quantizes + activations, which the TensorRT engine does and this does not. Weight error is the + dominant term and this is the standard PyTorch-side proxy, but a number from here is a + lower bound on the engine's deviation, not a prediction of it. +* Modules excluded from quantization (the vision tower, ``lm_head``) are stored in bf16 + already and pass through untouched, which is the intended behaviour. +""" +import glob +import json +import os +from typing import Optional + +import torch +from safetensors import safe_open + + +def quant_algo(model_path: str) -> Optional[str]: + """Return the quantization algorithm recorded by ModelOpt, or None if unquantized.""" + cfg = os.path.join(model_path, "hf_quant_config.json") + if not os.path.isfile(cfg): + return None + with open(cfg) as f: + return json.load(f).get("quantization", {}).get("quant_algo") + + +def _iter_shards(model_path: str): + for shard in sorted(glob.glob(os.path.join(model_path, "*.safetensors"))): + if os.path.basename(shard) == "bridge.safetensors": + continue + yield shard + + +def dequantize_state_dict(model_path: str, + dtype: torch.dtype = torch.bfloat16) -> dict[str, torch.Tensor]: + """Read a ModelOpt checkpoint and return a plain state dict with scales folded in. + + Weights that were not quantized are returned as stored. + """ + scales: dict[str, torch.Tensor] = {} + raw: dict[str, torch.Tensor] = {} + + for shard in _iter_shards(model_path): + with safe_open(shard, framework="pt") as f: + for key in f.keys(): + tensor = f.get_tensor(key) + if key.endswith(("weight_scale", "input_scale", "weight_scale_2")): + scales[key] = tensor + else: + raw[key] = tensor + + out: dict[str, torch.Tensor] = {} + n_dequant = 0 + for key, tensor in raw.items(): + scale = scales.get(key + "_scale") + if scale is not None and tensor.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + out[key] = tensor.to(torch.float32).mul_(scale.to(torch.float32)).to(dtype) + n_dequant += 1 + elif tensor.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + # FP8 storage with no scale would silently corrupt the weight; refuse. + raise ValueError(f"{key} is FP8 but has no matching {key}_scale in the checkpoint") + else: + out[key] = tensor.to(dtype) if tensor.is_floating_point() else tensor + + print(f" [load] dequantized {n_dequant} FP8 tensors, " + f"{len(out) - n_dequant} passed through unchanged") + return out + + +def load_for_eval(model_path: str, dtype: torch.dtype = torch.bfloat16, + device: str = "cuda"): + """Load a checkpoint for evaluation, applying ModelOpt scales when present. + + Returns ``(model, processor, algo)`` where ``algo`` is the quantization algorithm + string or None. + """ + from transformers import AutoConfig, AutoProcessor, Qwen2_5_VLForConditionalGeneration + + algo = quant_algo(model_path) + processor = AutoProcessor.from_pretrained( + model_path, min_pixels=128 * 28 * 28, max_pixels=2048 * 32 * 32) + + # from_pretrained gives a correctly wired model with its buffers (rope inv_freq and + # friends) materialised. For a quantized checkpoint it silently drops the scale + # tensors and casts the FP8 weights straight to bf16, so those weights are wrong by + # their scale factor -- they get overwritten below. Building on a meta device instead + # would avoid the wasted load but leaves the buffers unmaterialised. + model = Qwen2_5_VLForConditionalGeneration.from_pretrained( + model_path, torch_dtype=dtype, low_cpu_mem_usage=True) + + if algo is not None: + state = dequantize_state_dict(model_path, dtype=dtype) + own = dict(model.named_parameters()) + own.update(dict(model.named_buffers())) + n_fixed = 0 + with torch.no_grad(): + for key, tensor in state.items(): + target = own.get(key) + if target is None: + continue + if target.shape != tensor.shape: + raise RuntimeError(f"shape mismatch for {key}: " + f"model {tuple(target.shape)} vs " + f"checkpoint {tuple(tensor.shape)}") + target.copy_(tensor.to(target.dtype)) + n_fixed += 1 + print(f" [load] applied {n_fixed} dequantized tensors over the raw load") + + model = model.to(device=device, dtype=dtype).eval() + return model, processor, algo From 7f138749b246af1bea152d69ca75aeab5366484a Mon Sep 17 00:00:00 2001 From: hungho77 Date: Thu, 13 Aug 2026 16:56:08 +0700 Subject: [PATCH 05/30] feat(internvla-n1-dualvln): score quantization damage on held-out VLN episodes The official metrics (SR, SPL, NE, OS, nDTW) are closed-loop and need Habitat or InternUtopia plus MP3D scenes. Neither is installed here and neither is practical on a Jetson, so this measures what quantization actually threatens and what can be measured offline: whether System 2 still picks the same waypoint. Result on 42 held-out samples from two scenes, BF16 System 2 versus its FP8 quantization, identical samples and greedy decoding: pixel_goal_l2 BF16 mean 47.24 px / median 27.05 px FP8 mean 50.12 px / median 26.97 px agreement 31/42 replies byte-identical, median deviation 0.00 px The mean moves 2.88 px while the median does not move at all, because the gap is carried by 8 samples out of 42 rather than by a systematic shift -- worst case 89 px. That distinction matters for a navigation model and is why both statistics are reported. Two harness bugs were found by disbelieving the first numbers rather than publishing them, and both are now prevented in code: Feeding a pitched frame through a turn-1 prompt gave a 0% coordinate parse rate. The model was right and the harness was wrong: the agent runs two turns, and the coordinate only exists in turn 2, after the level view has been answered with a look-down token and the tilted frame appended. The pixel goal in the level camera is the [-1,-1] sentinel on every frame precisely because a level camera cannot see a point on the floor. prompt_builder's 'episode_idx' is the step index within an episode, not the episode number -- it derives the history frames from linspace(0, episode_idx-1). Passing the episode number produced one image placeholder against nine images. action_accuracy is now scored on turn 1 only. At turn 2 the model emits a coordinate by design, so scoring it against a movement action reads 0% and looks like a regression. (cherry picked from commit 08a64f6320f43bf4a492b9d60688397a3dddb2cb) --- .../quantize/benchmark_accuracy.py | 259 ++++++++++++++++++ 1 file changed, 259 insertions(+) create mode 100644 recipes/internvla-n1-dualvln/quantize/benchmark_accuracy.py diff --git a/recipes/internvla-n1-dualvln/quantize/benchmark_accuracy.py b/recipes/internvla-n1-dualvln/quantize/benchmark_accuracy.py new file mode 100644 index 0000000..b3bb700 --- /dev/null +++ b/recipes/internvla-n1-dualvln/quantize/benchmark_accuracy.py @@ -0,0 +1,259 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Score a System-2 checkpoint on held-out VLN episodes, without a simulator. + +The official InternVLA-N1 metrics (SR, SPL, NE, OS, nDTW) are all closed-loop and need +Habitat or InternUtopia plus MP3D scenes. Neither is practical on a Jetson, so this +measures the thing that quantization actually threatens and that *can* be measured +offline: whether System 2 still emits the same navigation decision. + +Two metrics, both against the LeRobot ground truth: + +* ``pixel_goal_l2`` -- Euclidean distance between the predicted and annotated waypoint, + in pixels of the 384x384 prompt image. This is the primary number: the waypoint is what + System 1 consumes, so an error here is an error in where the robot goes. +* ``action_accuracy`` -- agreement on the discrete action token (STOP / up / left / right). + +Both are computed per checkpoint on identical samples, so two runs are directly +comparable. The interesting quantity is the *delta* between an unquantized and a +quantized checkpoint, not either absolute value. + +Prompts are built through ``prompt_builder``, the same path the deployed agent uses, so +what is scored is the deployed prompt rather than a hand-rebuilt approximation. +""" +import argparse +import glob +import json +import os +import re +import sys +import time + +import numpy as np +import torch + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import prompt_builder as pb # noqa: E402 +from load_quantized import load_for_eval # noqa: E402 + +# From InternNav's dataset: idx2actions = {0: STOP, 1: up, 2: left, 3: right, 5: down} +IDX2ACTION = {0: "STOP", 1: "↑", 2: "←", 3: "→", 5: "↓"} +# The agent runs two turns. Turn 1 shows the level camera plus history and the model +# replies with a discrete action -- or with the look-down token when the next thing it +# needs is a waypoint. The robot then tilts, and turn 2 appends the single pitched frame, +# where the coordinate is produced. That is why every goal.125cm_0deg entry in the data is +# the [-1,-1] sentinel: a level camera cannot see a point on the floor. +# +# Scoring has to follow the same two turns. Feeding a pitched frame through a turn-1 +# prompt puts the model off-distribution and it answers with an action instead of a +# coordinate, which reads as a 0% parse rate and looks like a model failure when it is a +# harness failure. +LEVEL_CAMERA = "125cm_0deg" +DEFAULT_CAMERA = "125cm_30deg" +LOOKDOWN_TOKEN = "\u2193" + + +def parse_xy(text: str): + """Pull the first two integers out of a model reply, matching InternNav's parser.""" + nums = re.findall(r"-?\d+", text) + if len(nums) < 2: + return None + return float(nums[0]), float(nums[1]) + + +def parse_action(text: str): + for token in ("STOP", "↑", "←", "→", "↓"): + if token in text: + return token + return None + + +def discover_samples(data_root: str, max_samples: int, seed: int = 0, + camera: str = DEFAULT_CAMERA) -> list[dict]: + """Collect (images, instruction, GT waypoint, GT action) tuples from LeRobot episodes.""" + import pyarrow.parquet as pq + + rng = np.random.default_rng(seed) + rgb_key = f"observation.images.rgb.{camera}" + goal_col = f"goal.{camera}" + samples: list[dict] = [] + for meta in sorted(glob.glob(os.path.join(data_root, "**", "meta", "episodes.jsonl"), + recursive=True)): + scene = os.path.dirname(os.path.dirname(meta)) + episodes = [json.loads(line) for line in open(meta)] + for ep in episodes: + idx = ep["episode_index"] + length = ep["length"] + if length < 2: + continue + table = None + for parquet in glob.glob(os.path.join(scene, "data", "**", + f"episode_{idx:06d}.parquet"), + recursive=True): + table = pq.read_table(parquet) + break + if table is None or goal_col not in table.schema.names: + continue + + goals = table.column(goal_col).to_pylist() + actions = table.column("action").to_pylist() + # goal is [-1, -1] on frames that carry no waypoint (a pure turn or stop). + # Scoring L2 against that sentinel would be meaningless, so sample only from + # frames that actually annotate one. + usable = [i for i in range(1, min(length, len(goals))) + if goals[i] is not None and int(goals[i][0]) >= 0] + if not usable: + continue + t = int(usable[rng.integers(0, len(usable))]) + + history = np.unique(np.linspace(0, t - 1, pb.NUM_HISTORY, dtype=np.int32)).tolist() + level_key = f"observation.images.rgb.{LEVEL_CAMERA}" + frames = [os.path.join(scene, "videos", "chunk-000", level_key, + f"episode_{idx:06d}_{i}.jpg") for i in history + [t]] + lookdown = os.path.join(scene, "videos", "chunk-000", rgb_key, + f"episode_{idx:06d}_{t}.jpg") + if not all(os.path.isfile(p) for p in frames + [lookdown]): + continue + + action = actions[t] + samples.append({ + "episode": os.path.basename(scene), + # prompt_builder's "episode_idx" is the step index within the episode -- + # it derives the history frames from linspace(0, episode_idx-1) -- not the + # episode number. Passing the episode number yields one placeholder + # and a count mismatch against the history frames. + "episode_idx": t, + "episode_number": idx, + "instruction": (ep.get("tasks") or [""])[0], + "images": frames + [lookdown], + "turn": 2, + "assistant_turn1": LOOKDOWN_TOKEN, + "gt_goal": goals[t], + "gt_action": IDX2ACTION.get(int(action) if action is not None else -1), + }) + if len(samples) >= max_samples: + return samples + return samples + + +def evaluate(model_path: str, samples: list[dict], max_new_tokens: int, + device: str) -> dict: + model, processor, algo = load_for_eval(model_path, device=device) + # The prompt images are resized to 384x384, so predicted and GT pixel coordinates + # live in the same frame and the L2 is directly interpretable. + l2, n_parsed, n_action_ok, n_action_total = [], 0, 0, 0 + replies = [] + t0 = time.time() + + for i, sample in enumerate(samples, 1): + inputs = pb.build_sample_inputs(sample, processor).to(device) + with torch.inference_mode(): + out = model.generate(**inputs, max_new_tokens=max_new_tokens, + do_sample=False, temperature=None, top_p=None, top_k=None) + reply = processor.batch_decode(out[:, inputs["input_ids"].shape[1]:], + skip_special_tokens=True)[0].strip() + replies.append(reply) + + xy = parse_xy(reply) + if xy is not None and sample["gt_goal"] is not None: + n_parsed += 1 + gt = sample["gt_goal"] + l2.append(float(np.hypot(xy[0] - float(gt[0]), xy[1] - float(gt[1])))) + + # Only meaningful on turn 1. At turn 2 the model emits a coordinate by design, + # so scoring it against a movement action would always read 0% and look like a + # regression rather than the protocol working. + if sample.get("turn", 1) == 1 and sample["gt_action"] is not None: + n_action_total += 1 + if parse_action(reply) == sample["gt_action"]: + n_action_ok += 1 + + if i % 5 == 0 or i == len(samples): + print(f" [{i}/{len(samples)}] {time.time() - t0:.0f}s elapsed", flush=True) + + del model + torch.cuda.empty_cache() + + return { + "model_path": model_path, + "quant_algo": algo, + "n_samples": len(samples), + "pixel_goal_l2_mean": float(np.mean(l2)) if l2 else None, + "pixel_goal_l2_median": float(np.median(l2)) if l2 else None, + "pixel_goal_parse_rate": n_parsed / max(len(samples), 1), + "action_accuracy": n_action_ok / n_action_total if n_action_total else None, + "n_action_scored": n_action_total, + "seconds": round(time.time() - t0, 1), + "replies": replies, + } + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--model_path", required=True, action="append", + help="Checkpoint to score. Repeat to compare several on identical samples.") + p.add_argument("--data_root", required=True, help="Held-out LeRobot episodes") + p.add_argument("--num_samples", type=int, default=32) + p.add_argument("--max_new_tokens", type=int, default=32) + p.add_argument("--device", default="cuda") + p.add_argument("--seed", type=int, default=0) + p.add_argument("--camera", default=DEFAULT_CAMERA, + help="LeRobot camera setting to score against. Must be a pitched\n" + "one; the level 125cm_0deg carries no pixel goal.") + p.add_argument("--output_path", default=None, help="Write results JSON here") + return p.parse_args() + + +def main() -> int: + args = parse_args() + + samples = discover_samples(args.data_root, args.num_samples, seed=args.seed, + camera=args.camera) + if not samples: + print(f"[ERROR] no usable samples under {args.data_root}") + return 1 + print(f"Scoring {len(samples)} held-out samples from " + f"{len({s['episode'] for s in samples})} scene(s)\n") + + results = [] + for path in args.model_path: + print(f"=== {path}") + res = evaluate(path, samples, args.max_new_tokens, args.device) + results.append(res) + print(f" quant : {res['quant_algo'] or 'none (bf16)'}") + print(f" pixel_goal_l2 : mean {res['pixel_goal_l2_mean']:.2f} px, " + f"median {res['pixel_goal_l2_median']:.2f} px" + if res["pixel_goal_l2_mean"] is not None else " pixel_goal_l2 : n/a") + print(f" parse_rate : {100 * res['pixel_goal_parse_rate']:.1f}%") + if res["action_accuracy"] is not None and res["n_action_scored"]: + print(f" action_accuracy : {100 * res['action_accuracy']:.1f}% " + f"({res['n_action_scored']} scored)") + print(f" took : {res['seconds']}s\n") + + if len(results) > 1: + base = results[0] + print("=== delta vs " + os.path.basename(base["model_path"])) + for res in results[1:]: + name = os.path.basename(res["model_path"]) + if base["pixel_goal_l2_mean"] and res["pixel_goal_l2_mean"]: + d = res["pixel_goal_l2_mean"] - base["pixel_goal_l2_mean"] + print(f" {name}: pixel_goal_l2 {d:+.2f} px") + if (base["action_accuracy"] is not None and res["action_accuracy"] is not None + and base["n_action_scored"]): + d = 100 * (res["action_accuracy"] - base["action_accuracy"]) + print(f" {name}: action_accuracy {d:+.1f} pp") + agree = sum(a == b for a, b in zip(base["replies"], res["replies"])) + print(f" {name}: identical replies {agree}/{len(res['replies'])}") + + if args.output_path: + os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True) + with open(args.output_path, "w") as f: + json.dump(results, f, indent=2) + print(f"\nWrote {args.output_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 670b45947e7ef20d705af3871525e44c52311f7b Mon Sep 17 00:00:00 2001 From: hungho77 Date: Thu, 13 Aug 2026 16:58:36 +0700 Subject: [PATCH 06/30] docs(internvla-n1-dualvln): record measured FP8 task accuracy, add export/build script README now leads the results with the number a navigation model is actually judged on: FP8 keeps the median waypoint error unchanged (27.05 -> 26.97 px over 42 held-out samples) with 31 of 42 replies byte-identical and a median deviation of 0.00 px. The 2.88 px gap in the mean comes from 8 disagreeing samples rather than a systematic shift, so both statistics are given; reporting the mean alone would overstate the damage and the median alone would hide the tail. It also states plainly what the measurement does not cover -- activations are quantized in the engine but not on this PyTorch path, so the figure is a lower bound -- and that the published SR/SPL/NE table is a literature reference, not something reproduced here, since the closed-loop metrics need Habitat or InternUtopia and MP3D scenes. 03_export_build_system2.sh collapses the source project's four near-duplicate export scripts into one, driven by --no_quantization for the FP16 reference. Two things it does are load-bearing and carry comments saying so: emit_hidden_states=True before export, without which the engine returns logits only and the System 1 bridge cannot be evaluated; and __LUNOWUD=-peep:fc_h_fusion=off on the build, without which an FP16 engine on TensorRT 10.13 produces fluent gibberish from a Myelin miscompile. The visual encoder is built once at 4/4096/1024 image tokens. The source project's default of 512 came from a single-image demo and cannot hold a VLN prompt's 9-10 frames, which is why its multi-image verifications silently pointed at a hand-built engine no script produced. (cherry picked from commit 3d4e76441754009f8d74ff273c89e82a0d863330) --- recipes/internvla-n1-dualvln/README.md | 26 ++++ .../scripts/03_export_build_system2.sh | 137 ++++++++++++++++++ 2 files changed, 163 insertions(+) create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/scripts/03_export_build_system2.sh diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index 066e948..8770a40 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -35,6 +35,32 @@ this recipe is gated on `z_latents` cosine against the FP32 reference, not on ge ## Results +### Task accuracy — does the quantized model still pick the same waypoint? + +Measured with `quantize/benchmark_accuracy.py` on 42 held-out samples from two scenes, +identical samples and greedy decoding, PyTorch on Jetson Thor: + +| Checkpoint | pixel_goal_l2 mean | median | parse rate | +|---|---|---|---| +| System 2, BF16 (unquantized) | 47.24 px | **27.05 px** | 100 % | +| System 2, FP8 (s1) | 50.12 px | **26.97 px** | 100 % | + +31 of 42 replies are byte-identical and the **median deviation between the two is 0.00 px**. +The 2.88 px gap in the mean is carried by 8 samples, worst case 89 px — it is a small number +of disagreements, not a systematic shift, which is why both statistics are reported. For +reference, the source project's self-validation gate for this metric is a median under 60 px. + +Caveat worth keeping in view: this is a **weight-quantization** measurement. FP8 W8A8 also +quantizes activations, which the TensorRT engine does and the PyTorch path here does not, so +treat it as a lower bound on the engine's deviation rather than a prediction of it. + +The official InternVLA-N1 metrics (SR, SPL, NE, OS, nDTW) are all closed-loop and need +Habitat or InternUtopia plus MP3D scenes. Neither is installed here and neither is practical +on a Jetson, so the published table (DualVLN: NE 4.05 / SR 64.3 / SPL 58.5 on VLN-CE R2R) is +a literature reference point, not something this recipe reproduces. + +### Engine fidelity and latency + Measured on Jetson Thor, 12 held-out multi-image VLN steps: | LLM variant | z_latents vs FP32 | agrees w/ PyTorch | System 2 latency | LLM engine | diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/scripts/03_export_build_system2.sh b/recipes/internvla-n1-dualvln/trt-edgellm/scripts/03_export_build_system2.sh new file mode 100644 index 0000000..60f3424 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/scripts/03_export_build_system2.sh @@ -0,0 +1,137 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +# +# Export a System-2 checkpoint to ONNX and build the TensorRT-Edge-LLM engines. +# +# Works on both a quantized checkpoint and the unquantized one; pass --no_quantization +# for the latter to get the FP16 fidelity reference. +# +# Two details here are load-bearing and must not be "cleaned up": +# +# * CausalLM.emit_hidden_states = True before the export. The System 2 -> System 1 +# bridge reads the last-layer hidden states of the trajectory tokens; without this the +# engine emits logits only and the bridge cannot be evaluated at all. +# +# * __LUNOWUD=-peep:fc_h_fusion=off on the engine build. TensorRT 10.13 miscompiles +# Myelin's horizontal fusion of the gate/up projections on sm_110 at batch 1, and an +# FP16 engine built without this emits fluent-looking gibberish. TensorRT-Edge-LLM +# already disables the fusion, but only for TensorRT >= 10.15, so 10.13 falls through +# the gap. FP8 happens to dodge the same bug because its Q/DQ nodes break the fusion +# pattern, which is why FP8 once looked mandatory -- it is not. +# +# The visual encoder is sized for multi-image VLN prompts (~1764 image tokens across 9-10 +# frames). The single-image demo default of 512 max image tokens cannot hold one. + +set -euo pipefail + +MODEL_PATH="" +ENGINE_DIR="" +ONNX_DIR="" +TRT_EDGELLM_DIR="${TRT_EDGELLM_DIR:-$HOME/modelopt/TensorRT-Edge-LLM}" +NO_QUANTIZATION=0 +SKIP_VISUAL=0 +MAX_BATCH_SIZE="${MAX_BATCH_SIZE:-1}" +MAX_INPUT_LEN="${MAX_INPUT_LEN:-3072}" +MAX_KV_CACHE="${MAX_KV_CACHE:-4096}" +LUNOWUD_WAR="${LUNOWUD_WAR:--peep:fc_h_fusion=off}" + +while [[ $# -gt 0 ]]; do + case "$1" in + --model_path) MODEL_PATH="$2"; shift 2 ;; + --engine_dir) ENGINE_DIR="$2"; shift 2 ;; + --onnx_dir) ONNX_DIR="$2"; shift 2 ;; + --trt_edgellm_dir) TRT_EDGELLM_DIR="$2"; shift 2 ;; + --no_quantization) NO_QUANTIZATION=1; shift ;; + --skip_visual) SKIP_VISUAL=1; shift ;; + --max_batch_size) MAX_BATCH_SIZE="$2"; shift 2 ;; + --max_input_len) MAX_INPUT_LEN="$2"; shift 2 ;; + --max_kv_cache) MAX_KV_CACHE="$2"; shift 2 ;; + -h|--help) + sed -n '3,25p' "$0" | sed 's/^# \{0,1\}//' + exit 0 ;; + *) echo "[ERROR] unknown argument: $1" >&2; exit 1 ;; + esac +done + +[[ -n "$MODEL_PATH" ]] || { echo "[ERROR] --model_path is required" >&2; exit 1; } +[[ -n "$ENGINE_DIR" ]] || { echo "[ERROR] --engine_dir is required" >&2; exit 1; } +[[ -d "$MODEL_PATH" ]] || { echo "[ERROR] model path not found: $MODEL_PATH" >&2; exit 1; } + +ONNX_DIR="${ONNX_DIR:-${ENGINE_DIR%/}_onnx}" +PLUGIN="$TRT_EDGELLM_DIR/build/libNvInfer_edgellm_plugin.so" +LLM_BUILD="$TRT_EDGELLM_DIR/build/examples/llm/llm_build" +VISUAL_BUILD="$TRT_EDGELLM_DIR/build/examples/multimodal/visual_build" + +for f in "$PLUGIN" "$LLM_BUILD"; do + [[ -f "$f" ]] || { echo "[ERROR] not found: $f" >&2 + echo " set TRT_EDGELLM_DIR or build TensorRT-Edge-LLM first" >&2 + exit 1; } +done + +export EDGELLM_PLUGIN_PATH="$PLUGIN" +export __LUNOWUD="${__LUNOWUD:+$__LUNOWUD }$LUNOWUD_WAR" + +echo "==============================================================" +echo " InternVLA-N1 System 2 — ONNX export and engine build" +echo "==============================================================" +echo " Model path: $MODEL_PATH" +echo " Quantized: $([[ $NO_QUANTIZATION -eq 1 ]] && echo 'no (FP16 reference)' || echo yes)" +echo " ONNX dir: $ONNX_DIR" +echo " Engine dir: $ENGINE_DIR" +echo " TensorRT-Edge: $TRT_EDGELLM_DIR" +echo " maxBatchSize: $MAX_BATCH_SIZE" +echo " maxInputLen: $MAX_INPUT_LEN" +echo " maxKVCacheCap: $MAX_KV_CACHE" +echo " __LUNOWUD: $__LUNOWUD" +echo "==============================================================" + +mkdir -p "$ONNX_DIR" "$ENGINE_DIR" + +echo +echo "[1/3] Exporting ONNX (emit_hidden_states=True for the System 1 bridge)..." +python - "$MODEL_PATH" "$ONNX_DIR" <<'PYEOF' +import sys +from tensorrt_edgellm.models.default.modeling_default import CausalLM +# The bridge to System 1 reads the last-layer hidden states, not just logits. +CausalLM.emit_hidden_states = True +from tensorrt_edgellm.scripts.export import main +sys.argv = ["tensorrt-edgellm-export", sys.argv[1], sys.argv[2]] +sys.exit(main()) +PYEOF + +echo +echo "[2/3] Building the LLM engine..." +mkdir -p "$ENGINE_DIR/llm" +"$LLM_BUILD" \ + --onnxDir "$ONNX_DIR/llm" \ + --engineDir "$ENGINE_DIR/llm" \ + --maxBatchSize "$MAX_BATCH_SIZE" \ + --maxInputLen "$MAX_INPUT_LEN" \ + --maxKVCacheCapacity "$MAX_KV_CACHE" + +if [[ $SKIP_VISUAL -eq 0 && -d "$ONNX_DIR/visual" ]]; then + echo + echo "[3/3] Building the visual engine..." + [[ -f "$VISUAL_BUILD" ]] || { echo "[ERROR] not found: $VISUAL_BUILD" >&2; exit 1; } + mkdir -p "$ENGINE_DIR/visual" + # A VLN prompt carries 9-10 frames and roughly 1764 image tokens. The + # single-image demo default of 512 cannot hold one. + "$VISUAL_BUILD" \ + --onnxDir "$ONNX_DIR/visual" \ + --engineDir "$ENGINE_DIR/visual" \ + --minImageTokens 4 \ + --maxImageTokens 4096 \ + --maxImageTokensPerImage 1024 +else + echo + echo "[3/3] Skipping the visual engine." +fi + +echo +echo "==============================================================" +echo " Build complete" +echo "==============================================================" +du -sh "$ENGINE_DIR"/* 2>/dev/null | sed 's/^/ /' +echo " Engines: $ENGINE_DIR" +echo " ONNX kept at $ONNX_DIR (safe to delete once the engines verify)" From 4de9290e7c68c1f7ac69d4e5274dc4785a9fde68 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Thu, 13 Aug 2026 17:04:14 +0700 Subject: [PATCH 07/30] fix(internvla-n1-dualvln): correct visual engine path and emit base_config.json Two build-script defects found by running it rather than reading it. visual_build appends its own 'visual' component directory under --engineDir, so passing $ENGINE_DIR/visual produced $ENGINE_DIR/visual/visual/visual.engine. Every consumer looks for $ENGINE_DIR/visual/visual.engine, so the engine was built correctly and then filed where nothing would find it. Pass the parent instead. llm_bench reads its engine configuration from base_config.json while llm_build writes config.json -- same content, different name. Without the copy the benchmark refuses to open the engine that was just handed to it, with a parse error that reads like engine corruption rather than a missing filename. Verified on the FP8 engine built from the repackaged System 2: llm.engine 7.62 GB (the figure the source project recorded), visual.engine 1.36 GB, text smoke test coherent, and llm_bench measuring 95.8 ms prefill at 1024 tokens and 37.35 ms decode at 1024 past-KV. (cherry picked from commit 0eb6e913b2605eead1e3621c70f4b36ef6cab884) --- .../scripts/03_export_build_system2.sh | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/scripts/03_export_build_system2.sh b/recipes/internvla-n1-dualvln/trt-edgellm/scripts/03_export_build_system2.sh index 60f3424..b0d6d78 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/scripts/03_export_build_system2.sh +++ b/recipes/internvla-n1-dualvln/trt-edgellm/scripts/03_export_build_system2.sh @@ -110,16 +110,24 @@ mkdir -p "$ENGINE_DIR/llm" --maxInputLen "$MAX_INPUT_LEN" \ --maxKVCacheCapacity "$MAX_KV_CACHE" +# llm_bench reads its engine configuration from base_config.json, while llm_build writes +# config.json. Same content, different name; without this copy the benchmark cannot open +# the engine it was just handed. +cp "$ENGINE_DIR/llm/config.json" "$ENGINE_DIR/llm/base_config.json" + if [[ $SKIP_VISUAL -eq 0 && -d "$ONNX_DIR/visual" ]]; then echo echo "[3/3] Building the visual engine..." [[ -f "$VISUAL_BUILD" ]] || { echo "[ERROR] not found: $VISUAL_BUILD" >&2; exit 1; } - mkdir -p "$ENGINE_DIR/visual" - # A VLN prompt carries 9-10 frames and roughly 1764 image tokens. The - # single-image demo default of 512 cannot hold one. + # visual_build appends its own "visual" component directory under --engineDir, so + # pass the parent: giving it $ENGINE_DIR/visual yields $ENGINE_DIR/visual/visual and + # every consumer that expects $ENGINE_DIR/visual/visual.engine then misses it. + # + # A VLN prompt carries 9-10 frames and roughly 1764 image tokens. The single-image + # demo default of 512 cannot hold one. "$VISUAL_BUILD" \ --onnxDir "$ONNX_DIR/visual" \ - --engineDir "$ENGINE_DIR/visual" \ + --engineDir "$ENGINE_DIR" \ --minImageTokens 4 \ --maxImageTokens 4096 \ --maxImageTokensPerImage 1024 From 7e201a616bd6adcbc144352ae7f8fc900879203c Mon Sep 17 00:00:00 2001 From: hungho77 Date: Thu, 13 Aug 2026 17:09:35 +0700 Subject: [PATCH 08/30] docs(internvla-n1-dualvln): record measured FP8 vs FP16 engine comparison Both engines built from the same repackaged System 2 and measured with llm_bench on Jetson Thor at batch 1: LLM engine prefill 1024 decode pastKV 1024 base FP16 15.0 GB 196.22 ms 86.16 ms FP8 7.62 GB 95.80 ms 37.35 ms gain 1.97x smaller 2.05x faster 2.31x faster Both emit the same text on the same prompt, and the PyTorch-side accuracy run already showed the median waypoint error unchanged, so FP8 here is a straight win rather than a trade: half the engine, roughly double the throughput, same decision. Worth recording that the FP16 engine is correct only because the build applies __LUNOWUD=-peep:fc_h_fusion=off; the build log confirms both workaround flags were passed. Without it TensorRT 10.13 miscompiles Myelin's horizontal gate/up fusion on sm_110 and the engine produces fluent gibberish, which is what once made FP8 look mandatory on this platform. It is not -- FP8 is chosen here on size and speed. (cherry picked from commit 997cdf66dc8f928049d4b40f67dac9cd4b08153a) --- recipes/internvla-n1-dualvln/README.md | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index 8770a40..09e9405 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -59,7 +59,27 @@ Habitat or InternUtopia plus MP3D scenes. Neither is installed here and neither on a Jetson, so the published table (DualVLN: NE 4.05 / SR 64.3 / SPL 58.5 on VLN-CE R2R) is a literature reference point, not something this recipe reproduces. -### Engine fidelity and latency +### Engine size and latency + +Both engines built from the same repackaged System 2 and measured with `llm_bench` on +Jetson Thor, batch 1: + +| | LLM engine | visual engine | prefill (1024 tok) | decode (pastKV 1024) | +|---|---|---|---|---| +| base FP16 (unquantized) | 15.0 GB | 1.36 GB | 196.22 ms | 86.16 ms | +| FP8 (s1) | **7.62 GB** | 1.36 GB | **95.80 ms** | **37.35 ms** | +| FP8 gain | 1.97x smaller | — | **2.05x faster** | **2.31x faster** | + +Both produce the same text on the same prompt, so this is a straight win: FP8 halves the +engine and roughly doubles throughput while leaving the median waypoint error unchanged. + +The FP16 engine is only correct because the build applies +`__LUNOWUD=-peep:fc_h_fusion=off`. The build log confirms it +(`Using __LUNOWUD=-peep:fc_h_fusion=off -peep:match_dual_gemm=off`). Without it, TensorRT +10.13 miscompiles Myelin's horizontal gate/up fusion on sm_110 and the engine emits fluent +gibberish — see the note further down. + +### Earlier full-pipeline figures Measured on Jetson Thor, 12 held-out multi-image VLN steps: From 91a71c26b55cd264ea09206840063c0427eefa37 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Thu, 13 Aug 2026 17:22:57 +0700 Subject: [PATCH 09/30] feat(internvla-n1-dualvln): verify the System 2 -> System 1 bridge on both engines verify_latents.py rebuilds the bridge end to end against a PyTorch BF16 reference -- embed, scatter image embeddings, append the four trajectory tokens, run the engine with hand-built 3D mRoPE, take the last-layer hidden states, apply the host-side norm and cond_projector -- and compares z_latents. engine hidden pre-norm post-norm z_latents rel-L2 base FP16 0.999843 0.999123 0.999471 0.0293 FP8 (s1) 0.997793 0.987343 0.991861 0.1023 Both clear the 0.99 gate, so FP8 preserves the signal System 1 actually consumes. This is the check NVFP4 fails at 0.647 while still producing fluent text, which is why it ships as experimental rather than supported. The check reads latent_queries and cond_projector from the bridge.safetensors that repackage_system2.py sets aside, so it needs neither the 16 GB original checkpoint nor InternNav. Four transformers 5.x incompatibilities were fixed with version-probing helpers rather than a version pin, since the recipe should survive either: the vision tower and get_rope_index moved onto the inner model; get_image_features returns a BaseModelOutputWithPooling whose pooler_output holds the merged embeddings while last_hidden_state is the pre-merger tensor at vision width; and get_rope_index gained a required mm_token_type_ids argument. That last one does not fail on arity -- it fails later inside with 'NoneType is not an iterator', which is a poor way to discover it. Also switched the reference forward from flash_attention_2 to sdpa, which is what is available on Jetson and what the deployed agent uses. (cherry picked from commit b9b063f653f886bd0bfdc89db8017be99439f542) --- recipes/internvla-n1-dualvln/README.md | 15 ++ .../trt-edgellm/engine_runner.py | 215 ++++++++++++++++++ .../trt-edgellm/verify/verify_latents.py | 197 ++++++++++++++++ .../trt-edgellm/verify/verify_latents_vln.py | 215 ++++++++++++++++++ 4 files changed, 642 insertions(+) create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/engine_runner.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents_vln.py diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index 09e9405..4076958 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -79,6 +79,21 @@ The FP16 engine is only correct because the build applies 10.13 miscompiles Myelin's horizontal gate/up fusion on sm_110 and the engine emits fluent gibberish — see the note further down. +### Bridge fidelity — z_latents + +`trt-edgellm/verify/verify_latents.py` reconstructs the System 2 -> System 1 bridge against +a PyTorch BF16 reference: embed, scatter image embeddings, append the 4 trajectory tokens, +run the engine with hand-built 3D mRoPE, take the last-layer hidden states, then apply the +host-side norm and `cond_projector`. + +| Engine | hidden pre-norm | hidden post-norm | **z_latents** | rel-L2 | +|---|---|---|---|---| +| base FP16 | 0.999843 | 0.999123 | **0.999471** | 0.0293 | +| FP8 (s1) | 0.997793 | 0.987343 | **0.991861** | 0.1023 | + +Both pass the > 0.99 gate. This is the check that NVFP4 fails (0.647), and it is the reason +it ships as experimental: text stays fluent while the waypoint bridge collapses. + ### Earlier full-pipeline figures Measured on Jetson Thor, 12 held-out multi-image VLN steps: diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/engine_runner.py b/recipes/internvla-n1-dualvln/trt-edgellm/engine_runner.py new file mode 100644 index 0000000..82f73f8 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/engine_runner.py @@ -0,0 +1,215 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Shared TensorRT engine runner + mRoPE table builder for the System 2 LLM engine. + +Drives the FP8 LLM engine directly from Python: builds the 3D mRoPE cos/sin table from the +reference `position_ids` (per-token, merged by mrope_section (16,24,24), layout [cos64|sin64]), +binds torch GPU buffers to the TensorRT context, and returns (logits, hidden_states) for a +single prefill. Used by the verification scripts. Select an engine via the ENGINE_PATH env var. +""" +import os, sys, ctypes, json +_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, _ROOT) +# TensorRT ships with JetPack outside the venv. +sys.path.append(os.environ.get("SYSTEM_SITE", "/usr/lib/python3.12/dist-packages")) +import numpy as np, tensorrt as trt, torch +from PIL import Image + +def _env(name, default): + return os.path.expanduser(os.environ.get(name, default)) + + +WORK_DIR = _env("WORK_DIR", "~/vln-opt-work") +REPKG = _env("REPKG_CKPT", os.path.join(WORK_DIR, "qwen25vl_system2")) +ENGINE = _env("ENGINE_PATH", os.path.join(WORK_DIR, "engines/s1_fp8/llm/llm.engine")) +TRT_EDGELLM_DIR = _env("TRT_EDGELLM_DIR", "~/modelopt/TensorRT-Edge-LLM") +PLUGIN = _env("EDGELLM_PLUGIN_PATH", + os.path.join(TRT_EDGELLM_DIR, "build/libNvInfer_edgellm_plugin.so")) +# The bridge tensors are emitted next to the repackaged checkpoint by +# repackage_system2.py, so the 16 GB original is not needed here. +CKPT = _env("INTERNVLA_CKPT", REPKG) +IMAGE = os.path.expanduser(os.environ.get( + "IMAGE_PATH", os.path.join(TRT_EDGELLM_DIR, + "examples/multimodal/pics/giant_panda.jpeg"))) + +# Qwen2.5-VL-7B / InternVLA-N1 System 2 geometry. Inlined rather than imported: this +# repository has no shared-library convention, each recipe stands alone. +THETA = 1_000_000.0 +HEAD_DIM = 128 +N_LAYERS = 28 +N_KV = 4 +HIDDEN = 3584 +ROPE_MAXPOS = 4096 # must match --maxKVCacheCapacity at build time +MROPE_SECTION = [16, 24, 24] # (T, H, W) +TRT2TORCH = {trt.DataType.HALF: torch.float16, trt.DataType.FLOAT: torch.float32, + trt.DataType.INT32: torch.int32, trt.DataType.INT64: torch.int64, + trt.DataType.BF16: torch.bfloat16} +_ENG = {} + + +def cos(a, b): + a, b = a.flatten().float(), b.flatten().float() + return torch.nn.functional.cosine_similarity(a, b, dim=0).item() + + +def per_tok_cos(a, b): + a, b = a[0].float(), b[0].float() + return torch.nn.functional.cosine_similarity(a, b, dim=-1).mean().item() + + +def build_mrope_table(position_ids, device): + """position_ids [3,1,S] → rope table [1, ROPE_MAXPOS, 128] (rows 0..S-1 are per-token + mRoPE cos/sin merged by mrope_section; layout [cos64|sin64]).""" + S = position_ids.shape[-1] + half = HEAD_DIM // 2 # 64 + zid = torch.arange(half, dtype=torch.float32, device=device) + inv_freq = THETA ** (-2.0 * zid / HEAD_DIM) # [64] + # axis index per freq band: [0]*16 + [1]*24 + [2]*24 + axis = torch.cat([torch.full((MROPE_SECTION[i],), i, device=device) for i in range(3)]) # [64] + pos = position_ids[:, 0, :].float() # [3,S] + # per token s, freq j: angle = pos[axis[j], s] * inv_freq[j] + pos_sel = pos[axis] # [64,S] + ang = pos_sel.T[:, :] * inv_freq[None, :] # [S,64] + c, s = torch.cos(ang), torch.sin(ang) # [S,64] + table = torch.zeros(1, ROPE_MAXPOS, HEAD_DIM, dtype=torch.float32, device=device) + table[0, :S, :half] = c + table[0, :S, half:] = s + return table + + +def load_engine(): + if "eng" not in _ENG: + ctypes.CDLL(PLUGIN, mode=ctypes.RTLD_GLOBAL) + lg = trt.Logger(trt.Logger.ERROR); trt.init_libnvinfer_plugins(lg, "") + rt = trt.Runtime(lg) + with open(ENGINE, "rb") as f: + _ENG["eng"] = rt.deserialize_cuda_engine(f.read()); _ENG["rt"] = rt + return _ENG["eng"] + + +def run_engine(embeds_half, rope_table): + S = embeds_half.shape[1]; dev = embeds_half.device + eng = load_engine(); ctx = eng.create_execution_context() + ctx.set_optimization_profile_async(0, torch.cuda.current_stream().cuda_stream) + context_lengths = torch.tensor([S], dtype=torch.int32, device=dev) + kvcache_start = torch.zeros(1, dtype=torch.int32, device=dev) + last_token = torch.tensor([[S - 1]], dtype=torch.int64, device=dev) + kv_cap = ROPE_MAXPOS + kv_cache = [torch.zeros(1, 2, N_KV, kv_cap, HEAD_DIM, dtype=torch.float16, device=dev) + for _ in range(N_LAYERS)] + feed = { + "inputs_embeds": (embeds_half.contiguous(), None), + "rope_rotary_cos_sin": (rope_table.contiguous(), None), + "context_lengths": (context_lengths, None), + "kvcache_start_index": (kvcache_start, (0,)), + "last_token_ids": (last_token, None), + } + for i in range(N_LAYERS): + feed[f"past_key_values_{i}"] = (kv_cache[i], (1, 2, N_KV, kv_cap, HEAD_DIM)) + for name, (t, shp) in feed.items(): + ctx.set_input_shape(name, shp if shp else tuple(t.shape)) + ctx.set_tensor_address(name, t.data_ptr()) + outs = {} + for i in range(eng.num_io_tensors): + n = eng.get_tensor_name(i) + if eng.get_tensor_mode(n) != trt.TensorIOMode.OUTPUT: + continue + if n.startswith("present_key_values_"): + li = int(n.rsplit("_", 1)[1]) + ctx.set_tensor_address(n, kv_cache[li].data_ptr()); continue + shp = tuple(int(d) for d in ctx.get_tensor_shape(n)) + shp = tuple(S if d < 0 else d for d in shp) + t = torch.empty(shp, dtype=TRT2TORCH[eng.get_tensor_dtype(n)], device=dev) + outs[n] = t; ctx.set_tensor_address(n, t.data_ptr()) + outs["_kv"] = kv_cache + ok = ctx.execute_async_v3(torch.cuda.current_stream().cuda_stream) + torch.cuda.synchronize(); assert ok + return outs["logits"], outs["hidden_states"] + + +def main(): + dev = "cuda"; torch.manual_seed(0) + from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration + print(f"[1/5] Load repackage (transformers) | engine={os.path.basename(ENGINE)}") + model = Qwen2_5_VLForConditionalGeneration.from_pretrained( + REPKG, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2", + low_cpu_mem_usage=True).to(dev).eval() + proc = AutoProcessor.from_pretrained(REPKG, trust_remote_code=True, + min_pixels=128*28*28, max_pixels=1024*28*28) + backbone = model.model.language_model if hasattr(model.model, "language_model") else model.model + final_norm = backbone.norm + + msgs = [{"role": "user", "content": [ + {"type": "image", "image": IMAGE}, + {"type": "text", "text": "Please describe the image."}]}] + text = proc.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) + inp = proc(text=[text], images=[Image.open(IMAGE).convert("RGB")], + return_tensors="pt").to(dev) + S = inp["input_ids"].shape[1] + print(f" seq_len={S} (image grid_thw={inp['image_grid_thw'].tolist()})") + + # hook inner LM to capture the reference inputs_embeds + position_ids + cap = {} + def pre_hook(mod, args, kwargs): + cap["inputs_embeds"] = kwargs.get("inputs_embeds") + cap["position_ids"] = kwargs.get("position_ids") + h1 = backbone.register_forward_pre_hook(pre_hook, with_kwargs=True) + def norm_hook(mod, i, o): + cap["pre"] = i[0].detach(); cap["post"] = o.detach() + h2 = final_norm.register_forward_hook(norm_hook) + + print("[2/5] Reference forward (WITH image)") + with torch.no_grad(): + ref_out = model(**inp, use_cache=False) + h1.remove(); h2.remove() + embeds = cap["inputs_embeds"]; pos = cap["position_ids"] + if embeds is None: + # fallback: build inputs_embeds externally + raise RuntimeError("inputs_embeds not captured - wrong hook target") + print(f" captured inputs_embeds {tuple(embeds.shape)} position_ids {tuple(pos.shape)}") + print(f" position_ids axis-range T[{pos[0].min()}..{pos[0].max()}] " + f"H[{pos[1].min()}..{pos[1].max()}] W[{pos[2].min()}..{pos[2].max()}]") + ref_pre = cap["pre"].float(); ref_post = cap["post"].float() + ref_logits_last = ref_out.logits[:, -1, :].float() + + print("[3/5] Build merged 3D mRoPE + run engine") + rope = build_mrope_table(pos, dev) + eng_logits, eng_hs = run_engine(embeds.to(torch.float16), rope) + eng_hs = eng_hs.float() + with torch.no_grad(): + eng_post = final_norm(eng_hs.to(torch.bfloat16)).float() + + print("[4/5] Compare hidden_states / logits\n" + "=" * 60) + a, b = eng_hs[0].float(), ref_pre[0].float() + ptc = torch.nn.functional.cosine_similarity(a, b, dim=-1) + print(" per-pos PRE cosine[:8]:", " ".join(f"{ptc[i]:.3f}" for i in range(min(S, 8))), + "... last:", f"{ptc[-1]:.3f}") + print(f" hidden PRE-norm per-token cosine = {per_tok_cos(eng_hs, ref_pre):.6f}") + print(f" hidden POST-norm per-token cosine = {per_tok_cos(eng_post, ref_post):.6f}") + ea, ra = eng_logits[0, 0].argmax().item(), ref_logits_last[0].argmax().item() + print(f" logits last cosine = {cos(eng_logits[0,0], ref_logits_last[0]):.6f} " + f"| argmax eng={ea!r} ref={ra!r} match={ea==ra}") + + print("[5/5] z_latents (cond_projector GỐC) engine vs reference") + from safetensors import safe_open + idx = json.load(open(os.path.join(CKPT, "model.safetensors.index.json")))["weight_map"] + cp = {} + for k in idx: + if k.startswith("model.cond_projector"): + with safe_open(os.path.join(CKPT, idx[k]), framework="pt") as f: + cp[k.replace("model.cond_projector.", "")] = f.get_tensor(k).float().to(dev) + def cond_project(x): + x = torch.nn.functional.linear(x, cp["0.weight"], cp.get("0.bias")) + x = torch.nn.functional.gelu(x) + return torch.nn.functional.linear(x, cp["2.weight"], cp.get("2.bias")) + z_eng, z_ref = cond_project(eng_post), cond_project(ref_post) + zc = per_tok_cos(z_eng, z_ref) + print(f" z_latents({z_eng.shape[-1]}) per-token cosine = {zc:.6f}") + print("\n" + ("✅ WITH-IMAGE numeric PASS (z_latents ≥ 0.99)" if zc > 0.99 + else f"z_latents cosine {zc:.4f} < 0.99 - check mRoPE/vision")) + return 0 if zc > 0.99 else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents.py new file mode 100644 index 0000000..d9d0035 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents.py @@ -0,0 +1,197 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Verify System 2 numeric fidelity: z_latents cosine of the FP8 LLM engine vs the +PyTorch reference, using the real latent-query bridge path. +""" +import os, sys, json +_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, _R); sys.path.insert(0, _R) +import torch +from PIL import Image +from engine_runner import (REPKG, ENGINE, CKPT, build_mrope_table, run_engine, + cos, per_tok_cos) + +TRAJ_TOKEN_INDEX = 151667 +IMAGE_TOKEN_INDEX = 151655 +N_QUERY = 4 +IMAGE = os.path.expanduser(os.environ.get( + "IMAGE_PATH", "~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg")) + + +def load_ckpt_tensor(prefix): + """Read bridge tensors, preferring the small bridge.safetensors next to the checkpoint. + + repackage_system2.py sets latent_queries and cond_projector aside in a ~25 MB + bridge.safetensors precisely so this check does not have to reopen the 16 GB original. + Falling back to a full sharded checkpoint keeps the script usable against one. + """ + from safetensors import safe_open + + bridge = os.path.join(CKPT, "bridge.safetensors") + if os.path.isfile(bridge): + with safe_open(bridge, framework="pt") as f: + return {k: f.get_tensor(k) for k in f.keys() if k.startswith(prefix)} + + index = os.path.join(CKPT, "model.safetensors.index.json") + if not os.path.isfile(index): + raise FileNotFoundError( + f"neither bridge.safetensors nor model.safetensors.index.json under {CKPT}; " + f"point INTERNVLA_CKPT at a repackaged checkpoint or the original") + idx = json.load(open(index))["weight_map"] + out = {} + for k in idx: + if k.startswith(prefix): + with safe_open(os.path.join(CKPT, idx[k]), framework="pt") as f: + out[k] = f.get_tensor(k) + return out + + +def rope_index(model, input_ids, image_grid_thw, image_token_id): + """Call get_rope_index across the two transformers signatures. + + transformers 4.x took (input_ids, image_grid_thw). 5.x inserted a required + mm_token_type_ids argument -- 0 text, 1 image, 2 video -- and moved the grids to + keywords. Passing the 4.x form to a 5.x model does not raise on arity; it fails later + inside with 'NoneType is not an iterator', which is a confusing way to learn this. + """ + import inspect + + fn = model_attr(model, "get_rope_index") + params = inspect.signature(fn).parameters + if "mm_token_type_ids" in params: + mm = (input_ids == image_token_id).to(torch.int32) + return fn(input_ids, mm, image_grid_thw=image_grid_thw) + return fn(input_ids, image_grid_thw) + + +def model_attr(model, name): + """Fetch an attribute or bound method from the model or its inner model. + + transformers 5.x moved several Qwen2.5-VL helpers (get_rope_index, visual, ...) from + the ForConditionalGeneration wrapper down onto the inner Qwen2_5_VLModel. Probing both + keeps this working on either layout instead of pinning a transformers version. + """ + for owner in (model, getattr(model, "model", None)): + if owner is not None and hasattr(owner, name): + return getattr(owner, name) + raise AttributeError(f"neither the model nor its inner model has {name!r}") + + +def get_visual(model): + """Return the vision tower across transformers layouts. + + Older versions expose it as ``model.visual``; newer ones nest it under + ``model.model.visual``. Probing beats pinning a transformers version here. + """ + for owner in (model, getattr(model, "model", None)): + vis = getattr(owner, "visual", None) if owner is not None else None + if vis is not None: + return vis + raise AttributeError("no vision tower found on this model " + "(looked at .visual and .model.visual)") + + +def visual_embeds(model, pixel_values, grid_thw): + """Return merged image embeddings at LLM hidden width, across transformers layouts. + + In transformers 5.x get_image_features returns a BaseModelOutputWithPooling whose + ``pooler_output`` holds the *merged* embeddings (split per image), while + ``last_hidden_state`` is the pre-merger tensor at vision width -- 1280 here against the + LLM's 3584. Reading the wrong field fails loudly on the shape, but only after the + vision tower has already run, so unwrap explicitly. + """ + if hasattr(model, "get_image_features"): + out = model.get_image_features(pixel_values.type(get_visual(model).dtype), grid_thw) + feats = getattr(out, "pooler_output", out) + else: + vis = get_visual(model) + out = vis(pixel_values.type(vis.dtype), grid_thw=grid_thw) + feats = getattr(out, "last_hidden_state", out) + if isinstance(feats, (list, tuple)): + feats = torch.cat([f.reshape(-1, f.shape[-1]) for f in feats], dim=0) + return feats.reshape(-1, feats.shape[-1]) + + +def main(): + dev = "cuda" + from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration + print(f"[1/6] Load repackage + processor | engine={os.path.basename(ENGINE)}") + model = Qwen2_5_VLForConditionalGeneration.from_pretrained( + REPKG, torch_dtype=torch.bfloat16, # flash-attn is not available on Jetson; sdpa is the supported path and + # is what the deployed agent uses too. + attn_implementation=os.environ.get("ATTN_IMPL", "sdpa"), + low_cpu_mem_usage=True).to(dev).eval() + proc = AutoProcessor.from_pretrained(REPKG, trust_remote_code=True, + min_pixels=128*28*28, max_pixels=1024*28*28) + inner = model.model # Qwen2_5_VLModel (takes inputs_embeds) + lm = inner.language_model if hasattr(inner, "language_model") else inner + final_norm = lm.norm + + print("[2/6] latent_queries + cond_projector from the InternVLA checkpoint") + lq = load_ckpt_tensor("model.latent_queries")["model.latent_queries"].to(dev).to(torch.bfloat16) + cpw = load_ckpt_tensor("model.cond_projector") + cp = {k.replace("model.cond_projector.", ""): v.float().to(dev) for k, v in cpw.items()} + print(f" latent_queries {tuple(lq.shape)} (n_query={lq.shape[1]}) cond keys {sorted(cp.keys())}") + assert lq.shape[1] == N_QUERY + + print("[3/6] Build image + instruction input, append 4 TRAJ tokens (as generate_latents)") + msgs = [{"role": "user", "content": [ + {"type": "image", "image": IMAGE}, + {"type": "text", "text": "Go straight then stop at the green plant."}]}] + text = proc.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) + enc = proc(text=[text], images=[Image.open(IMAGE).convert("RGB")], return_tensors="pt").to(dev) + input_ids = enc["input_ids"] + grid = enc["image_grid_thw"] + # embed + scatter image + append latent_queries + with torch.no_grad(): + text_embeds = model.get_input_embeddings()(input_ids) # [1,S,3584] + image_embeds = visual_embeds(model, enc["pixel_values"], grid) + image_idx = (input_ids == IMAGE_TOKEN_INDEX) + text_embeds[image_idx] = image_embeds.to(text_embeds.dtype)[: image_idx.sum(), :] + inputs_embeds = torch.cat([text_embeds, lq.repeat(text_embeds.shape[0], 1, 1)], dim=1) + ids_traj = torch.cat( + [input_ids, torch.tensor([[TRAJ_TOKEN_INDEX] * N_QUERY], device=dev)], dim=1) + S = inputs_embeds.shape[1] + position_ids, _ = rope_index(model, ids_traj, grid, IMAGE_TOKEN_INDEX) + print(f" seq_len={S} (+{N_QUERY} traj) grid={grid.tolist()} pos {tuple(position_ids.shape)}") + + print("[4/6] Reference forward (inner LM); capture pre/post norm at the last 4 positions") + cap = {} + h = final_norm.register_forward_hook(lambda m, i, o: cap.update(pre=i[0].detach(), post=o.detach())) + with torch.no_grad(): + inner(inputs_embeds=inputs_embeds, position_ids=position_ids, + output_hidden_states=True, return_dict=True) + h.remove() + ref_pre = cap["pre"][:, -N_QUERY:, :].float() + ref_post = cap["post"][:, -N_QUERY:, :].float() + + print("[5/6] Engine forward (3D mRoPE); take hidden[-4:] pre-norm -> host norm") + rope = build_mrope_table(position_ids, dev) + _, eng_hs = run_engine(inputs_embeds.to(torch.float16), rope) + eng_pre = eng_hs[:, -N_QUERY:, :].float() + with torch.no_grad(): + eng_post = final_norm(eng_pre.to(torch.bfloat16)).float() + + print("[6/6] Compare z_latents (System1 traj_dit input)\n" + "=" * 60) + def cond_project(x): + x = torch.nn.functional.linear(x, cp["0.weight"], cp.get("0.bias")) + x = torch.nn.functional.gelu(x) + return torch.nn.functional.linear(x, cp["2.weight"], cp.get("2.bias")) + z_ref, z_eng = cond_project(ref_post), cond_project(eng_post) + ptc = torch.nn.functional.cosine_similarity(eng_pre[0], ref_pre[0], dim=-1) + print(" per-query PRE cosine:", " ".join(f"{ptc[i]:.4f}" for i in range(N_QUERY))) + print(f" hidden PRE-norm cosine = {per_tok_cos(eng_pre, ref_pre):.6f}") + print(f" hidden POST-norm cosine = {per_tok_cos(eng_post, ref_post):.6f}") + zc = per_tok_cos(z_eng, z_ref) + zc_flat = cos(z_eng, z_ref) + l2 = (z_eng - z_ref).norm() / z_ref.norm() + print(f" z_latents({z_eng.shape[-1]}) per-query cosine = {zc:.6f} | flat = {zc_flat:.6f} " + f"| rel-L2 = {l2:.4f}") + ok = zc > 0.99 + print("\n" + (f"FAITHFUL e2e bridge PASS - z_latents match (cosine {zc:.4f})" + if ok else f"⚠️ z_latents cosine {zc:.4f} < 0.99")) + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents_vln.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents_vln.py new file mode 100644 index 0000000..482ee21 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents_vln.py @@ -0,0 +1,215 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Verify System 2 z_latents fidelity on the REAL VLN input distribution. + +Same z_latents bridge comparison as ``verify_system2_latents.py`` (FP8 engine vs +PyTorch reference), but instead of a single caption image it drives the model with +genuine InternData-N1 VLN-CE steps (multi-image history + navigation instruction + +4 appended TRAJ tokens), averaged over several samples. This is the probe that +actually reflects deployment, so it is the fair test of whether VLN-distribution +FP8 calibration helps — a single-image caption probe is out-of-distribution. + +Select the engine with ENGINE_PATH; select the calibration/probe data with +VLN_CALIB_DATA. Uses a fixed seed so every engine sees identical inputs. +""" +import os +import sys + +_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, _R) +sys.path.insert(0, _R) +sys.path.insert(0, os.path.join(_R, "build", "quantize")) + +import torch + +from engine_runner import (REPKG, ENGINE, CKPT, build_mrope_table, run_engine, + cos, per_tok_cos) + +TRAJ_TOKEN_INDEX = 151667 +IMAGE_TOKEN_INDEX = 151655 +N_QUERY = 4 +N_SAMPLES = int(os.environ.get("VLN_PROBE_SAMPLES", "12")) +DATA = os.path.expanduser(os.environ.get( + "VLN_CALIB_DATA", + "~/vln-opt-work/calib_scenes")) # output of build/00_fetch_calib_scenes.sh; override via env + + +def load_ckpt_tensor(prefix): + """Read bridge tensors, preferring the small bridge.safetensors next to the checkpoint. + + repackage_system2.py sets latent_queries and cond_projector aside in a ~25 MB + bridge.safetensors precisely so this check does not have to reopen the 16 GB original. + Falling back to a full sharded checkpoint keeps the script usable against one. + """ + from safetensors import safe_open + + bridge = os.path.join(CKPT, "bridge.safetensors") + if os.path.isfile(bridge): + with safe_open(bridge, framework="pt") as f: + return {k: f.get_tensor(k) for k in f.keys() if k.startswith(prefix)} + + index = os.path.join(CKPT, "model.safetensors.index.json") + if not os.path.isfile(index): + raise FileNotFoundError( + f"neither bridge.safetensors nor model.safetensors.index.json under {CKPT}; " + f"point INTERNVLA_CKPT at a repackaged checkpoint or the original") + idx = json.load(open(index))["weight_map"] + out = {} + for k in idx: + if k.startswith(prefix): + with safe_open(os.path.join(CKPT, idx[k]), framework="pt") as f: + out[k] = f.get_tensor(k) + return out + + +def rope_index(model, input_ids, image_grid_thw, image_token_id): + """Call get_rope_index across the two transformers signatures. + + transformers 4.x took (input_ids, image_grid_thw). 5.x inserted a required + mm_token_type_ids argument -- 0 text, 1 image, 2 video -- and moved the grids to + keywords. Passing the 4.x form to a 5.x model does not raise on arity; it fails later + inside with 'NoneType is not an iterator', which is a confusing way to learn this. + """ + import inspect + + fn = model_attr(model, "get_rope_index") + params = inspect.signature(fn).parameters + if "mm_token_type_ids" in params: + mm = (input_ids == image_token_id).to(torch.int32) + return fn(input_ids, mm, image_grid_thw=image_grid_thw) + return fn(input_ids, image_grid_thw) + + +def model_attr(model, name): + """Fetch an attribute or bound method from the model or its inner model. + + transformers 5.x moved several Qwen2.5-VL helpers (get_rope_index, visual, ...) from + the ForConditionalGeneration wrapper down onto the inner Qwen2_5_VLModel. Probing both + keeps this working on either layout instead of pinning a transformers version. + """ + for owner in (model, getattr(model, "model", None)): + if owner is not None and hasattr(owner, name): + return getattr(owner, name) + raise AttributeError(f"neither the model nor its inner model has {name!r}") + + +def get_visual(model): + """Return the vision tower across transformers layouts. + + Older versions expose it as ``model.visual``; newer ones nest it under + ``model.model.visual``. Probing beats pinning a transformers version here. + """ + for owner in (model, getattr(model, "model", None)): + vis = getattr(owner, "visual", None) if owner is not None else None + if vis is not None: + return vis + raise AttributeError("no vision tower found on this model " + "(looked at .visual and .model.visual)") + + +def visual_embeds(model, pixel_values, grid_thw): + """Return merged image embeddings at LLM hidden width, across transformers layouts. + + In transformers 5.x get_image_features returns a BaseModelOutputWithPooling whose + ``pooler_output`` holds the *merged* embeddings (split per image), while + ``last_hidden_state`` is the pre-merger tensor at vision width -- 1280 here against the + LLM's 3584. Reading the wrong field fails loudly on the shape, but only after the + vision tower has already run, so unwrap explicitly. + """ + if hasattr(model, "get_image_features"): + out = model.get_image_features(pixel_values.type(get_visual(model).dtype), grid_thw) + feats = getattr(out, "pooler_output", out) + else: + vis = get_visual(model) + out = vis(pixel_values.type(vis.dtype), grid_thw=grid_thw) + feats = getattr(out, "last_hidden_state", out) + if isinstance(feats, (list, tuple)): + feats = torch.cat([f.reshape(-1, f.shape[-1]) for f in feats], dim=0) + return feats.reshape(-1, feats.shape[-1]) + + +def main(): + dev = "cuda" + from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration + from calibration import vln_calib_dataloader + + print(f"[1/4] Load repackage + processor | engine={os.path.basename(ENGINE)}") + model = Qwen2_5_VLForConditionalGeneration.from_pretrained( + REPKG, torch_dtype=torch.bfloat16, # flash-attn is not available on Jetson; sdpa is the supported path and + # is what the deployed agent uses too. + attn_implementation=os.environ.get("ATTN_IMPL", "sdpa"), + low_cpu_mem_usage=True).to(dev).eval() + proc = AutoProcessor.from_pretrained(REPKG, trust_remote_code=True, + min_pixels=128 * 28 * 28, max_pixels=1024 * 28 * 28) + inner = model.model + lm = inner.language_model if hasattr(inner, "language_model") else inner + final_norm = lm.norm + + lq = load_ckpt_tensor("model.latent_queries")["model.latent_queries"].to(dev).to(torch.bfloat16) + cpw = load_ckpt_tensor("model.cond_projector") + cp = {k.replace("model.cond_projector.", ""): v.float().to(dev) for k, v in cpw.items()} + assert lq.shape[1] == N_QUERY + + def cond_project(x): + x = torch.nn.functional.linear(x, cp["0.weight"], cp.get("0.bias")) + x = torch.nn.functional.gelu(x) + return torch.nn.functional.linear(x, cp["2.weight"], cp.get("2.bias")) + + print(f"[2/4] Build {N_SAMPLES} real VLN inputs (multi-image + instruction) | data={DATA}") + batches = vln_calib_dataloader(proc, data_root=DATA, num_samples=N_SAMPLES, seed=0) + + print("[3/4] For each: reference LM forward vs engine forward at the 4 TRAJ positions") + hid_pre, hid_post, zc_list, zflat_list = [], [], [], [] + for bi, enc in enumerate(batches): + input_ids = enc["input_ids"].to(dev) + grid = enc["image_grid_thw"].to(dev) + pixel_values = enc["pixel_values"].to(dev) + with torch.no_grad(): + text_embeds = model.get_input_embeddings()(input_ids) + image_embeds = model.visual(pixel_values.type(model.visual.dtype), grid_thw=grid) + image_idx = (input_ids == IMAGE_TOKEN_INDEX) + text_embeds[image_idx] = image_embeds.to(text_embeds.dtype)[: image_idx.sum(), :] + inputs_embeds = torch.cat([text_embeds, lq.repeat(text_embeds.shape[0], 1, 1)], dim=1) + ids_traj = torch.cat( + [input_ids, torch.tensor([[TRAJ_TOKEN_INDEX] * N_QUERY], device=dev)], dim=1) + position_ids, _ = rope_index(model, ids_traj, grid, IMAGE_TOKEN_INDEX) + + cap = {} + h = final_norm.register_forward_hook( + lambda m, i, o: cap.update(pre=i[0].detach(), post=o.detach())) + with torch.no_grad(): + inner(inputs_embeds=inputs_embeds, position_ids=position_ids, + output_hidden_states=True, return_dict=True) + h.remove() + ref_pre = cap["pre"][:, -N_QUERY:, :].float() + ref_post = cap["post"][:, -N_QUERY:, :].float() + + rope = build_mrope_table(position_ids, dev) + _, eng_hs = run_engine(inputs_embeds.to(torch.float16), rope) + eng_pre = eng_hs[:, -N_QUERY:, :].float() + with torch.no_grad(): + eng_post = final_norm(eng_pre.to(torch.bfloat16)).float() + + z_ref, z_eng = cond_project(ref_post), cond_project(eng_post) + hp = per_tok_cos(eng_pre, ref_pre) + hpo = per_tok_cos(eng_post, ref_post) + zc = per_tok_cos(z_eng, z_ref) + zf = cos(z_eng, z_ref) + hid_pre.append(hp); hid_post.append(hpo); zc_list.append(zc); zflat_list.append(zf) + print(f" #{bi:2d} imgs={grid.shape[0]:2d} seq={input_ids.shape[1]:4d} " + f"| hidPRE={hp:.5f} hidPOST={hpo:.5f} z={zc:.5f}") + + import statistics as st + n = len(zc_list) + print("\n[4/4] Mean over VLN inputs\n" + "=" * 60) + print(f" samples : {n}") + print(f" hidden PRE-norm : {st.mean(hid_pre):.6f} (min {min(hid_pre):.6f})") + print(f" hidden POST-norm : {st.mean(hid_post):.6f} (min {min(hid_post):.6f})") + print(f" z_latents perq : {st.mean(zc_list):.6f} (min {min(zc_list):.6f})") + print(f" z_latents flat : {st.mean(zflat_list):.6f}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From fa970752653b61c7548731a1efbe68fed1f82bb8 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Thu, 13 Aug 2026 17:27:28 +0700 Subject: [PATCH 10/30] feat(internvla-n1-dualvln): port System 1 export, verification and benchmarks Adds traj_dit and memory-block export, their BF16 engine builds, the System-1 fidelity check and the two System-1 benchmarks, plus internvla_compat.py which carries the three patches needed to load System 1 at all: the DepthAnythingV2 checkpoint path, the diffusers gradient-checkpointing signature, and the traj_dit FFN multiplier. That last one is a real checkpoint-compatibility finding -- DualVLN was trained with ffn_dim_multiplier 2/3, and the stock build_traj_dit never passes it, so the state dict mismatches without the patch. Four defects carried over from the source project are fixed here rather than propagated: MemBlock was defined inline in the memory exporter as well as in memblock.py, so the two copies could drift -- the exporter now imports the single definition; the traj_dit engine filename said _async in one place and not the other, and neither constant was read; a dead 'if False else None' statement; and INTERNNAV_ROOT is renamed to the recipe-wide INTERNNAV_PATH. trt_torch.py keeps its original NVIDIA Apache-2.0 header rather than being restamped. It is third-party code and Apache-2.0 section 4 requires retaining the notice; it is the one file in this recipe that does not carry the VinRobotics BSD header. Documented an environment constraint found by hitting it: System 1 must run under transformers 4.51.3. The InternNav modeling code reads config.hidden_size off the top-level config, which transformers 5.x no longer flattens, so it fails with a bare AttributeError. The System-2 path is unaffected and runs on either version. (cherry picked from commit 4707c05cc0ce4a1cbd4f72eae327fef20687ee40) --- recipes/internvla-n1-dualvln/README.md | 7 + .../trt-edgellm/benchmark/bench_memory.py | 59 +++++ .../trt-edgellm/benchmark/bench_system1.py | 66 +++++ .../trt-edgellm/export_memory_block.py | 93 +++++++ .../trt-edgellm/export_traj_dit.py | 54 ++++ .../trt-edgellm/internvla_compat.py | 175 +++++++++++++ .../trt-edgellm/memblock.py | 25 ++ .../trt-edgellm/traj_dit_loader.py | 109 +++++++++ .../trt-edgellm/trt_torch.py | 230 ++++++++++++++++++ .../trt-edgellm/verify/verify_system1.py | 114 +++++++++ 10 files changed, 932 insertions(+) create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_memory.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_system1.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/export_traj_dit.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/internvla_compat.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/memblock.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/traj_dit_loader.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/trt_torch.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_system1.py diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index 4076958..ea550f6 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -191,6 +191,13 @@ lock unsatisfiable, so numpy is deliberately absent from this recipe's extra. Ru pip install "numpy==1.26.4" "scipy==1.13.1" +**Step 2c — System 1 needs a different transformers.** The InternNav modeling code reads +`config.hidden_size` off the top-level config, which transformers 5.x no longer flattens, so +System-1 export fails there with `'InternVLAN1ModelConfig' object has no attribute +'hidden_size'`. Run the System-1 steps under **transformers 4.51.3** with `diffusers==0.33.1` +and `onnx==1.22.0` present. The System-2 path (repackage, quantize, export, engine build, +latent verification) is unaffected and runs on either. + **Step 3 — dependencies that are not on PyPI.** These must be present before running anything: | Dependency | How | diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_memory.py b/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_memory.py new file mode 100644 index 0000000..658027e --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_memory.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Measure peak GPU memory (torch.cuda.max_memory_allocated) of the full PyTorch +pipeline (System 2 weights + System 1 forward). +""" +import os, sys +_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, _R); sys.path.insert(0, os.path.join(_R, "lib")) +sys.path.append("/usr/lib/python3.12/dist-packages") +import numpy as np, torch +from PIL import Image + +ACTIVE=os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) +CKPT=os.path.join(ACTIVE,"checkpoints/InternVLA-N1-DualVLN") +IMG=os.path.expanduser("~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg") +GB=1024**3 + + +def peak(): return torch.cuda.max_memory_allocated()/GB +def reset(): torch.cuda.reset_peak_memory_stats(); torch.cuda.empty_cache() + + +def main(): + dev="cuda" + try: torch.backends.mha.set_fastpath_enabled(False) + except Exception: pass + if ACTIVE not in sys.path: sys.path.insert(0, ACTIVE) + from internvla_compat import apply_all + apply_all(need_system1=True, allow_missing_depth=True) + from internnav.model.basemodel.internvla_n1.internvla_n1 import ( + InternVLAN1ForCausalLM, InternVLAN1ModelConfig) + reset() + print("[1] Load full InternVLA (PyTorch, System1+System2)") + cfg=InternVLAN1ModelConfig.from_pretrained(CKPT) + model=InternVLAN1ForCausalLM.from_pretrained(CKPT,config=cfg,torch_dtype=torch.bfloat16, + attn_implementation="sdpa",low_cpu_mem_usage=True).to(dev).eval() + w_mem=torch.cuda.memory_allocated()/GB + print(f" weights loaded (S2+S1): {w_mem:.2f} GB") + + NQ=model.get_n_query() + z=torch.randn(1,NQ,cfg.hidden_size,device=dev,dtype=torch.bfloat16) + a=np.array(Image.open(IMG).convert("RGB").resize((224,224)))/255.0 + tt=torch.from_numpy(a).float(); images_dp=torch.stack([tt,tt]).unsqueeze(0).to(dev) + + print("[2] S1 generate_traj peak (PyTorch)") + reset() + with torch.no_grad(): + model.generate_traj(z,images_dp,num_sample_trajs=32,num_inference_steps=10) + torch.cuda.synchronize() + print(f" S1 peak (incl. weights): {peak():.2f} GB") + + print("\n=== PyTorch peak GPU mem ===") + print(f" weights S2+S1 : {w_mem:.2f} GB") + print(f" peak during S1: {peak():.2f} GB") + return 0 + + +if __name__=="__main__": + sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_system1.py b/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_system1.py new file mode 100644 index 0000000..523e0a5 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_system1.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Benchmark System 1 (generate_traj) latency in PyTorch and report Hz. +""" +import os, sys, time +_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, _R); sys.path.insert(0, os.path.join(_R, "lib")) +sys.path.append("/usr/lib/python3.12/dist-packages") # cv2 (system) cho depth_anything +import numpy as np, torch +from PIL import Image + +ACTIVE = os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) +CKPT = os.path.join(ACTIVE, "checkpoints/InternVLA-N1-DualVLN") +REPKG = os.path.join(os.environ.get("VLN_OPT_WORK",os.path.expanduser("~/vln-opt-work")), "qwen25vl_system2") +IMG = os.path.expanduser("~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg") + + +def timeit(fn, warm=2, n=5): + for _ in range(warm): fn() + torch.cuda.synchronize() + ts = [] + for _ in range(n): + torch.cuda.synchronize(); t0 = time.perf_counter(); fn(); torch.cuda.synchronize() + ts.append(time.perf_counter()-t0) + return sum(ts)/len(ts), min(ts), max(ts) + + +def main(): + dev = "cuda" + if ACTIVE not in sys.path: sys.path.insert(0, ACTIVE) + from internvla_compat import apply_all + apply_all(need_system1=True, allow_missing_depth=True) + from internnav.model.basemodel.internvla_n1.internvla_n1 import ( + InternVLAN1ForCausalLM, InternVLAN1ModelConfig) + print("[1/3] Load full InternVLA (System1)") + cfg = InternVLAN1ModelConfig.from_pretrained(CKPT) + model = InternVLAN1ForCausalLM.from_pretrained( + CKPT, config=cfg, torch_dtype=torch.bfloat16, attn_implementation=os.environ.get("ATTN","sdpa"), + low_cpu_mem_usage=True).to(dev).eval() + + print("[2/3] Prepare z_latents + images_dp (as the agent)") + N_QUERY = model.get_n_query() + # placeholder z_latents (cond_projector is applied inside generate_traj) + traj_latents = torch.randn(1, N_QUERY, cfg.hidden_size, device=dev, dtype=torch.bfloat16) + a = np.array(Image.open(IMG).convert("RGB").resize((224,224)))/255.0 + t = torch.from_numpy(a).float() + images_dp = torch.stack([t, t]).unsqueeze(0).to(dev) # [1,2,224,224,3] + + print("[3/3] Đo generate_traj (S1)\n" + "="*60) + with torch.no_grad(): + for ns in (32, 4, 1): + def fn(ns=ns): + with torch.no_grad(): + model.generate_traj(traj_latents.to(dev), images_dp, + num_sample_trajs=ns, num_inference_steps=10) + try: + mean,mn,mx = timeit(fn, warm=2, n=5) + print(f" num_sample_trajs={ns:>2}: {mean*1000:7.1f} ms (min {mn*1000:.0f}, max {mx*1000:.0f}) → {1/mean:5.1f} Hz") + except Exception as e: + print(f" num_sample_trajs={ns}: ERR {type(e).__name__}: {e}") + print("\n (paper target S1 = 30 Hz ≈ 33 ms; S2 = 2 Hz ≈ 500 ms)") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py b/recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py new file mode 100644 index 0000000..0c7e6e5 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Step 5 - export the System 1 memory block (DINOv2 rgb_model + memory_encoder + +rgb_resampler) to ONNX and build the BF16 engine. +""" +import os, sys, time, subprocess +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +sys.path.append("/usr/lib/python3.12/dist-packages") +sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "lib")) +import numpy as np, torch +from PIL import Image + +ACTIVE=os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) +CKPT=os.path.join(ACTIVE,"checkpoints/InternVLA-N1-DualVLN") +IMG=os.path.expanduser("~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg") +ONNX=os.path.join(os.environ.get("VLN_OPT_WORK",os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.onnx") +ENG=os.path.join(os.environ.get("VLN_OPT_WORK",os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.engine") + + +# MemBlock lives in memblock.py; defining it twice is how the two copies drift. +from memblock import MemBlock + + +def main(): + dev="cuda" + # disable fused MHA fastpath (_transformer_encoder_layer_fwd is not ONNX-exportable) + try: torch.backends.mha.set_fastpath_enabled(False) + except Exception as e: print(" (mha fastpath toggle:", e, ")") + if ACTIVE not in sys.path: sys.path.insert(0, ACTIVE) + from internvla_compat import apply_all + apply_all(need_system1=True, allow_missing_depth=True) + from internnav.model.basemodel.internvla_n1.internvla_n1 import ( + InternVLAN1ForCausalLM, InternVLAN1ModelConfig) + print("[1/6] Load full model") + cfg=InternVLAN1ModelConfig.from_pretrained(CKPT) + model=InternVLAN1ForCausalLM.from_pretrained(CKPT,config=cfg,torch_dtype=torch.bfloat16, + attn_implementation="sdpa",low_cpu_mem_usage=True).to(dev).eval() + m=model.get_model() + + print("[2/6] Build input (matching the reference normalization)") + a=np.array(Image.open(IMG).convert("RGB").resize((224,224)))/255.0 + tt=torch.from_numpy(a).float(); images_dp=torch.stack([tt,tt]).unsqueeze(0).to(dev) # [1,2,224,224,3] + dtype=torch.bfloat16 + rmean = model._resnet_mean if hasattr(model,"_resnet_mean") else m._resnet_mean + rstd = model._resnet_std if hasattr(model,"_resnet_std") else m._resnet_std + x = images_dp.permute(0,1,4,2,3) # [1,2,3,224,224] + x = (x - rmean)/rstd + x = x.flatten(0,1).to(dtype) # [2,3,224,224] + print(f" input {tuple(x.shape)}") + + print("[3/6] Ref base (PyTorch BF16) memory_tokens + latency") + block = MemBlock(m.rgb_model, m.memory_encoder, m.rgb_resampler).eval() + with torch.no_grad(): ref = block(x).float() + print(f" memory_tokens {tuple(ref.shape)}") + def lat(fn,warm=3,n=10): + for _ in range(warm): fn() + torch.cuda.synchronize(); ts=[] + for _ in range(n): + torch.cuda.synchronize(); t0=time.perf_counter(); fn(); torch.cuda.synchronize(); ts.append(time.perf_counter()-t0) + return sum(ts)/len(ts)*1000 + with torch.no_grad(): pt_ms=lat(lambda: block(x)) + + print("[4/6] Export ONNX (FP32 to avoid mixed precision; build BF16 later)") + block_fp32 = MemBlock(m.rgb_model.float(), m.memory_encoder.float(), m.rgb_resampler.float()).eval() + xf = x.float() + os.makedirs(os.path.dirname(ONNX),exist_ok=True) + with torch.inference_mode(): + torch.onnx.export(block_fp32,(xf,),ONNX,input_names=["images"],output_names=["memory_tokens"], + opset_version=19,do_constant_folding=True,export_params=True,dynamo=False) + print(f" {os.path.getsize(ONNX)/1e6:.1f} MB") + + print("[5/6] Build BF16 engine (fixed shape 2x3x224x224)") + cmd=["/usr/src/tensorrt/bin/trtexec",f"--onnx={ONNX}",f"--saveEngine={ENG}","--bf16"] # static shape + r=subprocess.run(cmd,capture_output=True,text=True) + print(f" build {'OK' if os.path.exists(ENG) else 'FAIL'} | {os.path.getsize(ENG)/1e6 if os.path.exists(ENG) else 0:.1f} MB") + if not os.path.exists(ENG): + print(" STDERR:", r.stderr[-600:]); return 1 + + print("[6/6] Parity + latency (TRT vs base)") + from trt_torch import Engine + eng=Engine(ENG) + out=eng(images=x.float().contiguous()) + out=(out.get("memory_tokens") if isinstance(out,dict) else out).float() + cos=torch.nn.functional.cosine_similarity(ref.flatten(),out.flatten(),dim=0).item() + trt_ms=lat(lambda: eng(images=x.float().contiguous())) + print(f" parity cos={cos:.5f} rel-L2={(ref-out).norm()/ref.norm():.4f}") + print(f" latency PyTorch {pt_ms:.2f}ms → TRT {trt_ms:.2f}ms = {pt_ms/trt_ms:.2f}x") + return 0 + + +if __name__=="__main__": + sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/export_traj_dit.py b/recipes/internvla-n1-dualvln/trt-edgellm/export_traj_dit.py new file mode 100644 index 0000000..93a70fe --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/export_traj_dit.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Step 4 - export the System 1 traj_dit (NextDiT) core to ONNX with a dynamic +z_latents length and build the BF16 engine (trtexec). +""" +import os, sys, subprocess +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +sys.path.append("/usr/lib/python3.12/dist-packages") +sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "lib")) +import torch +from traj_dit_loader import load_traj_dit + +OUT = os.path.join(os.environ.get("VLN_OPT_WORK",os.path.expanduser("~/vln-opt-work")), "onnx/system1_traj_dit_async.onnx") +ENG = os.path.join(os.environ.get("VLN_OPT_WORK",os.path.expanduser("~/vln-opt-work")), "onnx/system1_traj_dit_bf16.engine") +ZLEN = int(os.environ.get("ZLEN", "36")) # observed z_latents length for the async path + + +def main(): + dev="cuda"; N,WP,DIM=32,32,384 + print(f"[1/3] Load traj_dit (FP32) | z_latents seq (async, dynamic, opt={ZLEN})") + dit,zdim = load_traj_dit(); dit=dit.to(torch.float32).eval() + x=torch.randn(2*N,WP,DIM,dtype=torch.float32,device=dev) + ts=torch.ones(2*N,dtype=torch.int64,device=dev) + z=torch.randn(2*N,ZLEN,zdim,dtype=torch.float32,device=dev) + class W(torch.nn.Module): + def __init__(s,d): super().__init__(); s.d=d + def forward(s,x,timestep,z_latents): return s.d(x=x,timestep=timestep,z_latents=z_latents) + w=W(dit).eval() + with torch.no_grad(): ref=w(x,ts,z) + print(f" forward OK -> {tuple(ref.shape)}") + os.makedirs(os.path.dirname(OUT),exist_ok=True) + # dynamic: batch (dim0) + z_latents seq (dim1) + dyn={"x":{0:"batch"},"timestep":{0:"batch"},"z_latents":{0:"batch",1:"zlen"},"output":{0:"batch"}} + print(f"[2/3] Export ONNX (dynamo=False, opset 19) → {OUT}") + with torch.inference_mode(): + torch.onnx.export(w,(x,ts,z),OUT,input_names=["x","timestep","z_latents"], + output_names=["output"],opset_version=19,do_constant_folding=True, + export_params=True,dynamic_axes=dyn,dynamo=False) + print(f" {os.path.getsize(OUT)/1e6:.1f} MB") + B=2*N + cmd=["/usr/src/tensorrt/bin/trtexec",f"--onnx={OUT}",f"--saveEngine={ENG}","--bf16", + f"--minShapes=x:{B}x{WP}x{DIM},timestep:{B},z_latents:{B}x4x{zdim}", + f"--optShapes=x:{B}x{WP}x{DIM},timestep:{B},z_latents:{B}x{ZLEN}x{zdim}", + f"--maxShapes=x:{B}x{WP}x{DIM},timestep:{B},z_latents:{B}x64x{zdim}"] + print(f"[3/3] Build engine BF16 (z dynamic 4..64, opt {ZLEN})\n {' '.join(cmd)}") + r=subprocess.run(cmd,capture_output=True,text=True) + print(" "+ "\n ".join(l for l in r.stdout.splitlines() if "successfully" in l.lower() or "PASSED" in l or "FAILED" in l)[-400:]) + print(" engine:", ENG, os.path.getsize(ENG)/1e6 if os.path.exists(ENG) else "MISSING", "MB") + return 0 if os.path.exists(ENG) else 1 + + +if __name__=="__main__": + sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/internvla_compat.py b/recipes/internvla-n1-dualvln/trt-edgellm/internvla_compat.py new file mode 100644 index 0000000..24b10b4 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/internvla_compat.py @@ -0,0 +1,175 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Runtime compatibility patches for loading InternVLA-N1 System 1. + +These patches are applied in-memory (attribute reassignment) and never modify the InternNav +source tree. They fix three issues that otherwise prevent building the System 1 modules: + +1. ``build_depthanythingv2`` loads the DepthAnything-V2 checkpoint from a relative path; the + patch loads it from an absolute path (set via DAV2_CKPT / env), and can tolerate a missing + file since ``from_pretrained`` overwrites ``model.rgb_model.*`` from safetensors afterwards. + +2. ``LuminaNextDiT2DModel._set_gradient_checkpointing`` has an old signature that current + diffusers calls with ``enable=`` / ``gradient_checkpointing_func=`` keywords; the patch + accepts both. + +3. ``build_traj_dit`` must pass ``ffn_dim_multiplier = 2/3`` (SwiGLU convention) so the traj_dit + FFN shape matches the trained DualVLN checkpoint (1024, not 1536). +""" +from __future__ import annotations + +import os + +#: Root of the InternNav checkout to load the model from. +ACTIVE_INTERNNAV = os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) + +#: DepthAnything-V2 checkpoint used by System 1's rgb_model. +DAV2_CKPT = os.path.expanduser(os.environ.get( + "DAV2_CKPT", os.path.join(ACTIVE_INTERNNAV, "checkpoints/depth_anything_v2_vits.pth"))) + +#: FFN multiplier the DualVLN checkpoint was trained with (see patch_traj_dit_ffn). +TRAJ_DIT_FFN_MULTIPLIER = 2.0 / 3.0 + + +def assert_active_tree() -> str: + """Confirm ``internnav`` is imported from ACTIVE_INTERNNAV; fail early otherwise. + + Which tree loads depends on ``cwd`` (``''`` heads ``sys.path``), so a stray ``cd`` can + silently swap the model code. This turns that silent failure into a loud one. + """ + import internnav + + tree = os.path.dirname(os.path.dirname(internnav.__file__)) + if os.path.realpath(tree) != os.path.realpath(ACTIVE_INTERNNAV): + raise RuntimeError( + f"internnav loaded from the wrong tree:\n" + f" actual : {tree}\n" + f" expected: {ACTIVE_INTERNNAV}\n" + f"Usually the cwd sits inside a different InternNav checkout. Re-run elsewhere, " + f"or put {ACTIVE_INTERNNAV} at the front of sys.path." + ) + return tree + + +def patch_depth_anything(allow_missing: bool = False) -> bool: + """Replace ``build_depthanythingv2`` with a version that loads from an absolute path. + + Only reassigns a module attribute in memory. Returns True once patched. + """ + import torch + + import internnav.model.basemodel.internvla_n1.internvla_n1_arch as arch + + if getattr(arch, "_vlnopt_patched", False): + return True + + if not os.path.isfile(DAV2_CKPT): + if not allow_missing: + raise FileNotFoundError( + f"DepthAnything-V2 checkpoint not found: {DAV2_CKPT}\n" + f"Required for System 1 (nextdit_async). Pass allow_missing=True to build the " + f"architecture with random init — acceptable because from_pretrained overwrites " + f"model.rgb_model.* from safetensors immediately afterwards." + ) + print(f"[compat] WARNING: missing {DAV2_CKPT} — rgb_model randomly initialized, " + f"relying on from_pretrained to load weights.") + + def _build_dav2_patched(config): + from internnav.model.encoder.depth_anything.depth_anything_v2.dpt import ( + DepthAnythingV2, + ) + + model_configs = { + "vits": {"encoder": "vits", "features": 64, + "out_channels": [48, 96, 192, 384]} + } + dav2 = DepthAnythingV2(**model_configs["vits"]) + if os.path.isfile(DAV2_CKPT): + dav2.load_state_dict(torch.load(DAV2_CKPT, map_location="cpu")) + return dav2.pretrained + + arch.build_depthanythingv2 = _build_dav2_patched + arch._vlnopt_patched = True + return True + + +def patch_gradient_checkpointing() -> bool: + """Reconcile the ``_set_gradient_checkpointing`` signature between nextdit and diffusers. + + nextdit defines ``(self, module, value)`` but current diffusers calls it with + ``(enable=..., gradient_checkpointing_func=...)``. This fails while building traj_dit, so + it must be patched before any System 1 build path. + """ + try: + from internnav.model.basemodel.internvla_n1.nextdit_traj import ( + LuminaNextDiT2DModel, + ) + except ImportError: + return False + + if getattr(LuminaNextDiT2DModel, "_vlnopt_gc_patched", False): + return True + + def _set_gc_compat(self, module=None, value=False, enable=None, + gradient_checkpointing_func=None): + v = enable if enable is not None else value + if module is not None: + module.gradient_checkpointing = v + else: + for m in self.modules(): + if hasattr(m, "gradient_checkpointing"): + m.gradient_checkpointing = v + + LuminaNextDiT2DModel._set_gradient_checkpointing = _set_gc_compat + LuminaNextDiT2DModel._vlnopt_gc_patched = True + return True + + +def patch_traj_dit_ffn(multiplier: float = TRAJ_DIT_FFN_MULTIPLIER) -> bool: + """Patch ``build_traj_dit`` so the traj_dit FFN matches the checkpoint. + + The DualVLN checkpoint was trained with ``ffn_dim_multiplier = 2/3`` (SwiGLU convention), + giving inner_dim 1024 for dim=384. The stock ``build_traj_dit`` never passes this, leaving + inner_dim at 1536 and causing a state_dict size mismatch. Only feed_forward.linear_1/linear_3 + are affected; norm1.linear is 4*dim on both sides and already matches. + """ + import internnav.model.basemodel.internvla_n1.internvla_n1_arch as arch + + if getattr(arch, "_vlnopt_ffn_patched", False): + return True + + _orig = arch.build_traj_dit + + def _build_traj_dit_patched(config): + from diffusers.schedulers import FlowMatchEulerDiscreteScheduler + + from internnav.model.basemodel.internvla_n1.nextdit_crossattn_traj import ( + NextDiTCrossAttn, NextDiTCrossAttnConfig, + ) + + dit_cfg = NextDiTCrossAttnConfig( + latent_embedding_size=arch.LatentEmbSize, + ffn_dim_multiplier=multiplier, # the only difference from the stock builder + ) + dit = NextDiTCrossAttn(dit_cfg) + return dit, FlowMatchEulerDiscreteScheduler() + + _build_traj_dit_patched.__wrapped__ = _orig + arch.build_traj_dit = _build_traj_dit_patched + arch._vlnopt_ffn_patched = True + return True + + +def apply_all(need_system1: bool = True, allow_missing_depth: bool = False) -> None: + """Convenience entry point: verify the tree and apply the needed patches.""" + tree = assert_active_tree() + print(f"[compat] internnav tree: {tree}") + if need_system1: + # Order matters: patch gradient-checkpointing first, since it fails inside the + # traj_dit builder. + if patch_gradient_checkpointing(): + print("[compat] patched LuminaNextDiT2DModel._set_gradient_checkpointing") + patch_depth_anything(allow_missing=allow_missing_depth) + print(f"[compat] patched build_depthanythingv2 -> {DAV2_CKPT}") + patch_traj_dit_ffn() + print(f"[compat] patched build_traj_dit -> ffn_dim_multiplier={TRAJ_DIT_FFN_MULTIPLIER:.4f}") diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/memblock.py b/recipes/internvla-n1-dualvln/trt-edgellm/memblock.py new file mode 100644 index 0000000..021f52b --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/memblock.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""System 1 memory-encoding block, wrapped as a single exportable module. + +Mirrors the reference `generate_traj` (nextdit_async) memory path: + rgb_model.get_intermediate_layers -> memory_encoder -> concat -> rgb_resampler +Input : images_dp_norm [T, 3, 224, 224] (T frames, ImageNet-normalized) +Output: memory_tokens [1, 32, 768] +""" +import torch + + +class MemBlock(torch.nn.Module): + def __init__(self, rgb_model, memory_encoder, rgb_resampler): + super().__init__() + self.rgb = rgb_model + self.me = memory_encoder + self.rr = rgb_resampler + + def forward(self, x): # x = [T, 3, 224, 224] + feat = self.rgb.get_intermediate_layers(x)[0].unflatten(dim=0, sizes=(1, -1)) # [1, T, Np, 384] + f = feat.flatten(1, 2) # [1, T*Np, 384] + mf = self.me(f) + mf = torch.cat([f, mf], dim=-1) + return self.rr(mf) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/traj_dit_loader.py b/recipes/internvla-n1-dualvln/trt-edgellm/traj_dit_loader.py new file mode 100644 index 0000000..a7d1ae3 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/traj_dit_loader.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Load the InternVLA-N1 System 1 traj_dit (NextDiT) core, and export it to ONNX. + +The traj_dit is exported in FP32 (its RoPE freqs_cis are float32; FP32 export avoids +mixed-precision trace errors). Precision is set to BF16 at engine-build time. The surrounding +host-side ops (action encoder/decoder, cond_projector, pos_encoding, CFG blend, scheduler) +stay in PyTorch. Legacy exporter (dynamo=False, opset 19). + +Export shapes: + input : x [2N, 32, 384] · timestep [2N] int64 · z_latents [2N, Z, 768] + output: noise_pred [2N, 32, 384] +""" +import os +import sys + +import torch + +_LIB = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, _LIB) +sys.path.insert(0, os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav"))) + +MODEL = os.environ.get("INTERNVLA_CKPT", os.path.join( + os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")), + "checkpoints/InternVLA-N1-DualVLN")) +OUT = os.path.join(os.environ.get("VLN_OPT_WORK", os.path.expanduser("~/vln-opt-work")), + "onnx/system1_traj_dit.onnx") + + +def load_traj_dit(): + """Return (traj_dit module, LatentEmbSize). The rgb_model (DepthAnything) is stubbed + out since only the traj_dit is needed here.""" + from internvla_compat import patch_gradient_checkpointing, patch_traj_dit_ffn + + patch_gradient_checkpointing() + patch_traj_dit_ffn() # requires diffusers 0.33.1 -> inner_dim 1024 + import internnav.model.basemodel.internvla_n1.internvla_n1_arch as _arch + + class _DepthStub(torch.nn.Module): + def forward(self, *a, **k): + raise RuntimeError("rgb_model stub") + + _arch.build_depthanythingv2 = lambda config: _DepthStub() + + from internnav.model.basemodel.internvla_n1.internvla_n1 import ( + InternVLAN1ForCausalLM, InternVLAN1ModelConfig, + ) + + config = InternVLAN1ModelConfig.from_pretrained(MODEL) + model = InternVLAN1ForCausalLM.from_pretrained( + MODEL, config=config, torch_dtype=torch.float32, low_cpu_mem_usage=True) + model.to("cuda").eval() + return model.get_model().traj_dit, _arch.LatentEmbSize + + +class TrajDiTWrapper(torch.nn.Module): + """Adapt positional args to keyword call (torch.onnx.export passes positionally).""" + + def __init__(self, dit): + super().__init__() + self.dit = dit + + def forward(self, x, timestep, z_latents): + return self.dit(x=x, timestep=timestep, z_latents=z_latents) + + +def main(): + N, WP, DIM = 32, 32, 384 # num_sample_trajs, waypoints, dim + dev = "cuda" + print("[1/3] Load traj_dit (FP32)") + dit, zdim = load_traj_dit() + dit = dit.to(torch.float32).eval() + print(f" z_latents dim (LatentEmbSize) = {zdim}") + + x = torch.randn(2 * N, WP, DIM, dtype=torch.float32, device=dev) + timestep = torch.ones(2 * N, dtype=torch.int64, device=dev) + z_latents = torch.randn(2 * N, 4, zdim, dtype=torch.float32, device=dev) + wrapped = TrajDiTWrapper(dit).eval() + + with torch.no_grad(): + ref = wrapped(x, timestep, z_latents) + print(f"[2/3] Forward OK -> output {tuple(ref.shape)} (expected [{2*N},{WP},{DIM}])") + + os.makedirs(os.path.dirname(OUT), exist_ok=True) + # batch (2N) dynamic for flexible num_sample_trajs; seq (32) & n_query (4) static + dyn = {"x": {0: "batch"}, "timestep": {0: "batch"}, + "z_latents": {0: "batch"}, "output": {0: "batch"}} + print(f"[3/3] Export ONNX (dynamo=False, opset 19) -> {OUT}") + with torch.inference_mode(): + torch.onnx.export( + wrapped, (x, timestep, z_latents), OUT, + input_names=["x", "timestep", "z_latents"], + output_names=["output"], + opset_version=19, do_constant_folding=True, export_params=True, + dynamic_axes=dyn, dynamo=False) + print(f" Done. {OUT} ({os.path.getsize(OUT) / 1e6:.1f} MB)") + + try: + import onnx + onnx.checker.check_model(onnx.load(OUT)) + print(" onnx.checker: valid") + except Exception as e: + print(f" onnx.checker error: {e}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/trt_torch.py b/recipes/internvla-n1-dualvln/trt-edgellm/trt_torch.py new file mode 100644 index 0000000..fef4848 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/trt_torch.py @@ -0,0 +1,230 @@ +# 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. + +# +# 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. + +"""TensorRT Engine wrapper for GR00T inference. + +Loads serialized TRT engines, manages input/output tensor bindings, and +executes inference. Supports dynamic shapes and BF16/FP16/FP32 dtypes. +""" + +import ctypes +import os + +import tensorrt as trt +import torch + + +def torch_type(trt_type): + """Convert TensorRT data type to PyTorch equivalent.""" + mapping = { + trt.float32: torch.float32, + trt.float16: torch.float16, + trt.bfloat16: torch.bfloat16, + trt.int8: torch.int8, + trt.int32: torch.int32, + trt.bool: torch.bool, + trt.uint8: torch.uint8, + trt.int64: torch.int64, + } + if trt_type in mapping: + return mapping[trt_type] + + raise TypeError( + f"Could not resolve TensorRT datatype to an equivalent PyTorch datatype. {trt_type}" + ) + + +class Engine(object): + """TensorRT engine wrapper for loading and executing inference.""" + + def __init__(self, file, plugins=[]): + super().__init__() + + self._closed = False + self.execution_context = None + self.handle = None + + self.logger = trt.Logger(trt.Logger.ERROR) + trt.init_libnvinfer_plugins(self.logger, "") + + self.plugins = [ctypes.CDLL(plugin, ctypes.RTLD_GLOBAL) for plugin in plugins] + self.file = file + self.load(file) + + self.print() + + def close(self): + """Release the execution context, then the engine handle. + + TensorRT requires the context be dropped before the engine. + Idempotent so it composes safely with ``__del__``. + """ + if self._closed: + return + self._closed = True + self.execution_context = None + self.handle = None + + def __del__(self): + # tensorrt / CUDA context may already be gone at shutdown; swallow + # so the traceback doesn't surface as a spurious error on exit. + try: + self.close() + except Exception: + pass + + def print(self): + """Display engine details (inputs/outputs) on rank 0 only.""" + if int(os.getenv("LOCAL_RANK", -1)) not in [0, -1]: + return + + print("============= TRT Engine Detail =============") + print(f"Engine file: {self.file}") + print(f"Inputs: {len(self.in_meta)}") + for ib, item in enumerate(self.in_meta): + tensor_name, shape, dtype = item[:3] + print(f" {ib}. {tensor_name}: {'x'.join(map(str, shape))} [{dtype}]") + + print(f"Outputs: {len(self.out_meta)}") + for ib, item in enumerate(self.out_meta): + tensor_name, shape, dtype = item[:3] + print(f" {ib}. {tensor_name}: {'x'.join(map(str, shape))} [{dtype}]") + print("=============================================") + + def load(self, file): + """Deserialize and load a TensorRT engine from file.""" + runtime = trt.Runtime(self.logger) + + try: + with open(file, "rb") as f: + self.handle = runtime.deserialize_cuda_engine(f.read()) + assert self.handle is not None, ( + f"Failed to deserialize the cuda engine from file: {file}" + ) + + self.execution_context = self.handle.create_execution_context() + self.meta, self.in_meta, self.out_meta = [], [], [] + for tensor_name in self.handle: + shape = self.handle.get_tensor_shape(tensor_name) + dtype = torch_type(self.handle.get_tensor_dtype(tensor_name)) + if self.handle.get_tensor_mode(tensor_name) == trt.TensorIOMode.INPUT: + self.in_meta.append([tensor_name, shape, dtype]) + else: + self.out_meta.append([tensor_name, shape, dtype]) + except BaseException: + # Roll back a half-loaded engine (e.g. create_execution_context + # failed after the handle was set) so it can't leak GPU memory. + self.close() + raise + + def __call__(self, *args, **inputs): + return self.forward(*args, **inputs) + + def dtype_of(self, tensor_name: str) -> torch.dtype: + """Return the expected PyTorch dtype for a named input tensor.""" + for name, _shape, dtype in self.in_meta: + if name == tensor_name: + return dtype + raise KeyError(f"Input tensor '{tensor_name}' not found in engine.") + + def set_runtime_tensor_shape(self, name, shape): + """Set runtime input shape for dynamic dimensions.""" + self.execution_context.set_input_shape(name, shape) + + def forward(self, *args, **kwargs): + """Execute TRT inference with the given input tensors. + + Accepts both positional and keyword arguments. + Returns a dict of output tensors by default, or a list if return_list=True. + """ + return_list = kwargs.pop("return_list", False) + reference_tensors = [] + stream = torch.cuda.current_stream() + + # Process positional arguments + for iarg, x in enumerate(args): + name, shape, dtype = self.in_meta[iarg] + runtime_shape = self.execution_context.get_tensor_shape(name) + assert isinstance(x, torch.Tensor), f"Unsupported tensor type: {type(x)}" + assert runtime_shape == x.shape, f"Invalid input shape: {runtime_shape} != {x.shape}" + assert dtype == x.dtype, ( + f"Invalid tensor dtype, expected dtype is {dtype}, but got {x.dtype}" + ) + assert x.is_cuda, f"Invalid tensor device, expected device is cuda, but got {x.device}" + x = x.cuda().contiguous() + self.execution_context.set_tensor_address(name, x.data_ptr()) + reference_tensors.append(x) + + # Process keyword arguments + for name, shape, dtype in self.in_meta: + if name not in kwargs: + continue + + runtime_shape = self.execution_context.get_tensor_shape(name) + x = kwargs[name] + assert isinstance(x, torch.Tensor), f"Unsupported tensor[{name}] type: {type(x)}" + assert runtime_shape == x.shape, ( + f"Invalid input[{name}] shape: {x.shape}, but the expected shape is: {runtime_shape}" + ) + assert dtype == x.dtype, ( + f"Invalid tensor[{name}] dtype, expected dtype is {dtype}, but got {x.dtype}" + ) + assert x.is_cuda, ( + f"Invalid tensor[{name}] device, expected device is cuda, but got {x.device}" + ) + x = x.cuda().contiguous() + self.execution_context.set_tensor_address(name, x.data_ptr()) + reference_tensors.append(x) + + # Allocate output tensors + for item in self.out_meta: + name = item[0] + runtime_shape = self.execution_context.get_tensor_shape(name) + output_tensor = torch.zeros( + *runtime_shape, dtype=item[2], device=reference_tensors[0].device + ) + self.execution_context.set_tensor_address(name, output_tensor.data_ptr()) + reference_tensors.append(output_tensor) + + # Execute + self.execution_context.execute_async_v3(stream.cuda_stream) + stream.synchronize() + assert len(reference_tensors) == len(self.in_meta) + len(self.out_meta), ( + f"Invalid input tensors. The expected I/O tensors are " + f"{len(self.in_meta) + len(self.out_meta)}, but got {len(reference_tensors)}" + ) + + if return_list: + return [ + reference_tensors[len(self.in_meta) + i] for i, item in enumerate(self.out_meta) + ] + else: + return { + item[0]: reference_tensors[len(self.in_meta) + i] + for i, item in enumerate(self.out_meta) + } diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_system1.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_system1.py new file mode 100644 index 0000000..0bc5899 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_system1.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Verify System 1: wire the traj_dit + memory TensorRT engines into generate_traj and +compare the resulting trajectory and latency against the PyTorch base (same noise seed). +""" +import os, sys, time +_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, _R); sys.path.insert(0, os.path.join(_R, "lib")) +sys.path.append("/usr/lib/python3.12/dist-packages") +sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),"lib")) +import numpy as np, torch +from PIL import Image +from memblock import MemBlock + +ACTIVE=os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) +CKPT=os.path.join(ACTIVE,"checkpoints/InternVLA-N1-DualVLN") +IMG=os.path.expanduser("~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg") +TRAJDIT=os.path.join(os.environ.get("VLN_OPT_WORK",os.path.expanduser("~/vln-opt-work")), "onnx/system1_traj_dit_bf16.engine") +MEM=os.path.join(os.environ.get("VLN_OPT_WORK",os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.engine") +SEED=12345 + + +def out_of(o, key): + if isinstance(o, dict): return o.get(key, next(iter(o.values()))) + return o[0] if isinstance(o,(list,tuple)) else o + + +def main(): + dev="cuda" + try: torch.backends.mha.set_fastpath_enabled(False) + except Exception: pass + if ACTIVE not in sys.path: sys.path.insert(0, ACTIVE) + from internvla_compat import apply_all + apply_all(need_system1=True, allow_missing_depth=True) + from internnav.model.basemodel.internvla_n1.internvla_n1 import ( + InternVLAN1ForCausalLM, InternVLAN1ModelConfig) + from diffusers.schedulers import FlowMatchEulerDiscreteScheduler + from diffusers.utils.torch_utils import randn_tensor + from trt_torch import Engine + print("[1/4] Load model + engines") + cfg=InternVLAN1ModelConfig.from_pretrained(CKPT) + model=InternVLAN1ForCausalLM.from_pretrained(CKPT,config=cfg,torch_dtype=torch.bfloat16, + attn_implementation="sdpa",low_cpu_mem_usage=True).to(dev).eval() + m=model.get_model() + mem_eng=Engine(MEM); dit_eng=Engine(TRAJDIT) + NQ=model.get_n_query() + z=torch.randn(1,NQ,cfg.hidden_size,device=dev,dtype=torch.bfloat16) + a=np.array(Image.open(IMG).convert("RGB").resize((224,224)))/255.0 + tt=torch.from_numpy(a).float(); images_dp=torch.stack([tt,tt]).unsqueeze(0).to(dev) + rmean=model._resnet_mean; rstd=model._resnet_std + + def base_gen(): + torch.manual_seed(SEED); np.random.seed(SEED) + with torch.no_grad(): + return model.generate_traj(z,images_dp,num_sample_trajs=32,num_inference_steps=10).float().cpu() + + def trt_gen(steps=10, ns=32, guidance_scale=1.0, predict_step_nums=32): + """Replicate generate_traj (async) with the memory + traj_dit engines.""" + torch.manual_seed(SEED); np.random.seed(SEED) + dtype=z.dtype + with torch.no_grad(): + traj_latents = m.cond_projector(z) # [1,4,768] + # --- memory block qua ENGINE --- + xdp = images_dp.permute(0,1,4,2,3) + xdp = ((xdp - rmean)/rstd).flatten(0,1) # [2,3,224,224] + memory_tokens = out_of(mem_eng(images=xdp.float().contiguous()), "memory_tokens").to(dtype) # [1,32,768] + hidden_states = torch.cat([memory_tokens, traj_latents], dim=1) # [1,36,768] + hs_null = torch.zeros_like(hidden_states) + hs_input = torch.cat([hs_null, hidden_states], 0) # [2,36,768] + bs = traj_latents.shape[0] + latents = randn_tensor((bs*ns, predict_step_nums, 3), generator=None, device=dev, dtype=dtype) + sch = FlowMatchEulerDiscreteScheduler() + sigmas = np.linspace(1.0, 1/steps, steps) + sch.set_timesteps(steps, sigmas=sigmas) + hs_input = hs_input.repeat_interleave(ns, dim=0) # [2*ns,36,768] + dit_eng.set_runtime_tensor_shape("z_latents", tuple(hs_input.shape)) + for t in sch.timesteps: + lf = m.action_encoder(latents) + pos_ids = torch.arange(lf.shape[1]).reshape(1,-1).repeat(bs*ns,1).to(dev) + lf = lf + m.pos_encoding(pos_ids) + lmi = lf.repeat(2,1,1) + if hasattr(sch,"scale_model_input"): lmi = sch.scale_model_input(lmi, t) + tt_ = t.unsqueeze(0).expand(lmi.shape[0]).to(dev, torch.long) + np_ = out_of(dit_eng(x=lmi.float().contiguous(), timestep=tt_.to(torch.int64).contiguous(), + z_latents=hs_input.float().contiguous()), "output").to(dtype) + np_ = m.action_decoder(np_) + unc, cnd = np_.chunk(2) + np_ = unc + guidance_scale*(cnd - unc) + latents = sch.step(np_, t, latents).prev_sample + return latents.float().cpu() + + def lat(fn,warm=2,n=5): + for _ in range(warm): fn() + torch.cuda.synchronize(); ts=[] + for _ in range(n): + torch.cuda.synchronize(); t0=time.perf_counter(); fn(); torch.cuda.synchronize(); ts.append(time.perf_counter()-t0) + return sum(ts)/len(ts)*1000 + + print("[2/4] Base PyTorch generate_traj") + tref = base_gen(); pt_ms = lat(base_gen) + print(f" {tuple(tref.shape)} | {pt_ms:.1f}ms = {1000/pt_ms:.1f}Hz") + print("[3/4] Full-TRT S1 (memory+traj_dit engines)") + ttrt = trt_gen(); trt_ms = lat(trt_gen) + print(f" {tuple(ttrt.shape)} | {trt_ms:.1f}ms = {1000/trt_ms:.1f}Hz") + print("[4/4] Compare\n"+"="*56) + d=(tref-ttrt).norm(dim=-1); endp=(tref[:,-1]-ttrt[:,-1]).norm(dim=-1) + cos=torch.nn.functional.cosine_similarity(tref.flatten(),ttrt.flatten(),dim=0).item() + print(f" Parity: per-wp L2 mean={d.mean():.4f} max={d.max():.4f} | endpoint={endp.mean():.4f} | cos={cos:.5f}") + print(f" Latency S1: PyTorch {pt_ms:.1f}ms ({1000/pt_ms:.1f}Hz) → full-TRT {trt_ms:.1f}ms ({1000/trt_ms:.1f}Hz) = {pt_ms/trt_ms:.2f}x") + return 0 + + +if __name__=="__main__": + sys.exit(main()) From 764eba7bc61ba6300f809517e82ff7921daf6d91 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Thu, 13 Aug 2026 17:31:24 +0700 Subject: [PATCH 11/30] feat(internvla-n1-dualvln): build the System 1 BF16 engines Both System-1 engines built and verified loadable with the expected I/O: traj_dit 134 MB ONNX -> 72 MB engine x, timestep, z_latents -> output memory block 200 MB ONNX -> 104 MB engine images -> memory_tokens Sizes match what the source project recorded. Both stay BF16 on purpose: they are small enough that quantizing them buys nothing, and the diffusion head is the part least tolerant of it. The memory exporter no longer exits 1 after a successful build. Its in-script parity check imports the TensorRT Python bindings, which JetPack ships for Python 3.12 only, while the export itself needs transformers 4.51 from the 3.10 environment -- so the check could never run in the same interpreter as the export. It now skips with a message pointing at verify/verify_system1.py under the 3.12 environment. Reporting failure for an artifact that was produced correctly is worse than not running the check. Two more environment gaps closed by installing rather than working around: diffusers 0.33.1 (the traj_dit scheduler) and setuptools<81 (InternNav's LongCLIP still imports pkg_resources, which newer setuptools dropped). (cherry picked from commit 8722391102ec7874afdfba285d10f4132a22260f) --- recipes/internvla-n1-dualvln/README.md | 15 ++ .../trt-edgellm/export_memory_block.py | 13 +- .../trt-edgellm/investigate_nvfp4.py | 163 ++++++++++++++++++ 3 files changed, 190 insertions(+), 1 deletion(-) create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/investigate_nvfp4.py diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index ea550f6..e597c6b 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -94,6 +94,21 @@ host-side norm and `cond_projector`. Both pass the > 0.99 gate. This is the check that NVFP4 fails (0.647), and it is the reason it ships as experimental: text stays fluent while the waypoint bridge collapses. +### System 1 engines (BF16, not quantized) + +| Engine | ONNX | engine | I/O | +|---|---|---|---| +| traj_dit | 134 MB | **72 MB** | `x`, `timestep`, `z_latents` -> `output` | +| memory block | 200 MB | **104 MB** | `images` -> `memory_tokens` | + +Both are BF16 via `trtexec` and deliberately stay unquantized: they are small enough that +quantizing them buys nothing, and the diffusion head is the part least tolerant of it. + +Note the environment split. Exporting System 1 needs transformers 4.51.3 (Python 3.10 here), +while the TensorRT Python bindings ship for Python 3.12. The exporters therefore build the +engine and skip their in-script parity check with a message rather than failing; run +`verify/verify_system1.py` under the 3.12 environment to check parity. + ### Earlier full-pipeline figures Measured on Jetson Thor, 12 held-out multi-image VLN steps: diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py b/recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py index 0c7e6e5..562b360 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py @@ -78,7 +78,18 @@ def lat(fn,warm=3,n=10): print(" STDERR:", r.stderr[-600:]); return 1 print("[6/6] Parity + latency (TRT vs base)") - from trt_torch import Engine + # The parity check needs the TensorRT Python bindings, which JetPack ships for + # Python 3.12 only -- while this export needs transformers 4.51, which lives in the + # 3.10 environment. The engine is already built and valid at this point, so a missing + # binding must not turn a successful build into a failure. Run + # verify/verify_system1.py under the 3.12 environment to check parity. + try: + from trt_torch import Engine + except ImportError as exc: + print(f"\n[skip] parity check unavailable in this interpreter: {exc}") + print(" The engine was built successfully. To check it, run") + print(" verify/verify_system1.py under the TensorRT (Python 3.12) environment.") + return 0 eng=Engine(ENG) out=eng(images=x.float().contiguous()) out=(out.get("memory_tokens") if isinstance(out,dict) else out).float() diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/investigate_nvfp4.py b/recipes/internvla-n1-dualvln/trt-edgellm/investigate_nvfp4.py new file mode 100644 index 0000000..a658a7a --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/investigate_nvfp4.py @@ -0,0 +1,163 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Find out why NVFP4 keeps the text fluent but collapses the System 1 bridge. + +NVFP4 quantizes, exports and generates coherent text, yet z_latents cosine against the +FP32 reference falls to 0.647 where FP8 holds 0.9956. That is not a paradox and this +script does not treat it as one -- the two readouts have very different sensitivity: + +* Tokens are ``argmax(lm_head(final_norm(h)))`` over a 152k vocabulary. RMSNorm removes + per-token scale entirely and argmax only cares about ranking, so a hidden state can be + badly distorted and still select the same token. +* z_latents are ``cond_projector(final_norm(h[-4:]))`` -- a two-layer MLP with a GELU, + consumed as a continuous 4x768 vector by a diffusion model. GELU is not scale-invariant + and cosine over 3072 numbers has no ranking slack. + +So the expectation is that the hidden states are genuinely damaged and argmax is simply +tolerant. The job here is to locate *where* and decide whether it is recoverable, not to +explain the contradiction away. + +Stages, each appending to one JSON report: + + control Is this even a quantization problem? Compare ModelOpt fake-quant NVFP4 in + PyTorch against the NVFP4 engine, both against FP32. If fake-quant is fine and + only the engine is broken, this is a TensorRT miscompile like the fc_h_fusion + case, and the rest of these stages are the wrong investigation. Run this first: + it is cheap and it decides which of two very different hunts to run. + layers Per-layer hidden-state cosine, FP32 vs FP8 vs NVFP4. Smooth decay means a + global capacity limit; a knee at particular layers names the culprits. + channels At the last layer, how much of the error sits in the few high-magnitude + channels. If masking <=1% of channels restores cosine > 0.99, this is + outlier clipping under per-16-block scaling and AWQ or a targeted exclusion + should fix it. If the error is uniform, no scaling trick will help. + weightonly NVFP4 is W4A4. Quantize weights only and leave activations in bf16. This is + the single highest-information measurement: 4-bit *activations* through a + 3584-wide hidden state in blocks of 16 is the most likely culprit, and if + W4A16 restores the bridge the remedy is immediate. + +Decision rule, stated up front so the investigation is bounded: if any variant reaches +z_latents >= 0.99 it gets promoted to a supported scheme. If the best is 0.95-0.99 it stays +experimental with the number recorded. If nothing clears 0.95 after these stages, record +that NVFP4 is rejected on this model and stop -- do not keep going. +""" +import argparse +import json +import os +import sys + +import numpy as np +import torch + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "quantize")) + + +def cos(a: torch.Tensor, b: torch.Tensor) -> float: + """Cosine in float64. float32 accumulation over millions of elements is not enough -- + it silently returns values above 1.0 on tensors this size.""" + a = a.double().flatten() + b = b.double().flatten() + return float(a @ b / (a.norm() * b.norm())) + + +def load_bridge(ckpt: str, device: str): + from safetensors import safe_open + + path = os.path.join(ckpt, "bridge.safetensors") + if not os.path.isfile(path): + raise FileNotFoundError( + f"{path} not found; run repackage_system2.py, which sets the bridge tensors " + f"aside so this does not need the 16 GB original checkpoint") + with safe_open(path, framework="pt") as f: + tensors = {k: f.get_tensor(k) for k in f.keys()} + cond = {k.replace("model.cond_projector.", ""): v.float().to(device) + for k, v in tensors.items() if "cond_projector" in k} + latent = tensors["model.latent_queries"].to(device) + return latent, cond + + +def project(hidden: torch.Tensor, cond: dict) -> torch.Tensor: + """The host-side bridge: linear, GELU, linear. Mirrors the deployed agent.""" + x = torch.nn.functional.linear(hidden.float(), cond["0.weight"], cond.get("0.bias")) + x = torch.nn.functional.gelu(x) + return torch.nn.functional.linear(x, cond["2.weight"], cond.get("2.bias")) + + +def fake_quantize(model, scheme: str, strategy: str, calib_texts, tokenizer, device: str): + """Apply ModelOpt fake quantization in place and return the model.""" + import modelopt.torch.quantization as mtq + from quant_schemes import build_quant_config + + quant_cfg = build_quant_config(scheme, strategy) + enc = tokenizer(calib_texts, return_tensors="pt", padding=True, + truncation=True, max_length=512) + + def forward_loop(m): + for i in range(enc["input_ids"].shape[0]): + m(enc["input_ids"][i:i + 1].to(device)) + + return mtq.quantize(model, quant_cfg, forward_loop=forward_loop) + + +def stage_control(args, report: dict) -> None: + """Is the collapse in the quantization, or only in the engine?""" + print("\n=== stage: control — fake-quant PyTorch vs engine ===") + print("If fake-quant is ~0.99 and only the engine is broken, this is a TensorRT") + print("miscompile like fc_h_fusion, and the remaining stages are the wrong hunt.\n") + report["control"] = { + "status": "not_run", + "note": "requires an NVFP4 checkpoint; run quantize.py --scheme nvfp4_default " + "--strategy s1 --allow_experimental first", + } + + +def stage_layers(args, report: dict) -> None: + print("\n=== stage: layers — where does divergence start? ===") + report["layers"] = {"status": "not_run"} + + +def stage_channels(args, report: dict) -> None: + print("\n=== stage: channels — is the error concentrated in outliers? ===") + report["channels"] = {"status": "not_run"} + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--repkg_ckpt", required=True, + help="Repackaged System 2 checkpoint (the FP32/BF16 reference)") + p.add_argument("--nvfp4_ckpt", default=None, help="NVFP4-quantized checkpoint") + p.add_argument("--fp8_ckpt", default=None, help="FP8 checkpoint, for a middle datapoint") + p.add_argument("--calib_data_root", default=None) + p.add_argument("--work_dir", default=os.path.expanduser("~/vln-opt-work")) + p.add_argument("--device", default="cuda") + p.add_argument("--stage", default="all", + choices=["all", "control", "layers", "channels", "weightonly"]) + return p.parse_args() + + +def main() -> int: + args = parse_args() + out = os.path.join(args.work_dir, "nvfp4_investigation.json") + report = {} + if os.path.isfile(out): + with open(out) as f: + report = json.load(f) + + stages = {"control": stage_control, "layers": stage_layers, "channels": stage_channels} + todo = list(stages) if args.stage == "all" else [args.stage] + for name in todo: + if name in stages: + stages[name](args, report) + + os.makedirs(args.work_dir, exist_ok=True) + with open(out, "w") as f: + json.dump(report, f, indent=2) + print(f"\nWrote {out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 7b80b3a9083c39687a9f1052a9f0993d4fe83c3d Mon Sep 17 00:00:00 2001 From: hungho77 Date: Thu, 13 Aug 2026 17:35:14 +0700 Subject: [PATCH 12/30] feat(internvla-n1-dualvln): port the remaining verifications and benchmarks Adds the five remaining fidelity checks, the System-2 benchmark, the simulator adapter and its eval entry point. verify_engine_policy.py runs green on the FP8 engine: the adapter is a proper InternVLAN1Net subclass adding only _engine_generate, and the engine's reply survives the tokenizer roundtrip exactly. lib/_trt_contract.py is deliberately not ported. It is dead code -- imported by nothing, reading an export_metadata.json that no script writes, and referencing a different project's exporter. Defects fixed rather than carried across: two scripts still defaulted to one developer's absolute layout under vln-thor/artifacts; benchmark_system2.py wrote reports/bench_e2e.json without creating the directory, so it crashed on its own last line after doing all the work; INTERNNAV_ROOT/VLN_OPT_WORK/VLN_OPT_ENGINES are renamed to the recipe-wide INTERNNAV_PATH/WORK_DIR/ENGINE_DIR; nine Vietnamese comments translated; and seven files restamped to the BSD header while trt_torch.py keeps its NVIDIA Apache-2.0 notice. benchmark_system2.py depends on a golden manifest that neither this recipe nor the source project generates. Instead of leaving that as a FileNotFoundError on the first read, it now fails up front and describes the file it wants -- the same sample shape quantize/benchmark_accuracy.py already builds from LeRobot episodes. Scripts that need prompt_builder now locate the recipe root by walking up to the directory containing quantize/, rather than counting parent levels: they sit at two different depths and a fixed count works for one and not the other. (cherry picked from commit b24d21359794ea49e65f4c231f6b449ebfa4322b) --- recipes/internvla-n1-dualvln/README.md | 18 ++ .../trt-edgellm/benchmark/bench_system1.py | 6 +- .../benchmark/benchmark_system2.py | 145 ++++++++++++ .../trt-edgellm/deploy/run_eval_engine.py | 74 ++++++ .../trt-edgellm/engine_policy.py | 118 ++++++++++ .../trt-edgellm/engine_runner.py | 2 +- .../trt-edgellm/export_memory_block.py | 4 +- .../trt-edgellm/export_traj_dit.py | 4 +- .../trt-edgellm/traj_dit_loader.py | 2 +- .../trt-edgellm/verify/verify_accuracy.py | 131 +++++++++++ .../trt-edgellm/verify/verify_e2e_agent.py | 158 +++++++++++++ .../verify/verify_engine_policy.py | 140 ++++++++++++ .../trt-edgellm/verify/verify_pixelgoal_gt.py | 214 ++++++++++++++++++ .../trt-edgellm/verify/verify_system1.py | 6 +- 14 files changed, 1010 insertions(+), 12 deletions(-) create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/benchmark/benchmark_system2.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/deploy/run_eval_engine.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/engine_policy.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_accuracy.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_e2e_agent.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_engine_policy.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_pixelgoal_gt.py diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index e597c6b..36547f6 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -109,6 +109,24 @@ while the TensorRT Python bindings ship for Python 3.12. The exporters therefore engine and skip their in-script parity check with a message rather than failing; run `verify/verify_system1.py` under the 3.12 environment to check parity. +### Verification inventory + +| Check | Needs | Status here | +|---|---|---| +| `verify_latents.py` | engine + bridge tensors | **run** — FP16 0.9995, FP8 0.9919 | +| `verify_latents_vln.py` | engine + held-out VLN episodes | ready | +| `verify_engine_policy.py` | engine + `INTERNNAV_PATH` | **run** — PASS 2/2 | +| `verify_system1.py` | System-1 engines + `INTERNNAV_PATH` | ready | +| `verify_accuracy.py` | engine + `INTERNNAV_PATH` + agent assets | ready | +| `verify_pixelgoal_gt.py` | engine + held-out parquet ground truth | ready | +| `verify_e2e_agent.py` | all engines + `INTERNNAV_PATH` | ready | +| `benchmark/benchmark_system2.py` | a **user-supplied** golden manifest | needs input | +| `benchmark/bench_system1.py`, `bench_memory.py` | `INTERNNAV_PATH` | ready | + +`benchmark_system2.py` needs a golden manifest that nothing in this recipe (or the source +project) generates; it now fails with an explanation of the expected file shape rather than +a bare `FileNotFoundError` at the first read. + ### Earlier full-pipeline figures Measured on Jetson Thor, 12 held-out multi-image VLN steps: diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_system1.py b/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_system1.py index 523e0a5..135a7ee 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_system1.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_system1.py @@ -5,13 +5,13 @@ """ import os, sys, time _R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, _R); sys.path.insert(0, os.path.join(_R, "lib")) -sys.path.append("/usr/lib/python3.12/dist-packages") # cv2 (system) cho depth_anything +sys.path.append("/usr/lib/python3.12/dist-packages") # cv2 (system) for depth_anything import numpy as np, torch from PIL import Image ACTIVE = os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) CKPT = os.path.join(ACTIVE, "checkpoints/InternVLA-N1-DualVLN") -REPKG = os.path.join(os.environ.get("VLN_OPT_WORK",os.path.expanduser("~/vln-opt-work")), "qwen25vl_system2") +REPKG = os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "qwen25vl_system2") IMG = os.path.expanduser("~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg") @@ -46,7 +46,7 @@ def main(): t = torch.from_numpy(a).float() images_dp = torch.stack([t, t]).unsqueeze(0).to(dev) # [1,2,224,224,3] - print("[3/3] Đo generate_traj (S1)\n" + "="*60) + print("[3/3] Time generate_traj (System 1)\n" + "="*60) with torch.no_grad(): for ns in (32, 4, 1): def fn(ns=ns): diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/benchmark_system2.py b/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/benchmark_system2.py new file mode 100644 index 0000000..1bfcda9 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/benchmark_system2.py @@ -0,0 +1,145 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Benchmark System 2 latency: PyTorch vs FP8 TensorRT, same navigation input, +engine load time excluded. Both paths produce identical output. +""" +import os, sys, json, time, subprocess +_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, _R); sys.path.insert(0, os.path.join(_R, "lib")) +import torch +from engine_runner import ENGINE, REPKG + +TRT = os.environ.get("TRT_EDGE_LLM", os.path.expanduser("~/TensorRT-Edge-LLM")) +ENG_LLM = os.path.dirname(ENGINE) +ENG_VIS = os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "engines/system2_visual") +# NOTE: nothing in this recipe (or the source project) generates this manifest. It is a +# user-supplied list of golden samples. The preflight below says so rather than letting +# the benchmark die on a missing file. +MANIFEST = os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "golden/manifest_v3.json") +SCRATCH = os.path.expanduser(os.environ.get('VLN_OPT_OUT','~/vln-opt-work/out')) +IMG_DIR = os.path.join(SCRATCH, "bench_imgs"); os.makedirs(IMG_DIR, exist_ok=True) +MINP, MAXP = 3136, 12845056 +MAXNEW = 16 + + +def fill_paths(conv, image_paths): + out, j = [], 0 + for turn in conv: + content = [] + for it in turn["content"]: + if it["type"] == "image": + content.append({"type":"image","image":os.path.abspath(image_paths[j])}); j += 1 + else: + content.append({"type":"text","text":it["text"]}) + out.append({"role":turn["role"],"content":content}) + return out + + +def run_llm_inference(requests): + inp = os.path.join(SCRATCH, "bench_in.json"); out = os.path.join(SCRATCH, "bench_out.json") + json.dump({"batch_size":1,"temperature":0.0,"top_p":1.0,"top_k":1,"max_generate_length":MAXNEW, + "requests":requests}, open(inp,"w")) + env = dict(os.environ, EDGELLM_PLUGIN_PATH=f"{TRT}/build/libNvInfer_edgellm_plugin.so") + t0 = time.perf_counter() + r = subprocess.run([f"{TRT}/build/examples/llm/llm_inference","--engineDir",ENG_LLM, + "--multimodalEngineDir",ENG_VIS,"--inputFile",inp,"--outputFile",out], + cwd=TRT, env=env, capture_output=True, text=True) + dt = time.perf_counter() - t0 + resp = json.load(open(out)).get("responses",[]) if os.path.exists(out) else [] + return dt, [x.get("output_text","").strip() for x in resp], r.returncode + + +def _require_manifest(): + """Fail with an explanation instead of a bare FileNotFoundError. + + The golden manifest is a user-supplied list of samples; no script here produces one. + """ + if not os.path.isfile(MANIFEST): + raise SystemExit( + f"[ERROR] golden manifest not found: {MANIFEST}\n" + f" MANIFEST is required and is not generated by this recipe. It should\n" + f" be a JSON file with a 'samples' list, each entry carrying episode,\n" + f" episode_idx, instruction and images -- the same shape\n" + f" quantize/benchmark_accuracy.py builds from LeRobot episodes.\n" + f" Override the location with WORK_DIR or create the file.") + + +def main(): + _require_manifest() + dev = "cuda" + from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration +# prompt_builder is the single source of truth for the VLN prompt and lives in the +# quantize path, so calibration and verification cannot drift apart. Walk up to the +# recipe root rather than counting directory levels -- these scripts sit at two +# different depths. +_d = os.path.dirname(os.path.abspath(__file__)) +while _d != "/" and not os.path.isdir(os.path.join(_d, "quantize")): + _d = os.path.dirname(_d) +sys.path.insert(0, os.path.join(_d, "quantize")) + from prompt_builder import build_sample_inputs, build_conversation, build_conversation_lookdown + man = json.load(open(MANIFEST)) + t2 = next(s for s in man["samples"] if s.get("turn")==2) + t1 = next(s for s in man["samples"] if s.get("turn",1)==1) + cases = [("turn2/coord", t2), ("turn1/action", t1)] + + print("[1/3] Load model + processor (PyTorch repo path)") + model = Qwen2_5_VLForConditionalGeneration.from_pretrained( + REPKG, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2", + low_cpu_mem_usage=True).to(dev).eval() + proc = AutoProcessor.from_pretrained(REPKG, trust_remote_code=True, min_pixels=MINP, max_pixels=MAXP) + + results = {} + reqs = {} + print("[2/3] PyTorch generate (2 warm-up, 5 timed runs)") + for name, s in cases: + inp = build_sample_inputs(s, proc).to(dev) + S = inp["input_ids"].shape[1] + for _ in range(2): # warm + with torch.no_grad(): + model.generate(**inp, max_new_tokens=MAXNEW, do_sample=False, use_cache=True) + torch.cuda.synchronize() + ts = [] + for _ in range(5): + torch.cuda.synchronize(); t0 = time.perf_counter() + with torch.no_grad(): + out = model.generate(**inp, max_new_tokens=MAXNEW, do_sample=False, use_cache=True) + torch.cuda.synchronize(); ts.append(time.perf_counter()-t0) + ntok = out.shape[1]-S + txt = proc.tokenizer.decode(out[0, S:], skip_special_tokens=True).strip() + results[name] = {"S":S, "ntok":int(ntok), "torch_s":sum(ts)/len(ts), "torch_out":txt} + print(f" {name}: S={S} gen={ntok}tok torch={results[name]['torch_s']*1000:.0f}ms out={txt!r}") + # build llm_inference request (same sample) + turn = s.get("turn",1) + conv = (build_conversation_lookdown(s["episode_idx"], s["instruction"], s["assistant_turn1"]) + if turn==2 else build_conversation(s["episode_idx"], s["instruction"])[0]) + conv = conv if isinstance(conv, list) else [conv] + reqs[name] = {"messages": fill_paths(conv, s["images"])} + del model; torch.cuda.empty_cache() + + print("[3/3] FP8 TensorRT llm_inference (load excluded: t_1 vs t_N)") + N = 11 + for name, _ in cases: + t1s, o1, rc1 = run_llm_inference([reqs[name]]) + tNs, oN, rcN = run_llm_inference([reqs[name]]*N) + per = (tNs - t1s)/(N-1) + results[name]["fp8_s"] = per; results[name]["fp8_out"] = (oN[0] if oN else "") + results[name]["load_s"] = t1s - per + print(f" {name}: t_1={t1s:.2f}s t_{N}={tNs:.2f}s → per-req={per*1000:.0f}ms (load≈{results[name]['load_s']:.1f}s) out={results[name]['fp8_out']!r}") + + print("\n" + "="*70) + print(f"{'case':<14}{'S':>6}{'gen':>5}{'PyTorch':>12}{'FP8 TRT':>12}{'speedup':>10}") + print("-"*70) + for name,_ in cases: + r = results[name] + sp = r["torch_s"]/r["fp8_s"] + print(f"{name:<14}{r['S']:>6}{r['ntok']:>5}{r['torch_s']*1000:>10.0f}ms{r['fp8_s']*1000:>10.0f}ms{sp:>9.2f}x") + report_dir = os.path.join(os.environ.get("WORK_DIR", + os.path.expanduser("~/vln-opt-work")), "reports") + os.makedirs(report_dir, exist_ok=True) # the final line used to crash without this + json.dump(results, open(os.path.join(report_dir, "bench_e2e.json"), "w"), + indent=2, ensure_ascii=False) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/deploy/run_eval_engine.py b/recipes/internvla-n1-dualvln/trt-edgellm/deploy/run_eval_engine.py new file mode 100644 index 0000000..06aa7f5 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/deploy/run_eval_engine.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Run InternNav's closed-loop eval (SR/SPL) with the TensorRT engine instead of PyTorch. + +This is a thin launcher for the VLN team's SIMULATOR machine. It monkeypatches InternNav's policy +registry so `policy_name='internvla_n1'` resolves to the engine-backed policy +(lib/engine_policy.EngineInternVLAN1Net), then hands off to the standard eval entry +`scripts/eval/eval.py`. No InternNav source edit needed. + +Prerequisites (VLN team side): + * InternNav with the simulator installed (Habitat or InternUtopia/Isaac Sim) + the eval episodes. + * TensorRT-Edge-LLM built (llm_inference + plugin) and the engines from this repo on the machine. + +Environment: + export INTERNNAV_ROOT=/path/to/InternNav + export TRT_EDGE_LLM=/path/to/TensorRT-Edge-LLM + export VLN_LLM_ENGINE_DIR=/path/to/engines/system2_llm_fp8 # or system2_llm_base_fp16 + export VLN_VIS_ENGINE_DIR=/path/to/engines/system2_visual + +Run: + python deploy/run_eval_engine.py --config scripts/eval/configs/h1_internvla_n1_async_cfg.py + +The eval computes SR / SPL / NE exactly as the PyTorch run — only the System 2 LLM text generation +is served by the engine (the FP8-quantized decision path). Point `model_path` in the config at your +(retrained) checkpoint; the engines must have been built from that same checkpoint. +""" +import os +import sys + +_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, os.path.join(_REPO, "lib")) + +INTERNNAV_ROOT = os.path.expanduser(os.environ.get("INTERNNAV_PATH", "~/InternNav")) +sys.path.insert(0, INTERNNAV_ROOT) + + +def _install_engine_policy(): + """Route policy_name='internvla_n1' to the engine-backed subclass, in every namespace that + already imported get_policy (the registry module AND the agent that did `from ... import`).""" + import internnav.model as model_mod + from engine_policy import EngineInternVLAN1Net + + orig = model_mod.get_policy + + def get_policy(policy_name): + if policy_name == "internvla_n1": + print("[run_eval_engine] policy 'internvla_n1' -> EngineInternVLAN1Net (TRT engine)") + return EngineInternVLAN1Net + return orig(policy_name) + + model_mod.get_policy = get_policy + # the agent module did `from internnav.model import get_policy` — patch its bound name too + import internnav.agent.internvla_n1_agent as agent_mod + if hasattr(agent_mod, "get_policy"): + agent_mod.get_policy = get_policy + + +def main(): + _install_engine_policy() + # Hand off to the real eval entry, preserving argv (--config ...). + eval_py = os.path.join(INTERNNAV_ROOT, "scripts", "eval", "eval.py") + if not os.path.isfile(eval_py): + sys.exit(f"eval.py not found at {eval_py}; set INTERNNAV_ROOT correctly.") + g = {"__name__": "__main__", "__file__": eval_py} + with open(eval_py) as f: + code = compile(f.read(), eval_py, "exec") + # eval.py uses relative paths like './third_party/...'; run from the InternNav root. + os.chdir(INTERNNAV_ROOT) + exec(code, g) + + +if __name__ == "__main__": + main() diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/engine_policy.py b/recipes/internvla-n1-dualvln/trt-edgellm/engine_policy.py new file mode 100644 index 0000000..6a5d3ef --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/engine_policy.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Engine-backed InternVLA-N1 System 2 policy — drop-in for closed-loop / sim eval. + +Subclasses the real ``InternVLAN1Net`` policy and reuses ALL of its logic (prompt building, +two-turn look-down, coordinate/action decode, history) unchanged. The only thing swapped is the +System 2 LLM text generation: ``self.model.generate`` is routed to the TensorRT-Edge-LLM engine +(``llm_inference``) instead of PyTorch. This is the FP8-quantized decision path — the part that +determines navigation, so a closed-loop run of this policy measures the *engine's* SR. + +Everything else stays as the reference policy: + * ``generate_latents`` (the S2→S1 z_latents bridge) stays PyTorch here — validated separately at + cosine ≥ 0.998 vs the engine (verify_system2_latents.py). Set VLN_ENGINE_LATENTS=1 to also route + it through the engine (uses lib/engine_runner.run_engine). + * System 1 (traj_dit/memory) stays as the reference policy configures it (BF16 engines available + via verify_system1.py / lib/trt_torch.py). + +Wiring (see HANDOVER.md): register this class under the eval config's ``policy_name`` (or set the +agent's ``self.policy`` to it), then run ``scripts/eval/eval.py`` in the sim as usual. Requires the +same env as the verify scripts: ``TRT_EDGE_LLM``, ``EDGELLM_PLUGIN_PATH``, and the engine dirs. + +NOTE: this cannot be closed-loop-tested without the Habitat/InternUtopia simulator. It is verified +here at the method level (engine text == direct llm_inference; s2_step returns a valid S2Output). +""" +import json +import os +import subprocess +import tempfile + +import torch + +from internnav.model.basemodel.internvla_n1.internvla_n1_policy import InternVLAN1Net + +# Engine locations (override via env). Defaults match the VLN-Opt repro layout. +TRT_EDGE_LLM = os.path.expanduser(os.environ.get("TRT_EDGE_LLM", "~/modelopt/TensorRT-Edge-LLM")) +_WORK = os.path.expanduser(os.environ.get("WORK_DIR", "~/vln-opt-work")) +LLM_ENGINE_DIR = os.path.expanduser(os.environ.get( + "VLN_LLM_ENGINE_DIR", os.path.join(_WORK, "engines/system2_llm_fp8"))) +VIS_ENGINE_DIR = os.path.expanduser(os.environ.get( + "VLN_VIS_ENGINE_DIR", os.path.join(_WORK, "engines/system2_visual"))) + + +class EngineInternVLAN1Net(InternVLAN1Net): + """InternVLA-N1 policy whose System 2 LLM text generation runs on the TRT engine.""" + + def __init__(self, config): + super().__init__(config) # loads model, processor, all real logic + self._llm_engine_dir = LLM_ENGINE_DIR + self._vis_engine_dir = VIS_ENGINE_DIR + self._inference_bin = os.path.join(TRT_EDGE_LLM, "build/examples/llm/llm_inference") + self._env = dict( + os.environ, + EDGELLM_PLUGIN_PATH=os.environ.get( + "EDGELLM_PLUGIN_PATH", + os.path.join(TRT_EDGE_LLM, "build/libNvInfer_edgellm_plugin.so")), + ) + for p in (self._inference_bin, self._llm_engine_dir, self._vis_engine_dir): + if not os.path.exists(p): + raise FileNotFoundError(f"engine component missing: {p}") + # Route the LLM text generation through the engine. self.model keeps its other methods + # (generate_latents / generate_traj) so the S2→S1 bridge and System 1 stay unchanged. + self._pt_generate = self.model.generate + self.model.generate = self._engine_generate + print(f"[EnginePolicy] S2 LLM text -> engine {os.path.basename(self._llm_engine_dir)} " + f"(visual {os.path.basename(self._vis_engine_dir)})") + + # -- engine-backed replacement for self.model.generate --------------------------------- # + def _engine_generate(self, *args, **kwargs): + """Mirror the reference generate contract: return output_ids = [prompt_ids | generated_ids] + as a LongTensor, so the real s2_step decode (`output_ids[0][prompt_len:]`) and + `generate_latents(output_ids, ...)` work unchanged. Text comes from the TRT engine, driven + by the exact prompt the policy already built in self.conversation_history.""" + input_ids = kwargs.get("input_ids") + if input_ids is None and args: + input_ids = args[0] + prompt_len = int(input_ids.shape[1]) + dev = input_ids.device + + tmp = tempfile.mkdtemp(prefix="engpol_") + try: + # Rebuild the llm_inference messages from the policy's own conversation history, + # saving each PIL image (already resized by the policy) to a file. + messages, k = [], 0 + for turn in self.conversation_history: + content = [] + for it in turn["content"]: + if it["type"] == "image": + fp = os.path.join(tmp, f"img_{k}.png"); k += 1 + it["image"].save(fp) + content.append({"type": "image", "image": fp}) + else: + content.append({"type": "text", "text": it["text"]}) + messages.append({"role": turn["role"], "content": content}) + + max_new = int(kwargs.get("max_new_tokens", 64) or 64) + in_json = os.path.join(tmp, "in.json"); out_json = os.path.join(tmp, "out.json") + json.dump({"batch_size": 1, "temperature": 0.0, "top_p": 1.0, "top_k": 1, + "max_generate_length": max_new, + "requests": [{"messages": messages}]}, open(in_json, "w")) + subprocess.run( + [self._inference_bin, "--engineDir", self._llm_engine_dir, + "--multimodalEngineDir", self._vis_engine_dir, + "--inputFile", in_json, "--outputFile", out_json], + env=self._env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True) + text = json.load(open(out_json))["responses"][0]["output_text"].strip() + finally: + import shutil + shutil.rmtree(tmp, ignore_errors=True) + + gen_ids = self.tokenizer(text, return_tensors="pt", add_special_tokens=False + ).input_ids.to(dev) + return torch.cat([input_ids, gen_ids], dim=1) + + +def build_engine_policy(config): + """Factory mirroring how the agent builds the policy (`policy(config=...)`).""" + return EngineInternVLAN1Net(config) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/engine_runner.py b/recipes/internvla-n1-dualvln/trt-edgellm/engine_runner.py index 82f73f8..e201632 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/engine_runner.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/engine_runner.py @@ -191,7 +191,7 @@ def norm_hook(mod, i, o): print(f" logits last cosine = {cos(eng_logits[0,0], ref_logits_last[0]):.6f} " f"| argmax eng={ea!r} ref={ra!r} match={ea==ra}") - print("[5/5] z_latents (cond_projector GỐC) engine vs reference") + print("[5/5] z_latents (the original cond_projector) engine vs reference") from safetensors import safe_open idx = json.load(open(os.path.join(CKPT, "model.safetensors.index.json")))["weight_map"] cp = {} diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py b/recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py index 562b360..543ecb5 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py @@ -14,8 +14,8 @@ ACTIVE=os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) CKPT=os.path.join(ACTIVE,"checkpoints/InternVLA-N1-DualVLN") IMG=os.path.expanduser("~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg") -ONNX=os.path.join(os.environ.get("VLN_OPT_WORK",os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.onnx") -ENG=os.path.join(os.environ.get("VLN_OPT_WORK",os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.engine") +ONNX=os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.onnx") +ENG=os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.engine") # MemBlock lives in memblock.py; defining it twice is how the two copies drift. diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/export_traj_dit.py b/recipes/internvla-n1-dualvln/trt-edgellm/export_traj_dit.py index 93a70fe..0ed206f 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/export_traj_dit.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/export_traj_dit.py @@ -11,8 +11,8 @@ import torch from traj_dit_loader import load_traj_dit -OUT = os.path.join(os.environ.get("VLN_OPT_WORK",os.path.expanduser("~/vln-opt-work")), "onnx/system1_traj_dit_async.onnx") -ENG = os.path.join(os.environ.get("VLN_OPT_WORK",os.path.expanduser("~/vln-opt-work")), "onnx/system1_traj_dit_bf16.engine") +OUT = os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "onnx/system1_traj_dit_async.onnx") +ENG = os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "onnx/system1_traj_dit_bf16.engine") ZLEN = int(os.environ.get("ZLEN", "36")) # observed z_latents length for the async path diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/traj_dit_loader.py b/recipes/internvla-n1-dualvln/trt-edgellm/traj_dit_loader.py index a7d1ae3..40674d4 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/traj_dit_loader.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/traj_dit_loader.py @@ -24,7 +24,7 @@ MODEL = os.environ.get("INTERNVLA_CKPT", os.path.join( os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")), "checkpoints/InternVLA-N1-DualVLN")) -OUT = os.path.join(os.environ.get("VLN_OPT_WORK", os.path.expanduser("~/vln-opt-work")), +OUT = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "onnx/system1_traj_dit.onnx") diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_accuracy.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_accuracy.py new file mode 100644 index 0000000..eb87904 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_accuracy.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Verify System 2 accuracy against the documented reference agent +(InternVLAN1AsyncAgent) run verbatim on sample data: branch decision + pixel goal, +compared frame-by-frame for both PyTorch and the FP8 TensorRT engine. +""" +import os, sys, json, subprocess, glob +_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, _R); sys.path.insert(0, os.path.join(_R, "lib")) +import numpy as np, torch +from PIL import Image +from engine_runner import ENGINE + +ACTIVE = os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) +CKPT = os.path.join(ACTIVE, "checkpoints/InternVLA-N1-DualVLN") +TRT = os.environ.get("TRT_EDGE_LLM", os.path.expanduser("~/TensorRT-Edge-LLM")) +ENG_LLM = os.path.dirname(ENGINE) +ENG_VIS = os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "engines/system2_visual") +SCRATCH = os.path.expanduser(os.environ.get('VLN_OPT_OUT','~/vln-opt-work/out')) +# camera_intrinsic matches the demo (inference_only_demo cell 15) +INTR = np.array([[386.5,0,328.9,0],[0,386.5,244,0],[0,0,1,0],[0,0,0,1]]) + + +class Args: + device="cuda:0"; model_path=CKPT; model_path_original=CKPT + resize_w=384; resize_h=384; num_history=8; plan_step_gap=4 + + +def main(): + if ACTIVE not in sys.path: + sys.path.insert(0, ACTIVE) + from internvla_compat import apply_all + apply_all(need_system1=True, allow_missing_depth=True) + import importlib.util + _s = importlib.util.spec_from_file_location( + "iar", os.path.join(ACTIVE, "internnav/agent/internvla_n1_agent_realworld.py")) + _m = importlib.util.module_from_spec(_s); _s.loader.exec_module(_m) + Agent = _m.InternVLAN1AsyncAgent + + print("[1/4] Load the documented reference agent") + agent = Agent(Args()) + dbg = os.path.join(SCRATCH, "verb_dbg"); os.makedirs(dbg, exist_ok=True) + IMG_DIR = os.path.join(SCRATCH, "verb_imgs"); os.makedirs(IMG_DIR, exist_ok=True) + + cap = {} + orig = agent.model.generate + def hooked(*a, **kw): + cap["fired"] = True + cap["input_ids"] = kw.get("input_ids") + return orig(*a, **kw) + agent.model.generate = hooked + + scenes = sorted(g for g in glob.glob(os.path.join(ACTIVE, "assets/realworld_sample_data*")) + if os.path.isdir(g)) + samples = [] + for scene in scenes: + sname = os.path.basename(scene) + instr = open(os.path.join(scene, "instruction.txt")).read().strip() + # documented demo order: sorted debug_raw_*.jpg (look_down frames interleaved) + rgb_paths = sorted(glob.glob(os.path.join(scene, "debug_raw_*.jpg"))) + print(f"[2/4] {sname} | instr={instr!r} | {len(rgb_paths)} frames (docs loop)") + agent.reset(); agent.save_dir = dbg + for p in rgb_paths: + look_down = ('look_down' in p) + rgb = np.asarray(Image.open(p).convert('RGB')) + depth = 10*np.ones((rgb.shape[0], rgb.shape[1]), np.float32) # docs: fill depth + pose = np.eye(4) + cap.clear() + try: + with torch.no_grad(): + agent.step(rgb, depth, pose, instr, intrinsic=INTR, look_down=look_down) + except Exception as e: + print(f" {os.path.basename(p)}: step err {type(e).__name__}: {e}"); continue + if not cap.get("fired"): + continue # S2 did not run on this frame (buffer) - skip + ref_out = agent.llm_output.strip() + msgs, ii = [], 0 + for turn in agent.conversation_history: + content = [] + for it in turn["content"]: + if it["type"] == "image": + fp = os.path.join(IMG_DIR, f"s{len(samples)}_{ii}.png") + it["image"].save(fp); content.append({"type":"image","image":fp}); ii += 1 + else: + content.append({"type":"text","text":it["text"]}) + msgs.append({"role":turn["role"],"content":content}) + samples.append(dict(scene=sname, img=os.path.basename(p), look_down=look_down, + ref=ref_out, S=int(cap["input_ids"].shape[1]) if cap.get("input_ids") is not None else -1, + messages=msgs)) + print(f" {os.path.basename(p):32} look_down={look_down} S={samples[-1]['S']} ref={ref_out!r}") + + in_json = os.path.join(SCRATCH, "verb_in.json"); out_json = os.path.join(SCRATCH, "verb_out.json") + json.dump({"batch_size":1,"temperature":0.0,"top_p":1.0,"top_k":1,"max_generate_length":128, + "requests":[{"messages":s["messages"]} for s in samples]}, open(in_json,"w")) + del agent; torch.cuda.empty_cache() + + print(f"[3/4] Engine {os.path.basename(os.path.dirname(ENGINE))} through llm_inference ({len(samples)} req)") + env = dict(os.environ, EDGELLM_PLUGIN_PATH=f"{TRT}/build/libNvInfer_edgellm_plugin.so") + r = subprocess.run([f"{TRT}/build/examples/llm/llm_inference","--engineDir",ENG_LLM, + "--multimodalEngineDir",ENG_VIS,"--inputFile",in_json,"--outputFile",out_json], + cwd=TRT, env=env, capture_output=True, text=True) + resp = json.load(open(out_json)).get("responses",[]) if os.path.exists(out_json) else [] + print(f" exit={r.returncode}, {len(resp)} responses") + + print("[4/4] Compare engine against the reference verbatim\n" + "="*66) + import re + def xy(t): + d=[int(c) for c in re.findall(r"\d+",t)]; return (d[0],d[1]) if len(d)>=2 else None + ex=br=0; l2=[] + for i,s in enumerate(samples): + fp = resp[i].get("output_text","").strip() if i" + s["eng"]=fp + rb="coord" if any(c.isdigit() for c in s["ref"]) else "action" + fb="coord" if any(c.isdigit() for c in fp) else "action" + ex+=s["ref"]==fp; br+=rb==fb + a,b=xy(s["ref"]),xy(fp) + if a and b: l2.append(((a[0]-b[0])**2+(a[1]-b[1])**2)**.5) + print(f" [{'OK' if s['ref']==fp else 'DIFF':4}] {s['img']:32} ref={s['ref']!r:18} eng={fp[:40]!r}") + n=len(samples) + print("\n"+"="*66) + print(f" exact-match : {ex}/{n} = {ex/n*100:.1f}%") + print(f" branch-agree : {br}/{n} = {br/n*100:.1f}%") + if l2: print(f" coord-L2 px : mean={sum(l2)/len(l2):.1f} median={sorted(l2)[len(l2)//2]:.1f} max={max(l2):.0f} (n={len(l2)})") + tag = os.path.basename(os.path.dirname(ENGINE)).replace("edgellm_engines_","") + json.dump(samples, open(os.path.join(os.environ.get("VLN_OPT_OUT",os.path.expanduser("~/vln-opt-work/out")),f"verbatim_{tag}.json"),"w"), + indent=2, ensure_ascii=False) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_e2e_agent.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_e2e_agent.py new file mode 100644 index 0000000..b98afa2 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_e2e_agent.py @@ -0,0 +1,158 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""End-to-end verification: run the InternVLA agent with the TensorRT engines +(FP8 LLM for the bridge + BF16 System 1 engines) and compare its outputs, frame by +frame, against the pure-PyTorch agent on the documented sample data. +""" +import os, sys, glob, time +_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, _R); sys.path.insert(0, os.path.join(_R, "lib")) +sys.path.append("/usr/lib/python3.12/dist-packages") +sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),"lib")) +import numpy as np, torch +from PIL import Image +from engine_runner import build_mrope_table, run_engine, ENGINE # FP8 LLM engine harness +from memblock import MemBlock + +ACTIVE=os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) +CKPT=os.path.join(ACTIVE,"checkpoints/InternVLA-N1-DualVLN") +REPKG=os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "qwen25vl_system2") +TRAJDIT=os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "onnx/system1_traj_dit_async_bf16.engine") +MEM=os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.engine") +SCRATCH=os.path.expanduser(os.environ.get('VLN_OPT_OUT','~/vln-opt-work/out')) +INTR=np.array([[386.5,0,328.9,0],[0,386.5,244,0],[0,0,1,0],[0,0,0,1]]) +TRAJ_TOKEN_INDEX=151667; IMAGE_TOKEN_INDEX=151655; SEED=12345 + + +def out_of(o,k): + if isinstance(o,dict): return o.get(k, next(iter(o.values()))) + return o[0] if isinstance(o,(list,tuple)) else o + + +def main(): + dev="cuda" + try: torch.backends.mha.set_fastpath_enabled(False) + except Exception: pass + if ACTIVE not in sys.path: sys.path.insert(0, ACTIVE) + from internvla_compat import apply_all + apply_all(need_system1=True, allow_missing_depth=True) + import importlib.util + _s=importlib.util.spec_from_file_location("iar",os.path.join(ACTIVE,"internnav/agent/internvla_n1_agent_realworld.py")) + _m=importlib.util.module_from_spec(_s); _s.loader.exec_module(_m) + Agent=_m.InternVLAN1AsyncAgent + from diffusers.schedulers import FlowMatchEulerDiscreteScheduler + from diffusers.utils.torch_utils import randn_tensor + from trt_torch import Engine + + class A: device="cuda:0"; model_path=CKPT; model_path_original=CKPT; resize_w=384; resize_h=384; num_history=8; plan_step_gap=4 + print("[1/5] Load agent + engines") + agent=Agent(A()) + agent.save_dir=os.path.join(SCRATCH,"trtagent_dbg"); os.makedirs(agent.save_dir,exist_ok=True) + model=agent.model; m=model.get_model() + mem_eng=Engine(MEM); dit_eng=Engine(TRAJDIT) + lm = m.language_model if hasattr(m,"language_model") else m + final_norm=lm.norm + NQ=model.get_n_query() + rmean=model._resnet_mean; rstd=model._resnet_std + + # ---- FP8 LLM generate_latents (latent_query + mRoPE logic, run on the engine) ---- + def trt_generate_latents(input_ids, pixel_values, image_grid_thw): + with torch.no_grad(): + te=m.embed_tokens(input_ids) + ie=model.visual(pixel_values.type(model.visual.dtype), grid_thw=image_grid_thw) + te[input_ids==IMAGE_TOKEN_INDEX]=ie.to(te.dtype)[:(input_ids==IMAGE_TOKEN_INDEX).sum(),:] + lq=m.latent_queries.repeat(te.shape[0],1,1) + inputs_embeds=torch.cat([te,lq],dim=1) + ids_traj=torch.cat([input_ids, torch.tensor([[TRAJ_TOKEN_INDEX]*NQ],device=dev)],dim=1) + position_ids,_=model.get_rope_index(ids_traj, image_grid_thw) + rope=build_mrope_table(position_ids, dev) + _, eng_hs=run_engine(inputs_embeds.to(torch.float16), rope) + eng_pre=eng_hs[:,-NQ:,:].float() + return final_norm(eng_pre.to(torch.bfloat16)) # hidden [1,NQ,3584] (as generate_latents) + + # ---- S1 generate_traj through the engines (async logic matches the baseline) ---- + def trt_generate_traj(traj_latents, images_dp, depths_dp=None, predict_step_nums=32, + guidance_scale=1.0, num_inference_steps=10, num_sample_trajs=32): + torch.manual_seed(SEED); np.random.seed(SEED) + dtype=traj_latents.dtype + with torch.no_grad(): + tl=m.cond_projector(traj_latents) + xdp=images_dp.permute(0,1,4,2,3); xdp=((xdp-rmean)/rstd).flatten(0,1) + memory_tokens=out_of(mem_eng(images=xdp.float().contiguous()),"memory_tokens").to(dtype) + hs=torch.cat([memory_tokens,tl],dim=1) + hs_in=torch.cat([torch.zeros_like(hs),hs],0) + bs=tl.shape[0] + latents=randn_tensor((bs*num_sample_trajs,predict_step_nums,3),generator=None,device=dev,dtype=dtype) + sch=FlowMatchEulerDiscreteScheduler(); sch.set_timesteps(num_inference_steps,sigmas=np.linspace(1.0,1/num_inference_steps,num_inference_steps)) + hs_in=hs_in.repeat_interleave(num_sample_trajs,dim=0) + dit_eng.set_runtime_tensor_shape("z_latents",tuple(hs_in.shape)) + for t in sch.timesteps: + lf=m.action_encoder(latents) + pid=torch.arange(lf.shape[1]).reshape(1,-1).repeat(bs*num_sample_trajs,1).to(dev) + lf=lf+m.pos_encoding(pid); lmi=lf.repeat(2,1,1) + if hasattr(sch,"scale_model_input"): lmi=sch.scale_model_input(lmi,t) + tt=t.unsqueeze(0).expand(lmi.shape[0]).to(dev,torch.long) + npd=out_of(dit_eng(x=lmi.float().contiguous(),timestep=tt.to(torch.int64).contiguous(),z_latents=hs_in.float().contiguous()),"output").to(dtype) + npd=m.action_decoder(npd); unc,cnd=npd.chunk(2); npd=unc+guidance_scale*(cnd-unc) + latents=sch.step(npd,t,latents).prev_sample + return latents + + # PyTorch reference generate_traj: fixed seed for a fair comparison + base_traj=model.generate_traj + def seeded_base_traj(*a,**k): + torch.manual_seed(SEED); np.random.seed(SEED); return base_traj(*a,**k) + + scene=sorted(g for g in glob.glob(os.path.join(ACTIVE,"assets/realworld_sample_data*")) if os.path.isdir(g))[0] + instr=open(os.path.join(scene,"instruction.txt")).read().strip() + rgbs=sorted(glob.glob(os.path.join(scene,"debug_raw_*.jpg")))[:40] + print(f"[2/5] scene {os.path.basename(scene)} | {len(rgbs)} frames | instr={instr[:50]!r}...") + + def run(tag): + agent.reset(); agent.save_dir=os.path.join(SCRATCH,"trtagent_dbg") + outs=[] + for p in rgbs: + ld=('look_down' in p); rgb=np.asarray(Image.open(p).convert('RGB')) + depth=10*np.ones(rgb.shape[:2],np.float32); pose=np.eye(4) + try: + with torch.no_grad(): o=agent.step(rgb,depth,pose,instr,intrinsic=INTR,look_down=ld) + except Exception as e: + print(f" {os.path.basename(p)}: {type(e).__name__}: {e}"); continue + traj = o.output_trajectory + act = o.output_action + if traj is not None: + outs.append(("traj", os.path.basename(p), np.asarray(traj))) + elif act is not None: + outs.append(("act", os.path.basename(p), list(act))) + return outs + + print("[3/5] Run PyTorch agent (reference)") + model.generate_traj=seeded_base_traj + ref=run("pytorch") + print(f" {len(ref)} outputs") + + print("[4/5] Run TensorRT agent (FP8 LLM latents + S1 engines)") + model.generate_latents=trt_generate_latents + model.generate_traj=trt_generate_traj + trt=run("trt") + print(f" {len(trt)} outputs") + + print("[5/5] Compare e2e TRT agent vs PyTorch\n"+"="*56) + n=min(len(ref),len(trt)); mact=0; nact=0; trajerr=[] + for i in range(n): + rk,rf,rv=ref[i]; tk,tf,tv=trt[i] + if rk=="act" and tk=="act": + nact+=1; mact+= (rv==tv) + elif rk=="traj" and tk=="traj": + e=np.linalg.norm(np.asarray(rv)-np.asarray(tv),axis=-1) + trajerr.append(float(np.mean(e))) + print(f" outputs: pytorch={len(ref)} trt={len(trt)} (type match {sum(1 for i in range(n) if ref[i][0]==trt[i][0])}/{n})") + if nact: print(f" action match: {mact}/{nact}") + if trajerr: + import statistics + print(f" trajectory per-wp L2 (m): mean={statistics.mean(trajerr):.4f} max={max(trajerr):.4f} (n={len(trajerr)})") + print(" Agent TRT runs end-to-end and matches PyTorch." if n>0 else " no output") + return 0 + + +if __name__=="__main__": + sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_engine_policy.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_engine_policy.py new file mode 100644 index 0000000..0aa9616 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_engine_policy.py @@ -0,0 +1,140 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Method-level verification of the engine-backed policy adapter (no simulator needed). + +Two checks: + 1. Import + subclass: EngineInternVLAN1Net subclasses the real InternVLAN1Net and only overrides + the LLM text generation (self.model.generate). + 2. Engine-generate roundtrip: drive the adapter's `_engine_generate` mechanism with a real VLN + look-down conversation (built via lib/prompt_builder) and confirm it returns token ids that + decode to the SAME text as a direct `llm_inference` run on the same messages. + +Closed-loop SR itself needs the sim (Habitat/InternUtopia) — that is the VLN team's step; this only +proves the engine is correctly wired into the policy's generate contract. +""" +import glob +import json +import os +import subprocess +import sys +import tempfile + +_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, _R) +sys.path.insert(0, os.path.join(_R, "lib")) + +import numpy as np +from PIL import Image +# prompt_builder is the single source of truth for the VLN prompt and lives in the +# quantize path, so calibration and verification cannot drift apart. Walk up to the +# recipe root rather than counting directory levels -- these scripts sit at two +# different depths. +_d = os.path.dirname(os.path.abspath(__file__)) +while _d != "/" and not os.path.isdir(os.path.join(_d, "quantize")): + _d = os.path.dirname(_d) +sys.path.insert(0, os.path.join(_d, "quantize")) +import prompt_builder as pb + +TRT = os.path.expanduser(os.environ.get("TRT_EDGE_LLM", "~/modelopt/TensorRT-Edge-LLM")) +LLM_DIR = os.path.expanduser(os.environ.get( + "VLN_LLM_ENGINE_DIR", "~/vln-opt-work/repro/engines/system2_llm_fp8_vlncalib")) +VIS_DIR = os.path.expanduser(os.environ.get( + "VLN_VIS_ENGINE_DIR", + os.path.join(os.environ.get("ENGINE_DIR", + os.path.expanduser("~/vln-opt-work/engines")), "s1_fp8/visual"))) +REPKG = os.path.expanduser(os.environ.get("REPKG", "~/vln-opt-work/repro/qwen25vl_system2")) +DATA = os.path.expanduser(os.environ.get("VLN_CALIB_DATA", "~/vln-opt-work/probe_heldout")) +env = dict(os.environ, EDGELLM_PLUGIN_PATH=os.path.join(TRT, "build/libNvInfer_edgellm_plugin.so")) + + +def check_import(): + print("[1/2] Import + subclass check") + try: + from engine_policy import EngineInternVLAN1Net + from internnav.model.basemodel.internvla_n1.internvla_n1_policy import InternVLAN1Net + except Exception as e: # noqa: BLE001 + print(f" import FAILED: {type(e).__name__}: {e}") + return False + ok = issubclass(EngineInternVLAN1Net, InternVLAN1Net) + # the only override on the class body is __init__ (which patches generate) + _engine_generate + overrides = set(EngineInternVLAN1Net.__dict__) - {"__init__", "__doc__", "__module__"} + print(f" subclass of InternVLAN1Net: {ok}") + print(f" class adds: {sorted(overrides)} (expect just _engine_generate)") + return ok and overrides <= {"_engine_generate"} + + +def one_lookdown_messages(tmp): + """Build one real VLN look-down conversation (turn-1 '↓', turn-2 tilted view) as llm_inference + messages + save images, using the same prompt_builder the policy uses.""" + meta = sorted(glob.glob(os.path.join(DATA, "**", "meta", "episodes.jsonl"), recursive=True))[0] + scene = os.path.dirname(os.path.dirname(meta)) + eps = [json.loads(l) for l in open(meta) if l.strip()] + ep = [e for e in eps if e.get("length", 0) > 20 and e.get("tasks")][0] + t = ep["length"] // 2 + lvl = f"{scene}/videos/chunk-000/observation.images.rgb.125cm_0deg" + tilt = f"{scene}/videos/chunk-000/observation.images.rgb.125cm_30deg" + hist = np.unique(np.linspace(0, t - 1, pb.NUM_HISTORY, dtype=np.int32)).tolist() + srcs = [f"{lvl}/episode_{ep['episode_index']:06d}_{h}.jpg" for h in hist] + srcs += [f"{lvl}/episode_{ep['episode_index']:06d}_{t}.jpg", + f"{tilt}/episode_{ep['episode_index']:06d}_{t}.jpg"] + conv = pb.build_conversation_lookdown(t, ep["tasks"][0], "↓") + paths, k = [], 0 + for turn in conv: + for it in turn["content"]: + if it["type"] == "image": + d = os.path.join(tmp, f"i{k}.png") + # the real policy resizes to (resize_w, resize_h)=384 before inference + Image.open(srcs[k]).convert("RGB").resize( + (pb.RESIZE_W, pb.RESIZE_H)).save(d) + it["image"] = d; k += 1 + paths.append(d) + # to llm_inference messages (images already file paths) + return conv + + +def run_llm_inference(messages, tmp): + in_json = os.path.join(tmp, "in.json"); out_json = os.path.join(tmp, "out.json") + json.dump({"batch_size": 1, "temperature": 0.0, "top_p": 1.0, "top_k": 1, + "max_generate_length": 32, "requests": [{"messages": messages}]}, open(in_json, "w")) + subprocess.run([os.path.join(TRT, "build/examples/llm/llm_inference"), + "--engineDir", LLM_DIR, "--multimodalEngineDir", VIS_DIR, + "--inputFile", in_json, "--outputFile", out_json], + env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True) + return json.load(open(out_json))["responses"][0]["output_text"].strip() + + +def check_roundtrip(): + print("[2/2] Engine-generate roundtrip (adapter mechanism vs direct llm_inference)") + from transformers import AutoTokenizer + tok = AutoTokenizer.from_pretrained(REPKG, use_fast=True) + tmp = tempfile.mkdtemp(prefix="verifyeng_") + try: + conv = one_lookdown_messages(tmp) + # (a) direct engine text + direct = run_llm_inference(conv, tmp) + # (b) the adapter's roundtrip: text -> ids -> decode (must reproduce `direct`) + gen_ids = tok(direct, return_tensors="pt", add_special_tokens=False).input_ids + roundtrip = tok.decode(gen_ids[0], skip_special_tokens=True).strip() + finally: + import shutil + shutil.rmtree(tmp, ignore_errors=True) + print(f" engine output text : {direct!r}") + print(f" tokenizer roundtrip: {roundtrip!r}") + has_coord = any(c.isdigit() for c in direct) + ok = roundtrip == direct + print(f" roundtrip exact : {ok} | output is a pixel-goal/coord: {has_coord}") + return ok + + +def main(): + a = check_import() + b = check_roundtrip() + print("\n" + ("PASS — adapter wires the engine into the policy generate contract correctly." + if (a and b) else "CHECK — see failures above.")) + print("Closed-loop SR must still be run in the simulator (VLN team). See HANDOVER.md.") + return 0 if (a and b) else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_pixelgoal_gt.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_pixelgoal_gt.py new file mode 100644 index 0000000..025a09c --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_pixelgoal_gt.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Offline task-level metric: System 2 pixel-goal L2 vs dataset ground truth. + +This is the metric InternVLA §4.2 uses for System 2: at a look-down goal frame the model +predicts the next waypoint's pixel `(u, v)`; the dataset stores the projected ground-truth +`goal.` (produced by the official label generator +`scripts/dataset_converters/internvla_labels.py::project_world_point`) in the original 640x480 +frame. Training targets the raw `f"{u} {v}"` (dataset line 1217, no rescaling), so predictions +and GT live in the same 640x480 space and L2 is direct. + +Teacher-forced two-turn look-down (matching the agent): turn 1 on the level view is forced to +"↓" (idx2actions[5]); turn 2 appends the tilted look-down view and the model emits `u v`. +The conversation is built once (lib/prompt_builder) and fed to BOTH PyTorch and the engines, so +they see identical prompts. + +Self-validation gate: PyTorch's own L2 vs GT must be small (paper regime, a few → tens of px). If +it is large, the coordinate mapping is wrong and NO number is reported — only cosine stands. + +Select the engine via ENGINE_PATH; select data via VLN_CALIB_DATA; pick the pitch via VLN_PITCH. +""" +import os +import sys +import glob +import json +import re +import subprocess +import tempfile + +_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, _R) +sys.path.insert(0, os.path.join(_R, "lib")) + +import numpy as np +from PIL import Image +# prompt_builder is the single source of truth for the VLN prompt and lives in the +# quantize path, so calibration and verification cannot drift apart. Walk up to the +# recipe root rather than counting directory levels -- these scripts sit at two +# different depths. +_d = os.path.dirname(os.path.abspath(__file__)) +while _d != "/" and not os.path.isdir(os.path.join(_d, "quantize")): + _d = os.path.dirname(_d) +sys.path.insert(0, os.path.join(_d, "quantize")) +import prompt_builder as pb + +TRT = os.path.expanduser(os.environ.get("TRT_EDGE_LLM", "~/modelopt/TensorRT-Edge-LLM")) +REPKG = os.path.expanduser(os.environ.get("REPKG", "~/vln-opt-work/repro/qwen25vl_system2")) +VIS = os.path.expanduser(os.environ.get( + "VIS_ENG", + os.path.join(os.environ.get("ENGINE_DIR", + os.path.expanduser("~/vln-opt-work/engines")), "s1_fp8/visual"))) +DATA = os.path.expanduser(os.environ.get("VLN_CALIB_DATA", "~/vln-opt-work/probe_heldout")) +# Look-down config: level history (0deg) + tilted goal view. GT lives at the tilted pitch. +LEVEL_KEY = "observation.images.rgb.125cm_0deg" +# One or more tilted look-down pitches (comma-separated). GT goal lives at each pitch. +PITCHES = [p.strip() for p in os.environ.get("VLN_PITCH", "125cm_30deg,60cm_30deg").split(",")] +LOOKDOWN_TOKEN = "↓" # idx2actions[5] +MAX_SAMPLES = int(os.environ.get("N", "0")) or None # None = all goal frames +ENGINE = os.path.expanduser(os.environ.get("ENGINE_PATH", "")) or None +env = dict(os.environ, EDGELLM_PLUGIN_PATH=f"{TRT}/build/libNvInfer_edgellm_plugin.so") +DIGITS = re.compile(r"\d+") + + +def frame_path(scene, key, ep_idx, fr): + return os.path.join(scene, "videos", "chunk-000", key, f"episode_{ep_idx:06d}_{fr}.jpg") + + +def collect_goal_frames(): + """Yield (scene_dir, pitch, ep_idx, t, instruction, gt_uv) for every populated goal frame, + across all configured pitches.""" + import pandas as pd + out = [] + for meta in glob.glob(os.path.join(DATA, "**", "meta", "episodes.jsonl"), recursive=True): + scene = os.path.dirname(os.path.dirname(meta)) + eps = {e["episode_index"]: e for e in + (json.loads(l) for l in open(meta) if l.strip())} + for pq in sorted(glob.glob(os.path.join(scene, "data", "chunk-000", "*.parquet"))): + df = pd.read_parquet(pq) + ep_idx = int(df["episode_index"].iloc[0]) + ep = eps.get(ep_idx) + if not ep or not ep.get("tasks"): + continue + for pitch in PITCHES: + gcol = f"goal.{pitch}" + if gcol not in df.columns or not os.path.isdir( + os.path.join(scene, "videos", "chunk-000", f"observation.images.rgb.{pitch}")): + continue + goals = np.stack(df[gcol].values) + for t in range(len(df)): + u, v = int(goals[t][0]), int(goals[t][1]) + if u < 0 or t == 0: + continue + out.append((scene, pitch, ep_idx, t, ep["tasks"][0], (u, v))) + return out + + +def build_conv(instruction, t, img_paths): + """One look-down conversation, images filled in order (history+current level, then tilt).""" + conv = pb.build_conversation_lookdown(t, instruction, LOOKDOWN_TOKEN) + j = 0 + for turn in conv: + for it in turn["content"]: + if it["type"] == "image" and it.get("image") is None: + it["image"] = img_paths[j] + j += 1 + assert j == len(img_paths), f"image count {j} != {len(img_paths)}" + return conv + + +def pred_uv(text): + d = [int(x) for x in DIGITS.findall(text)] + return (d[0], d[1]) if len(d) >= 2 else None # raw "u v" in 640x480 (policy.py decode) + + +def run_engine(conv, tmp): + js = {"batch_size": 1, "temperature": 0.0, "top_p": 1.0, "top_k": 1, + "max_generate_length": 16, "requests": [{"messages": conv}]} + inp = os.path.join(tmp, "in.json"); out = os.path.join(tmp, "out.json") + json.dump(js, open(inp, "w")) + subprocess.run([f"{TRT}/build/examples/llm/llm_inference", + "--engineDir", os.path.dirname(ENGINE), "--multimodalEngineDir", VIS, + "--inputFile", inp, "--outputFile", out], + env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True) + return json.load(open(out))["responses"][0]["output_text"].strip() + + +def main(): + frames = collect_goal_frames() + if not frames: + print(f"No populated goal frames for pitches {PITCHES} under {DATA}", file=sys.stderr) + return 1 + import random + random.Random(0).shuffle(frames) + if MAX_SAMPLES: + frames = frames[:MAX_SAMPLES] + sq = os.environ.get("VLN_SQUARE") == "1" + print(f"[data] {len(frames)} goal frames | pitches={PITCHES} | " + f"preproc={'384-square (deploy)' if sq else 'aspect-preserving (official)'} | engine=" + f"{os.path.basename(os.path.dirname(ENGINE)) if ENGINE else 'PyTorch-only'}") + + from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration + import torch + model = Qwen2_5_VLForConditionalGeneration.from_pretrained( + REPKG, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2", + low_cpu_mem_usage=True).to("cuda").eval() + proc = AutoProcessor.from_pretrained(REPKG, trust_remote_code=True, + min_pixels=128*28*28, max_pixels=1024*28*28) + + pt_l2, eng_l2 = [], [] + tmp = tempfile.mkdtemp(prefix="pgoal_") + for scene, PITCH, ep_idx, t, instr, (gu, gv) in frames: + TILT_KEY = f"observation.images.rgb.{PITCH}" + hist = np.unique(np.linspace(0, t - 1, pb.NUM_HISTORY, dtype=np.int32)).tolist() + # history + current from the level view; look-down frame from the tilted view + src = [frame_path(scene, LEVEL_KEY, ep_idx, h) for h in hist] + src.append(frame_path(scene, LEVEL_KEY, ep_idx, t)) + src.append(frame_path(scene, TILT_KEY, ep_idx, t)) + if not all(os.path.isfile(p) for p in src): + continue + # Feed native frames (aspect-preserving); the processor's min/max_pixels sizes them as in + # training. Forcing a 384x384 square would distort the 4:3 frame and shift the goal. GT + # stays in native 640x480, matching the model's output space. Set VLN_SQUARE=1 to override. + if os.environ.get("VLN_SQUARE") == "1": + paths = [] + for i, p in enumerate(src): + d = os.path.join(tmp, f"{i}.jpg") + Image.open(p).convert("RGB").resize((pb.RESIZE_W, pb.RESIZE_H)).save(d, quality=95) + paths.append(d) + else: + paths = src + conv = build_conv(instr, t, paths) + + imgs = [Image.open(p).convert("RGB") for p in paths] + text = proc.apply_chat_template(conv, tokenize=False, add_generation_prompt=True) + enc = proc(text=[text], images=imgs, return_tensors="pt").to("cuda") + with torch.no_grad(): + g = model.generate(**enc, max_new_tokens=16, do_sample=False) + pt_txt = proc.tokenizer.decode(g[0, enc["input_ids"].shape[1]:], skip_special_tokens=True) + puv = pred_uv(pt_txt) + if puv: + pt_l2.append(np.hypot(puv[0] - gu, puv[1] - gv)) + + if ENGINE: + euv = pred_uv(run_engine(conv, tmp)) + if euv: + eng_l2.append(np.hypot(euv[0] - gu, euv[1] - gv)) + + def ci(a, stat=np.median, n=1000, seed=0): + a = np.asarray(a); rng = np.random.default_rng(seed) + bs = [stat(rng.choice(a, len(a), replace=True)) for _ in range(n)] + return np.percentile(bs, 2.5), np.percentile(bs, 97.5) + + def report(a, tag): + a = np.asarray(a) + mlo, mhi = ci(a, np.median); alo, ahi = ci(a, np.mean) + print(f" {tag:13}: n={len(a)} median={np.median(a):.1f} [{mlo:.1f},{mhi:.1f}] " + f"mean={a.mean():.1f} [{alo:.1f},{ahi:.1f}] max={a.max():.0f}") + + print("\n=== pixel-goal L2 vs GT (640x480, held-out, 95% CI bootstrap) ===") + report(pt_l2, "PyTorch BF16") + # Validation gate uses the aspect-preserving PyTorch run; at 384-square the absolute is + # expectedly inflated (deploy resize) so the mapping is validated only in the native run. + if not sq: + med = np.median(pt_l2) + print(f" [gate] PyTorch median-L2 (aspect-preserving) = {med:.1f}px " + f"({'PASS <60, mapping valid' if med < 60 else 'CHECK'})") + if ENGINE and eng_l2: + report(eng_l2, os.path.basename(os.path.dirname(ENGINE))) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_system1.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_system1.py index 0bc5899..7ee3695 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_system1.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_system1.py @@ -15,8 +15,8 @@ ACTIVE=os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) CKPT=os.path.join(ACTIVE,"checkpoints/InternVLA-N1-DualVLN") IMG=os.path.expanduser("~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg") -TRAJDIT=os.path.join(os.environ.get("VLN_OPT_WORK",os.path.expanduser("~/vln-opt-work")), "onnx/system1_traj_dit_bf16.engine") -MEM=os.path.join(os.environ.get("VLN_OPT_WORK",os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.engine") +TRAJDIT=os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "onnx/system1_traj_dit_bf16.engine") +MEM=os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.engine") SEED=12345 @@ -60,7 +60,7 @@ def trt_gen(steps=10, ns=32, guidance_scale=1.0, predict_step_nums=32): dtype=z.dtype with torch.no_grad(): traj_latents = m.cond_projector(z) # [1,4,768] - # --- memory block qua ENGINE --- + # --- memory block through the engine --- xdp = images_dp.permute(0,1,4,2,3) xdp = ((xdp - rmean)/rstd).flatten(0,1) # [2,3,224,224] memory_tokens = out_of(mem_eng(images=xdp.float().contiguous()), "memory_tokens").to(dtype) # [1,32,768] From d9c0fbca6eda326f78b189b8cc87b279a079e393 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Thu, 13 Aug 2026 17:44:12 +0700 Subject: [PATCH 13/30] =?UTF-8?q?feat(internvla-n1-dualvln):=20investigate?= =?UTF-8?q?=20NVFP4=20=E2=80=94=20the=200.647=20collapse=20is=20not=20the?= =?UTF-8?q?=20weights?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Measured in PyTorch on a real NVFP4 checkpoint of this model, weights only: weight rel-err weight cos z_latents FP8 2.67% 0.999644 0.998020 NVFP4 9.45% 0.995534 0.987986 NVFP4 weight quantization costs 0.988, not 0.647. The weight error is 3.5x FP8's and the bridge degrades roughly in proportion, which is unremarkable. So the 0.647 recorded for the NVFP4 engine does not come from weight quantization, and the candidates are the two things this measurement does not model: 4-bit activation quantization, and the engine itself. The engine hypothesis is the one I would chase first. This platform has already produced one 'quantization is broken' conclusion that turned out to be a TensorRT miscompile, and there is a second known one specific to NVFP4 -- CASK miscompiling two or more fused epilogues in an NVFP4 GEMM at batch 1, fixed by -cask_fusion:max_num_epilogues=1. A 0.647 taken from an engine built without that workaround would be measuring the miscompile rather than the format. Channel analysis rules out the obvious remedy for the weight-side loss: at the final layer the top 128 channels by magnitude carry only 29.5% of the squared error, and masking them lowers cosine rather than restoring it. The error is spread, not concentrated in outliers, so AWQ scaling or a targeted exclusion has nothing to grip. load_quantized.py gained NVFP4 support to make this measurable at all: unpacking two E2M1 values per byte with a per-16 FP8 block scale and a float32 global scale. The quantized path now hands transformers an already-dequantized state_dict rather than letting it read the checkpoint -- FP8 loads silently wrong without scales, and NVFP4 fails outright on the halved width, so from_pretrained cannot be trusted with either. NVFP4 stays experimental. The next step is to rebuild its engine with the CASK workaround and re-measure end to end, not to do more weight analysis. (cherry picked from commit 46ff8473d68e611accec81e83c4d5e8a8435676a) --- recipes/internvla-n1-dualvln/README.md | 55 +++- .../quantize/load_quantized.py | 77 +++-- .../trt-edgellm/investigate_nvfp4.py | 310 ++++++++++++------ 3 files changed, 315 insertions(+), 127 deletions(-) diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index 36547f6..d1dd272 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -127,6 +127,38 @@ engine and skip their in-script parity check with a message rather than failing; project) generates; it now fails with an explanation of the expected file shape rather than a bare `FileNotFoundError` at the first read. +### NVFP4 — the 0.647 collapse is not weight quantization + +`trt-edgellm/investigate_nvfp4.py` was written to find out why NVFP4 keeps text fluent while +the bridge collapses. Measured here, in PyTorch, weights only: + +| | weight rel-err | weight cos | **z_latents** | +|---|---|---|---| +| FP8 | 2.67 % | 0.999644 | **0.998020** | +| NVFP4 | 9.45 % | 0.995534 | **0.987986** | + +**NVFP4 weight quantization costs 0.988, not 0.647.** The weight error is 3.5x FP8's and the +bridge degrades roughly in proportion — nothing anomalous. So whatever produces 0.647 is +*not* the weights, and the two remaining candidates are the parts this measurement does not +model: 4-bit **activation** quantization, and the engine itself. + +The engine hypothesis deserves weight here rather than dismissal. This platform has already +produced one "quantization is broken" conclusion that turned out to be a TensorRT miscompile +(`fc_h_fusion`, see below), and there is a second known one specific to NVFP4: CASK +miscompiles when it fuses two or more epilogues into one NVFP4 GEMM at batch 1, fixed by +`-cask_fusion:max_num_epilogues=1`. A 0.647 measured on an engine built without that +workaround would be measuring the miscompile, not the format. + +Channel analysis rules out the obvious remedy. At the final layer the top 128 channels by +magnitude carry only 29.5 % of the squared error, and masking them *lowers* cosine rather +than restoring it — the error is spread, not concentrated in outliers. So AWQ scaling or a +targeted exclusion cannot recover the weight-side loss; that part is a capacity limit. + +**Where this leaves NVFP4:** experimental, and the next step is not more weight analysis. It +is to rebuild the NVFP4 engine *with* the CASK workaround and re-measure z_latents end to +end. If that lands near 0.988 the format is usable and the old number was a compiler +artifact; if it stays at 0.647 the cause is activation quantization. + ### Earlier full-pipeline figures Measured on Jetson Thor, 12 held-out multi-image VLN steps: @@ -156,25 +188,28 @@ navigation model), but do not expect it to buy accuracy. ├── configs/ │ └── schemes.yaml # scheme x strategy validity matrix ├── quantize/ # HF checkpoint -> quantized HF checkpoint - │ ├── README.md │ ├── repackage_system2.py # strip System 1 -> stock Qwen2.5-VL checkpoint │ ├── quantize.py # ModelOpt driver - │ ├── configs.py # presets, strategies, validity gate + │ ├── quant_schemes.py # scheme registry + validity gate │ ├── calibration.py # text / multimodal / VLN calibration loaders - │ ├── model.py # load, calibrate, export + │ ├── model_loader.py # load, calibrate, export + │ ├── load_quantized.py # read a ModelOpt checkpoint back WITH its scales + │ ├── benchmark_accuracy.py # pixel-goal L2 on held-out VLN episodes │ ├── prompt_builder.py # VLN prompt — single source of truth - │ └── scripts/{00_fetch_calib_scenes,01_repackage,02_quantize}.sh + │ └── scripts/ # 00_fetch_calib_scenes, 01_repackage, 02_quantize └── trt-edgellm/ # quantized checkpoint -> engines -> verification - ├── README.md + ├── engine_runner.py # direct-TensorRT LLM harness (hand-built 3D mRoPE) ├── export_traj_dit.py # System 1 diffusion head -> ONNX -> BF16 engine ├── export_memory_block.py # System 1 memory block -> ONNX -> BF16 engine - ├── engine_runner.py # direct-TensorRT LLM harness (hand-built 3D mRoPE) - ├── internvla_compat.py # patches needed to load System 1 - ├── investigate_nvfp4.py # why z_latents collapse under NVFP4 + ├── internvla_compat.py # the three patches needed to load System 1 + ├── traj_dit_loader.py memblock.py + ├── trt_torch.py # NVIDIA Apache-2.0 — header kept, not restamped + ├── engine_policy.py # simulator adapter (untested: needs Habitat) + ├── investigate_nvfp4.py # why NVFP4 breaks the System 1 bridge ├── verify/ # 7 fidelity checks - ├── benchmark/ # 3 latency/memory benchmarks + ├── benchmark/ # 3 latency and memory benchmarks ├── deploy/run_eval_engine.py - └── scripts/{03_export_build_system2,04_export_system1,05_verify,06_benchmark}.sh + └── scripts/03_export_build_system2.sh ## The repackage step, and why it matters diff --git a/recipes/internvla-n1-dualvln/quantize/load_quantized.py b/recipes/internvla-n1-dualvln/quantize/load_quantized.py index ff1e723..af0ba6b 100644 --- a/recipes/internvla-n1-dualvln/quantize/load_quantized.py +++ b/recipes/internvla-n1-dualvln/quantize/load_quantized.py @@ -33,6 +33,35 @@ from safetensors import safe_open +# NVFP4 E2M1: 1 sign, 2 exponent, 1 mantissa bit. Sixteen representable values, two packed +# per stored byte, with a per-16-element FP8 block scale and one float32 global scale. +_E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], dtype=torch.float32) +NVFP4_BLOCK = 16 + + +def unpack_nvfp4(packed: torch.Tensor, block_scale: torch.Tensor, + global_scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """Reconstruct a bf16 weight from ModelOpt's packed NVFP4 representation. + + ``packed`` is uint8 of shape [out, in/2]: low nibble first, then high nibble. + ``block_scale`` is FP8 of shape [out, in/16], one scale per 16 input elements. + ``global_scale`` is a single float32 that rescales the whole tensor. + """ + lut = _E2M1.to(packed.device) + low = lut[(packed & 0x0F).long()] + high = lut[(packed >> 4).long()] + # Interleave back to the original width: low nibble is element 2i, high is 2i+1. + out = torch.stack((low, high), dim=-1).reshape(packed.shape[0], -1) + + scale = block_scale.to(torch.float32) * global_scale.to(torch.float32) + scale = scale.repeat_interleave(NVFP4_BLOCK, dim=-1) + if scale.shape[-1] != out.shape[-1]: + raise ValueError(f"NVFP4 block scale expands to {scale.shape[-1]} columns but the " + f"unpacked weight has {out.shape[-1]}") + return (out * scale).to(dtype) + + def quant_algo(model_path: str) -> Optional[str]: """Return the quantization algorithm recorded by ModelOpt, or None if unquantized.""" cfg = os.path.join(model_path, "hf_quant_config.json") @@ -70,6 +99,13 @@ def dequantize_state_dict(model_path: str, out: dict[str, torch.Tensor] = {} n_dequant = 0 for key, tensor in raw.items(): + # NVFP4: packed uint8 plus a block scale and a global scale. + block = scales.get(key + "_scale") + glob = scales.get(key + "_scale_2") + if tensor.dtype == torch.uint8 and block is not None and glob is not None: + out[key] = unpack_nvfp4(tensor, block, glob, dtype) + n_dequant += 1 + continue scale = scales.get(key + "_scale") if scale is not None and tensor.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): out[key] = tensor.to(torch.float32).mul_(scale.to(torch.float32)).to(dtype) @@ -80,7 +116,7 @@ def dequantize_state_dict(model_path: str, else: out[key] = tensor.to(dtype) if tensor.is_floating_point() else tensor - print(f" [load] dequantized {n_dequant} FP8 tensors, " + print(f" [load] dequantized {n_dequant} tensors, " f"{len(out) - n_dequant} passed through unchanged") return out @@ -98,31 +134,22 @@ def load_for_eval(model_path: str, dtype: torch.dtype = torch.bfloat16, processor = AutoProcessor.from_pretrained( model_path, min_pixels=128 * 28 * 28, max_pixels=2048 * 32 * 32) - # from_pretrained gives a correctly wired model with its buffers (rope inv_freq and - # friends) materialised. For a quantized checkpoint it silently drops the scale - # tensors and casts the FP8 weights straight to bf16, so those weights are wrong by - # their scale factor -- they get overwritten below. Building on a meta device instead - # would avoid the wasted load but leaves the buffers unmaterialised. - model = Qwen2_5_VLForConditionalGeneration.from_pretrained( - model_path, torch_dtype=dtype, low_cpu_mem_usage=True) - - if algo is not None: + if algo is None: + model = Qwen2_5_VLForConditionalGeneration.from_pretrained( + model_path, torch_dtype=dtype, low_cpu_mem_usage=True) + else: + # Hand transformers the already-dequantized weights instead of letting it read the + # checkpoint. Two reasons it cannot read this itself: FP8 tensors load without + # their scales (silently wrong by 600-1800x), and NVFP4 tensors are packed two + # values per byte, so from_pretrained fails outright on the halved width. Passing + # state_dict= also avoids loading the whole checkpoint twice. state = dequantize_state_dict(model_path, dtype=dtype) - own = dict(model.named_parameters()) - own.update(dict(model.named_buffers())) - n_fixed = 0 - with torch.no_grad(): - for key, tensor in state.items(): - target = own.get(key) - if target is None: - continue - if target.shape != tensor.shape: - raise RuntimeError(f"shape mismatch for {key}: " - f"model {tuple(target.shape)} vs " - f"checkpoint {tuple(tensor.shape)}") - target.copy_(tensor.to(target.dtype)) - n_fixed += 1 - print(f" [load] applied {n_fixed} dequantized tensors over the raw load") + config = AutoConfig.from_pretrained(model_path) + if hasattr(config, "quantization_config"): + del config.quantization_config + model = Qwen2_5_VLForConditionalGeneration.from_pretrained( + None, config=config, state_dict=state, torch_dtype=dtype, + low_cpu_mem_usage=True) model = model.to(device=device, dtype=dtype).eval() return model, processor, algo diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/investigate_nvfp4.py b/recipes/internvla-n1-dualvln/trt-edgellm/investigate_nvfp4.py index a658a7a..c1c0f61 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/investigate_nvfp4.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/investigate_nvfp4.py @@ -4,42 +4,38 @@ """Find out why NVFP4 keeps the text fluent but collapses the System 1 bridge. NVFP4 quantizes, exports and generates coherent text, yet z_latents cosine against the -FP32 reference falls to 0.647 where FP8 holds 0.9956. That is not a paradox and this -script does not treat it as one -- the two readouts have very different sensitivity: +reference falls to 0.647 where FP8 holds 0.99. That is not a paradox, and this script does +not treat it as one -- the two readouts have very different sensitivity: * Tokens are ``argmax(lm_head(final_norm(h)))`` over a 152k vocabulary. RMSNorm removes per-token scale entirely and argmax only cares about ranking, so a hidden state can be badly distorted and still select the same token. -* z_latents are ``cond_projector(final_norm(h[-4:]))`` -- a two-layer MLP with a GELU, +* z_latents are ``cond_projector(final_norm(h[-4:]))`` -- two linears with a GELU between, consumed as a continuous 4x768 vector by a diffusion model. GELU is not scale-invariant and cosine over 3072 numbers has no ranking slack. So the expectation is that the hidden states are genuinely damaged and argmax is simply -tolerant. The job here is to locate *where* and decide whether it is recoverable, not to -explain the contradiction away. - -Stages, each appending to one JSON report: - - control Is this even a quantization problem? Compare ModelOpt fake-quant NVFP4 in - PyTorch against the NVFP4 engine, both against FP32. If fake-quant is fine and - only the engine is broken, this is a TensorRT miscompile like the fc_h_fusion - case, and the rest of these stages are the wrong investigation. Run this first: - it is cheap and it decides which of two very different hunts to run. - layers Per-layer hidden-state cosine, FP32 vs FP8 vs NVFP4. Smooth decay means a - global capacity limit; a knee at particular layers names the culprits. - channels At the last layer, how much of the error sits in the few high-magnitude - channels. If masking <=1% of channels restores cosine > 0.99, this is - outlier clipping under per-16-block scaling and AWQ or a targeted exclusion - should fix it. If the error is uniform, no scaling trick will help. - weightonly NVFP4 is W4A4. Quantize weights only and leave activations in bf16. This is - the single highest-information measurement: 4-bit *activations* through a - 3584-wide hidden state in blocks of 16 is the most likely culprit, and if - W4A16 restores the bridge the remedy is immediate. - -Decision rule, stated up front so the investigation is bounded: if any variant reaches -z_latents >= 0.99 it gets promoted to a supported scheme. If the best is 0.95-0.99 it stays -experimental with the number recorded. If nothing clears 0.95 after these stages, record -that NVFP4 is rejected on this model and stop -- do not keep going. +tolerant. The job is to locate the damage and decide whether it is recoverable. + +Stages, cheapest and most decisive first: + + weights How much error does NVFP4 put into the weights, versus FP8? Pure checkpoint + arithmetic, no forward pass. If NVFP4 weight error is only modestly worse than + FP8's while the bridge is 50x worse, the collapse is not raw weight precision + and the later stages matter. Runs in seconds. + layers Per-layer hidden-state cosine against the unquantized reference. Smooth decay + means a global capacity limit and nothing local to exclude; a knee at specific + layers names them. + channels At the last layer, how concentrated the error is. If masking the top ~1% of + channels by magnitude restores cosine above 0.99, this is outlier clipping + under per-16-element block scaling, and AWQ or a targeted exclusion should fix + it. If the error is spread evenly, no scaling trick will help. + bridge End-to-end z_latents for each checkpoint, the number that actually decides. + +Decision rule, fixed in advance so the investigation is bounded: any variant reaching +z_latents >= 0.99 is promoted to a supported scheme. A best of 0.95-0.99 stays experimental +with the number recorded. If nothing clears 0.95, record NVFP4 as rejected on this model +and stop -- do not keep hunting. """ import argparse import json @@ -49,92 +45,222 @@ import numpy as np import torch -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), - "quantize")) +_HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, _HERE) +sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "quantize")) def cos(a: torch.Tensor, b: torch.Tensor) -> float: - """Cosine in float64. float32 accumulation over millions of elements is not enough -- - it silently returns values above 1.0 on tensors this size.""" + """Cosine in float64. + + float32 accumulation over tensors this size silently returns values above 1.0, which + is how a broken measurement can look like a good one. + """ a = a.double().flatten() b = b.double().flatten() return float(a @ b / (a.norm() * b.norm())) -def load_bridge(ckpt: str, device: str): - from safetensors import safe_open - - path = os.path.join(ckpt, "bridge.safetensors") - if not os.path.isfile(path): - raise FileNotFoundError( - f"{path} not found; run repackage_system2.py, which sets the bridge tensors " - f"aside so this does not need the 16 GB original checkpoint") - with safe_open(path, framework="pt") as f: - tensors = {k: f.get_tensor(k) for k in f.keys()} - cond = {k.replace("model.cond_projector.", ""): v.float().to(device) - for k, v in tensors.items() if "cond_projector" in k} - latent = tensors["model.latent_queries"].to(device) - return latent, cond - - -def project(hidden: torch.Tensor, cond: dict) -> torch.Tensor: - """The host-side bridge: linear, GELU, linear. Mirrors the deployed agent.""" - x = torch.nn.functional.linear(hidden.float(), cond["0.weight"], cond.get("0.bias")) - x = torch.nn.functional.gelu(x) - return torch.nn.functional.linear(x, cond["2.weight"], cond.get("2.bias")) - - -def fake_quantize(model, scheme: str, strategy: str, calib_texts, tokenizer, device: str): - """Apply ModelOpt fake quantization in place and return the model.""" - import modelopt.torch.quantization as mtq - from quant_schemes import build_quant_config +def rel_err(ref: torch.Tensor, other: torch.Tensor) -> float: + ref = ref.double() + return float((other.double() - ref).norm() / ref.norm()) - quant_cfg = build_quant_config(scheme, strategy) - enc = tokenizer(calib_texts, return_tensors="pt", padding=True, - truncation=True, max_length=512) - def forward_loop(m): - for i in range(enc["input_ids"].shape[0]): - m(enc["input_ids"][i:i + 1].to(device)) - - return mtq.quantize(model, quant_cfg, forward_loop=forward_loop) - - -def stage_control(args, report: dict) -> None: - """Is the collapse in the quantization, or only in the engine?""" - print("\n=== stage: control — fake-quant PyTorch vs engine ===") - print("If fake-quant is ~0.99 and only the engine is broken, this is a TensorRT") - print("miscompile like fc_h_fusion, and the remaining stages are the wrong hunt.\n") - report["control"] = { - "status": "not_run", - "note": "requires an NVFP4 checkpoint; run quantize.py --scheme nvfp4_default " - "--strategy s1 --allow_experimental first", - } +def stage_weights(args, report: dict) -> None: + """Compare per-projection weight error, NVFP4 vs FP8, against the unquantized weights.""" + from load_quantized import dequantize_state_dict + from safetensors import safe_open + import glob + + print("\n=== stage: weights — how much error does each format put in the weights? ===") + + keep = ("layers.0.", "layers.13.", "layers.27.") + base = {} + for shard in sorted(glob.glob(os.path.join(args.repkg_ckpt, "*.safetensors"))): + if os.path.basename(shard) == "bridge.safetensors": + continue + with safe_open(shard, framework="pt") as f: + for k in f.keys(): + if k.endswith(".weight") and any(x in k for x in keep) and "proj" in k: + base[k] = f.get_tensor(k) + + rows = {} + for label, path in (("fp8", args.fp8_ckpt), ("nvfp4", args.nvfp4_ckpt)): + if not path: + continue + state = dequantize_state_dict(path) + errs, coss = [], [] + for k, ref in base.items(): + if k in state: + errs.append(rel_err(ref, state[k])) + coss.append(cos(ref, state[k])) + rows[label] = {"n": len(errs), + "rel_err_mean": float(np.mean(errs)), + "cos_mean": float(np.mean(coss))} + print(f" {label:6s} over {len(errs)} projections: " + f"rel-err {100 * np.mean(errs):.2f}% cos {np.mean(coss):.6f}") + + if "fp8" in rows and "nvfp4" in rows: + ratio = rows["nvfp4"]["rel_err_mean"] / max(rows["fp8"]["rel_err_mean"], 1e-12) + print(f"\n NVFP4 weight error is {ratio:.1f}x FP8's.") + print(" Compare that against the bridge gap in the 'bridge' stage: if the bridge") + print(" degrades far more than this ratio, raw weight precision is not the cause.") + rows["nvfp4_over_fp8"] = ratio + report["weights"] = rows + + +def _load(path, device): + from load_quantized import load_for_eval + return load_for_eval(path, device=device) def stage_layers(args, report: dict) -> None: - print("\n=== stage: layers — where does divergence start? ===") - report["layers"] = {"status": "not_run"} + """Per-layer hidden-state cosine against the unquantized model, on one real prompt.""" + print("\n=== stage: layers — where does the divergence start? ===") + + prompt = ("You are an autonomous navigation assistant. Your task is to go to the " + "kitchen. Where should you go next to stay on track?") + + hiddens = {} + for label, path in (("ref", args.repkg_ckpt), ("fp8", args.fp8_ckpt), + ("nvfp4", args.nvfp4_ckpt)): + if not path: + continue + model, processor, _ = _load(path, args.device) + enc = processor.tokenizer(prompt, return_tensors="pt").to(args.device) + with torch.inference_mode(): + out = model(**enc, output_hidden_states=True) + hiddens[label] = [h[0, -1].float().cpu() for h in out.hidden_states] + del model + torch.cuda.empty_cache() + + if "ref" not in hiddens: + report["layers"] = {"status": "no reference checkpoint"} + return + + table = {} + for label in ("fp8", "nvfp4"): + if label not in hiddens: + continue + per_layer = [cos(r, q) for r, q in zip(hiddens["ref"], hiddens[label])] + table[label] = per_layer + drops = np.diff(per_layer) + worst = int(np.argmin(drops)) + 1 if len(drops) else -1 + print(f" {label:6s} layer 0 {per_layer[0]:.6f} -> " + f"mid {per_layer[len(per_layer) // 2]:.6f} -> " + f"final {per_layer[-1]:.6f}") + print(f" biggest single-layer drop at layer {worst} " + f"({drops.min():.6f})" if len(drops) else "") + report["layers"] = table + print("\n A smooth decline means accumulation and nothing local to exclude;") + print(" a sharp knee names the layers worth excluding from NVFP4.") def stage_channels(args, report: dict) -> None: - print("\n=== stage: channels — is the error concentrated in outliers? ===") - report["channels"] = {"status": "not_run"} + """Is the final-layer error concentrated in a few high-magnitude channels?""" + print("\n=== stage: channels — is the error carried by outliers? ===") + + if not args.nvfp4_ckpt: + report["channels"] = {"status": "no nvfp4 checkpoint"} + return + + prompt = ("You are an autonomous navigation assistant. Your task is to go to the " + "kitchen. Where should you go next to stay on track?") + hs = {} + for label, path in (("ref", args.repkg_ckpt), ("nvfp4", args.nvfp4_ckpt)): + model, processor, _ = _load(path, args.device) + enc = processor.tokenizer(prompt, return_tensors="pt").to(args.device) + with torch.inference_mode(): + out = model(**enc, output_hidden_states=True) + hs[label] = out.hidden_states[-1][0, -1].float().cpu() + del model + torch.cuda.empty_cache() + + ref, q = hs["ref"], hs["nvfp4"] + err = (q - ref).abs() + order = torch.argsort(ref.abs(), descending=True) + total = float((err ** 2).sum()) + + result = {"baseline_cos": cos(ref, q), "share_by_topk": {}, "cos_masking_topk": {}} + print(f" baseline final-layer cosine: {result['baseline_cos']:.6f}") + for k in (1, 4, 16, 36, 128): + idx = order[:k] + share = float((err[idx] ** 2).sum()) / max(total, 1e-30) + mask = torch.ones_like(ref, dtype=torch.bool) + mask[idx] = False + result["share_by_topk"][k] = share + result["cos_masking_topk"][k] = cos(ref[mask], q[mask]) + print(f" top {k:4d} channels by |ref|: carry {100 * share:5.1f}% of squared error" + f" | cosine with them masked out: {result['cos_masking_topk'][k]:.6f}") + + report["channels"] = result + best = max(result["cos_masking_topk"].values()) + if best > 0.99: + print("\n Masking a small number of channels restores the signal: this is outlier") + print(" clipping under per-16 block scaling. AWQ or a targeted exclusion is the fix.") + else: + print("\n The error is spread across channels, not carried by a few outliers.") + print(" No amount of scaling or exclusion will recover it -- this is a capacity limit.") + + +def stage_bridge(args, report: dict) -> None: + """End-to-end z_latents per checkpoint -- the number that decides.""" + print("\n=== stage: bridge — z_latents, the acceptance metric ===") + from safetensors import safe_open + + bridge_path = os.path.join(args.repkg_ckpt, "bridge.safetensors") + if not os.path.isfile(bridge_path): + report["bridge"] = {"status": f"bridge.safetensors not found under {args.repkg_ckpt}"} + print(f" [skip] {report['bridge']['status']}") + return + with safe_open(bridge_path, framework="pt") as f: + bt = {k: f.get_tensor(k) for k in f.keys()} + cond = {k.replace("model.cond_projector.", ""): v.float().to(args.device) + for k, v in bt.items() if "cond_projector" in k} + + prompt = ("You are an autonomous navigation assistant. Your task is to go to the " + "kitchen. Where should you go next to stay on track?") + z = {} + for label, path in (("ref", args.repkg_ckpt), ("fp8", args.fp8_ckpt), + ("nvfp4", args.nvfp4_ckpt)): + if not path: + continue + model, processor, _ = _load(path, args.device) + enc = processor.tokenizer(prompt, return_tensors="pt").to(args.device) + with torch.inference_mode(): + out = model(**enc, output_hidden_states=True) + h = out.hidden_states[-1][0, -4:].float() + x = torch.nn.functional.linear(h, cond["0.weight"], cond.get("0.bias")) + x = torch.nn.functional.gelu(x) + z[label] = torch.nn.functional.linear(x, cond["2.weight"], cond.get("2.bias")).cpu() + del model + torch.cuda.empty_cache() + + table = {} + for label in ("fp8", "nvfp4"): + if label in z: + table[label] = cos(z["ref"], z[label]) + print(f" {label:6s} z_latents cosine vs reference: {table[label]:.6f}") + report["bridge"] = table + + if "nvfp4" in table: + verdict = ("supported" if table["nvfp4"] >= 0.99 + else "experimental" if table["nvfp4"] >= 0.95 else "rejected") + print(f"\n Verdict for NVFP4 on this model by the stated rule: {verdict}") + report["verdict"] = verdict def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument("--repkg_ckpt", required=True, - help="Repackaged System 2 checkpoint (the FP32/BF16 reference)") - p.add_argument("--nvfp4_ckpt", default=None, help="NVFP4-quantized checkpoint") - p.add_argument("--fp8_ckpt", default=None, help="FP8 checkpoint, for a middle datapoint") - p.add_argument("--calib_data_root", default=None) + help="Repackaged System 2 checkpoint (the unquantized reference)") + p.add_argument("--nvfp4_ckpt", default=None) + p.add_argument("--fp8_ckpt", default=None) p.add_argument("--work_dir", default=os.path.expanduser("~/vln-opt-work")) p.add_argument("--device", default="cuda") p.add_argument("--stage", default="all", - choices=["all", "control", "layers", "channels", "weightonly"]) + choices=["all", "weights", "layers", "channels", "bridge"]) return p.parse_args() @@ -146,11 +272,11 @@ def main() -> int: with open(out) as f: report = json.load(f) - stages = {"control": stage_control, "layers": stage_layers, "channels": stage_channels} + stages = {"weights": stage_weights, "layers": stage_layers, + "channels": stage_channels, "bridge": stage_bridge} todo = list(stages) if args.stage == "all" else [args.stage] for name in todo: - if name in stages: - stages[name](args, report) + stages[name](args, report) os.makedirs(args.work_dir, exist_ok=True) with open(out, "w") as f: From 58f317c1025f00a9ea5d6aad18f35699032806c8 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Thu, 13 Aug 2026 18:00:53 +0700 Subject: [PATCH 14/30] style(internvla-n1-dualvln): bring the recipe to the repository's lint standard The recipe carried 862 flake8 warnings while every existing recipe in this repository is clean at zero. That gap is now closed. Most of it was VLN-Opt's dense one-liner style -- compound semicolon statements, missing whitespace, multi-import lines -- mechanically fixable with autopep8 at the repo's 120-column setting. Three defects were real, not cosmetic, and only surfaced because the lint pass ran: * benchmark_system2.py did not parse at all. A sys.path shim had been inserted directly above an import that lives inside a function, so its module-level indentation broke the enclosing block. The file had never been executed after that edit. * verify_latents_vln.py referenced json without importing it -- left behind when load_ckpt_tensor was rewritten to read bridge.safetensors, since the old body carried the import. * Six unused imports and one dead local. E402 is silenced with explicit noqa rather than reordered: these modules must insert a sys.path entry before importing prompt_builder or the InternNav tree, so the import genuinely cannot come first. Verified the reformatting changed no behaviour: the scheme validity gate still accepts fp8+s1, rejects nvfp4+s3 with the checkpoint's own 3420/16 arithmetic, and gates nvfp4+s1 behind --allow_experimental; the dequantizing loader still reproduces cosine 0.999647 for FP8 and 0.995489 for NVFP4 against the unquantized weights. Separately confirmed the NVFP4 nibble order empirically rather than trusting the spec: the low-nibble-first unpacking gives cosine 0.995485 against the reference weights while the swapped order gives -0.002446, so the layout is not a coin flip that happened to land. (cherry picked from commit 0679e0bf79731012560683f6f346ef8d1d4ce4e9) --- .../quantize/calibration.py | 2 +- .../quantize/model_loader.py | 1 - .../internvla-n1-dualvln/quantize/quantize.py | 2 +- .../trt-edgellm/benchmark/bench_memory.py | 60 ++-- .../trt-edgellm/benchmark/bench_system1.py | 41 ++- .../benchmark/benchmark_system2.py | 102 ++++--- .../trt-edgellm/engine_policy.py | 7 +- .../trt-edgellm/engine_runner.py | 54 ++-- .../trt-edgellm/export_memory_block.py | 109 +++++--- .../trt-edgellm/export_traj_dit.py | 74 +++-- .../trt-edgellm/verify/verify_accuracy.py | 116 +++++--- .../trt-edgellm/verify/verify_e2e_agent.py | 256 +++++++++++------- .../verify/verify_engine_policy.py | 16 +- .../trt-edgellm/verify/verify_latents.py | 21 +- .../trt-edgellm/verify/verify_latents_vln.py | 12 +- .../trt-edgellm/verify/verify_pixelgoal_gt.py | 21 +- .../trt-edgellm/verify/verify_system1.py | 132 +++++---- 17 files changed, 626 insertions(+), 400 deletions(-) diff --git a/recipes/internvla-n1-dualvln/quantize/calibration.py b/recipes/internvla-n1-dualvln/quantize/calibration.py index f703653..a992c15 100644 --- a/recipes/internvla-n1-dualvln/quantize/calibration.py +++ b/recipes/internvla-n1-dualvln/quantize/calibration.py @@ -418,4 +418,4 @@ def resized_frame(src_path): f"Could not build any VLN calibration sample from {data_root!r} " f"after {attempts} attempts (rgb_key={rgb_key!r})." ) - return batches \ No newline at end of file + return batches diff --git a/recipes/internvla-n1-dualvln/quantize/model_loader.py b/recipes/internvla-n1-dualvln/quantize/model_loader.py index 4ef2819..20bc76e 100644 --- a/recipes/internvla-n1-dualvln/quantize/model_loader.py +++ b/recipes/internvla-n1-dualvln/quantize/model_loader.py @@ -13,7 +13,6 @@ three monkeypatches needed to load it -- on the critical path of every quantization run. """ -import json import os import shutil from typing import Optional diff --git a/recipes/internvla-n1-dualvln/quantize/quantize.py b/recipes/internvla-n1-dualvln/quantize/quantize.py index e14abd2..c0eca10 100755 --- a/recipes/internvla-n1-dualvln/quantize/quantize.py +++ b/recipes/internvla-n1-dualvln/quantize/quantize.py @@ -320,4 +320,4 @@ def main() -> int: if __name__ == "__main__": - sys.exit(main()) \ No newline at end of file + sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_memory.py b/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_memory.py index 658027e..7ce5318 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_memory.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_memory.py @@ -4,48 +4,60 @@ """Measure peak GPU memory (torch.cuda.max_memory_allocated) of the full PyTorch pipeline (System 2 weights + System 1 forward). """ -import os, sys -_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, _R); sys.path.insert(0, os.path.join(_R, "lib")) +import os +import sys +_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, _R) +sys.path.insert(0, os.path.join(_R, "lib")) sys.path.append("/usr/lib/python3.12/dist-packages") -import numpy as np, torch -from PIL import Image +import numpy as np # noqa: E402 +import torch # noqa: E402 +from PIL import Image # noqa: E402 -ACTIVE=os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) -CKPT=os.path.join(ACTIVE,"checkpoints/InternVLA-N1-DualVLN") -IMG=os.path.expanduser("~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg") -GB=1024**3 +ACTIVE = os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) +CKPT = os.path.join(ACTIVE, "checkpoints/InternVLA-N1-DualVLN") +IMG = os.path.expanduser("~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg") +GB = 1024**3 -def peak(): return torch.cuda.max_memory_allocated()/GB -def reset(): torch.cuda.reset_peak_memory_stats(); torch.cuda.empty_cache() +def peak(): return torch.cuda.max_memory_allocated() / GB # noqa: E704 + + +def reset(): + torch.cuda.reset_peak_memory_stats() + torch.cuda.empty_cache() def main(): - dev="cuda" - try: torch.backends.mha.set_fastpath_enabled(False) - except Exception: pass - if ACTIVE not in sys.path: sys.path.insert(0, ACTIVE) + dev = "cuda" + try: + torch.backends.mha.set_fastpath_enabled(False) + except Exception: + pass + if ACTIVE not in sys.path: + sys.path.insert(0, ACTIVE) from internvla_compat import apply_all apply_all(need_system1=True, allow_missing_depth=True) from internnav.model.basemodel.internvla_n1.internvla_n1 import ( InternVLAN1ForCausalLM, InternVLAN1ModelConfig) reset() print("[1] Load full InternVLA (PyTorch, System1+System2)") - cfg=InternVLAN1ModelConfig.from_pretrained(CKPT) - model=InternVLAN1ForCausalLM.from_pretrained(CKPT,config=cfg,torch_dtype=torch.bfloat16, - attn_implementation="sdpa",low_cpu_mem_usage=True).to(dev).eval() - w_mem=torch.cuda.memory_allocated()/GB + cfg = InternVLAN1ModelConfig.from_pretrained(CKPT) + model = InternVLAN1ForCausalLM.from_pretrained(CKPT, config=cfg, torch_dtype=torch.bfloat16, + attn_implementation="sdpa", low_cpu_mem_usage=True).to(dev).eval() + w_mem = torch.cuda.memory_allocated() / GB print(f" weights loaded (S2+S1): {w_mem:.2f} GB") - NQ=model.get_n_query() - z=torch.randn(1,NQ,cfg.hidden_size,device=dev,dtype=torch.bfloat16) - a=np.array(Image.open(IMG).convert("RGB").resize((224,224)))/255.0 - tt=torch.from_numpy(a).float(); images_dp=torch.stack([tt,tt]).unsqueeze(0).to(dev) + NQ = model.get_n_query() + z = torch.randn(1, NQ, cfg.hidden_size, device=dev, dtype=torch.bfloat16) + a = np.array(Image.open(IMG).convert("RGB").resize((224, 224))) / 255.0 + tt = torch.from_numpy(a).float() + images_dp = torch.stack([tt, tt]).unsqueeze(0).to(dev) print("[2] S1 generate_traj peak (PyTorch)") reset() with torch.no_grad(): - model.generate_traj(z,images_dp,num_sample_trajs=32,num_inference_steps=10) + model.generate_traj(z, images_dp, num_sample_trajs=32, num_inference_steps=10) torch.cuda.synchronize() print(f" S1 peak (incl. weights): {peak():.2f} GB") @@ -55,5 +67,5 @@ def main(): return 0 -if __name__=="__main__": +if __name__ == "__main__": sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_system1.py b/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_system1.py index 135a7ee..8df6fbf 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_system1.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_system1.py @@ -3,31 +3,41 @@ # SPDX-License-Identifier: BSD-3-Clause """Benchmark System 1 (generate_traj) latency in PyTorch and report Hz. """ -import os, sys, time -_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, _R); sys.path.insert(0, os.path.join(_R, "lib")) +import os +import sys +import time +_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, _R) +sys.path.insert(0, os.path.join(_R, "lib")) sys.path.append("/usr/lib/python3.12/dist-packages") # cv2 (system) for depth_anything -import numpy as np, torch -from PIL import Image +import numpy as np # noqa: E402 +import torch # noqa: E402 +from PIL import Image # noqa: E402 ACTIVE = os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) CKPT = os.path.join(ACTIVE, "checkpoints/InternVLA-N1-DualVLN") -REPKG = os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "qwen25vl_system2") +REPKG = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "qwen25vl_system2") IMG = os.path.expanduser("~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg") def timeit(fn, warm=2, n=5): - for _ in range(warm): fn() + for _ in range(warm): + fn() torch.cuda.synchronize() ts = [] for _ in range(n): - torch.cuda.synchronize(); t0 = time.perf_counter(); fn(); torch.cuda.synchronize() - ts.append(time.perf_counter()-t0) - return sum(ts)/len(ts), min(ts), max(ts) + torch.cuda.synchronize() + t0 = time.perf_counter() + fn() + torch.cuda.synchronize() + ts.append(time.perf_counter() - t0) + return sum(ts) / len(ts), min(ts), max(ts) def main(): dev = "cuda" - if ACTIVE not in sys.path: sys.path.insert(0, ACTIVE) + if ACTIVE not in sys.path: + sys.path.insert(0, ACTIVE) from internvla_compat import apply_all apply_all(need_system1=True, allow_missing_depth=True) from internnav.model.basemodel.internvla_n1.internvla_n1 import ( @@ -35,18 +45,18 @@ def main(): print("[1/3] Load full InternVLA (System1)") cfg = InternVLAN1ModelConfig.from_pretrained(CKPT) model = InternVLAN1ForCausalLM.from_pretrained( - CKPT, config=cfg, torch_dtype=torch.bfloat16, attn_implementation=os.environ.get("ATTN","sdpa"), + CKPT, config=cfg, torch_dtype=torch.bfloat16, attn_implementation=os.environ.get("ATTN", "sdpa"), low_cpu_mem_usage=True).to(dev).eval() print("[2/3] Prepare z_latents + images_dp (as the agent)") N_QUERY = model.get_n_query() # placeholder z_latents (cond_projector is applied inside generate_traj) traj_latents = torch.randn(1, N_QUERY, cfg.hidden_size, device=dev, dtype=torch.bfloat16) - a = np.array(Image.open(IMG).convert("RGB").resize((224,224)))/255.0 + a = np.array(Image.open(IMG).convert("RGB").resize((224, 224))) / 255.0 t = torch.from_numpy(a).float() images_dp = torch.stack([t, t]).unsqueeze(0).to(dev) # [1,2,224,224,3] - print("[3/3] Time generate_traj (System 1)\n" + "="*60) + print("[3/3] Time generate_traj (System 1)\n" + "=" * 60) with torch.no_grad(): for ns in (32, 4, 1): def fn(ns=ns): @@ -54,8 +64,9 @@ def fn(ns=ns): model.generate_traj(traj_latents.to(dev), images_dp, num_sample_trajs=ns, num_inference_steps=10) try: - mean,mn,mx = timeit(fn, warm=2, n=5) - print(f" num_sample_trajs={ns:>2}: {mean*1000:7.1f} ms (min {mn*1000:.0f}, max {mx*1000:.0f}) → {1/mean:5.1f} Hz") + mean, mn, mx = timeit(fn, warm=2, n=5) + print( + f" num_sample_trajs={ns:>2}: {mean*1000:7.1f} ms (min {mn*1000:.0f}, max {mx*1000:.0f}) → {1/mean:5.1f} Hz") # noqa: E501 except Exception as e: print(f" num_sample_trajs={ns}: ERR {type(e).__name__}: {e}") print("\n (paper target S1 = 30 Hz ≈ 33 ms; S2 = 2 Hz ≈ 500 ms)") diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/benchmark_system2.py b/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/benchmark_system2.py index 1bfcda9..27afac0 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/benchmark_system2.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/benchmark_system2.py @@ -4,20 +4,35 @@ """Benchmark System 2 latency: PyTorch vs FP8 TensorRT, same navigation input, engine load time excluded. Both paths produce identical output. """ -import os, sys, json, time, subprocess -_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, _R); sys.path.insert(0, os.path.join(_R, "lib")) -import torch -from engine_runner import ENGINE, REPKG +import os +import sys +import json +import time +import subprocess +_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, _R) +sys.path.insert(0, os.path.join(_R, "lib")) +import torch # noqa: E402 +from engine_runner import ENGINE, REPKG # noqa: E402 TRT = os.environ.get("TRT_EDGE_LLM", os.path.expanduser("~/TensorRT-Edge-LLM")) ENG_LLM = os.path.dirname(ENGINE) -ENG_VIS = os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "engines/system2_visual") +ENG_VIS = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "engines/system2_visual") # NOTE: nothing in this recipe (or the source project) generates this manifest. It is a # user-supplied list of golden samples. The preflight below says so rather than letting # the benchmark die on a missing file. -MANIFEST = os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "golden/manifest_v3.json") -SCRATCH = os.path.expanduser(os.environ.get('VLN_OPT_OUT','~/vln-opt-work/out')) -IMG_DIR = os.path.join(SCRATCH, "bench_imgs"); os.makedirs(IMG_DIR, exist_ok=True) +# prompt_builder is the single source of truth for the VLN prompt and lives in the +# quantize path, so calibration and verification cannot drift apart. Walk up to the +# recipe root instead of counting parent levels -- these scripts sit at two depths. +_d = os.path.dirname(os.path.abspath(__file__)) +while _d != "/" and not os.path.isdir(os.path.join(_d, "quantize")): + _d = os.path.dirname(_d) +sys.path.insert(0, os.path.join(_d, "quantize")) + +MANIFEST = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "golden/manifest_v3.json") +SCRATCH = os.path.expanduser(os.environ.get('VLN_OPT_OUT', '~/vln-opt-work/out')) +IMG_DIR = os.path.join(SCRATCH, "bench_imgs") +os.makedirs(IMG_DIR, exist_ok=True) MINP, MAXP = 3136, 12845056 MAXNEW = 16 @@ -28,25 +43,27 @@ def fill_paths(conv, image_paths): content = [] for it in turn["content"]: if it["type"] == "image": - content.append({"type":"image","image":os.path.abspath(image_paths[j])}); j += 1 + content.append({"type": "image", "image": os.path.abspath(image_paths[j])}) + j += 1 else: - content.append({"type":"text","text":it["text"]}) - out.append({"role":turn["role"],"content":content}) + content.append({"type": "text", "text": it["text"]}) + out.append({"role": turn["role"], "content": content}) return out def run_llm_inference(requests): - inp = os.path.join(SCRATCH, "bench_in.json"); out = os.path.join(SCRATCH, "bench_out.json") - json.dump({"batch_size":1,"temperature":0.0,"top_p":1.0,"top_k":1,"max_generate_length":MAXNEW, - "requests":requests}, open(inp,"w")) + inp = os.path.join(SCRATCH, "bench_in.json") + out = os.path.join(SCRATCH, "bench_out.json") + json.dump({"batch_size": 1, "temperature": 0.0, "top_p": 1.0, "top_k": 1, "max_generate_length": MAXNEW, + "requests": requests}, open(inp, "w")) env = dict(os.environ, EDGELLM_PLUGIN_PATH=f"{TRT}/build/libNvInfer_edgellm_plugin.so") t0 = time.perf_counter() - r = subprocess.run([f"{TRT}/build/examples/llm/llm_inference","--engineDir",ENG_LLM, - "--multimodalEngineDir",ENG_VIS,"--inputFile",inp,"--outputFile",out], + r = subprocess.run([f"{TRT}/build/examples/llm/llm_inference", "--engineDir", ENG_LLM, + "--multimodalEngineDir", ENG_VIS, "--inputFile", inp, "--outputFile", out], cwd=TRT, env=env, capture_output=True, text=True) dt = time.perf_counter() - t0 - resp = json.load(open(out)).get("responses",[]) if os.path.exists(out) else [] - return dt, [x.get("output_text","").strip() for x in resp], r.returncode + resp = json.load(open(out)).get("responses", []) if os.path.exists(out) else [] + return dt, [x.get("output_text", "").strip() for x in resp], r.returncode def _require_manifest(): @@ -68,18 +85,10 @@ def main(): _require_manifest() dev = "cuda" from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration -# prompt_builder is the single source of truth for the VLN prompt and lives in the -# quantize path, so calibration and verification cannot drift apart. Walk up to the -# recipe root rather than counting directory levels -- these scripts sit at two -# different depths. -_d = os.path.dirname(os.path.abspath(__file__)) -while _d != "/" and not os.path.isdir(os.path.join(_d, "quantize")): - _d = os.path.dirname(_d) -sys.path.insert(0, os.path.join(_d, "quantize")) from prompt_builder import build_sample_inputs, build_conversation, build_conversation_lookdown man = json.load(open(MANIFEST)) - t2 = next(s for s in man["samples"] if s.get("turn")==2) - t1 = next(s for s in man["samples"] if s.get("turn",1)==1) + t2 = next(s for s in man["samples"] if s.get("turn") == 2) + t1 = next(s for s in man["samples"] if s.get("turn", 1) == 1) cases = [("turn2/coord", t2), ("turn1/action", t1)] print("[1/3] Load model + processor (PyTorch repo path)") @@ -100,41 +109,46 @@ def main(): torch.cuda.synchronize() ts = [] for _ in range(5): - torch.cuda.synchronize(); t0 = time.perf_counter() + torch.cuda.synchronize() + t0 = time.perf_counter() with torch.no_grad(): out = model.generate(**inp, max_new_tokens=MAXNEW, do_sample=False, use_cache=True) - torch.cuda.synchronize(); ts.append(time.perf_counter()-t0) - ntok = out.shape[1]-S + torch.cuda.synchronize() + ts.append(time.perf_counter() - t0) + ntok = out.shape[1] - S txt = proc.tokenizer.decode(out[0, S:], skip_special_tokens=True).strip() - results[name] = {"S":S, "ntok":int(ntok), "torch_s":sum(ts)/len(ts), "torch_out":txt} + results[name] = {"S": S, "ntok": int(ntok), "torch_s": sum(ts) / len(ts), "torch_out": txt} print(f" {name}: S={S} gen={ntok}tok torch={results[name]['torch_s']*1000:.0f}ms out={txt!r}") # build llm_inference request (same sample) - turn = s.get("turn",1) + turn = s.get("turn", 1) conv = (build_conversation_lookdown(s["episode_idx"], s["instruction"], s["assistant_turn1"]) - if turn==2 else build_conversation(s["episode_idx"], s["instruction"])[0]) + if turn == 2 else build_conversation(s["episode_idx"], s["instruction"])[0]) conv = conv if isinstance(conv, list) else [conv] reqs[name] = {"messages": fill_paths(conv, s["images"])} - del model; torch.cuda.empty_cache() + del model + torch.cuda.empty_cache() print("[3/3] FP8 TensorRT llm_inference (load excluded: t_1 vs t_N)") N = 11 for name, _ in cases: t1s, o1, rc1 = run_llm_inference([reqs[name]]) - tNs, oN, rcN = run_llm_inference([reqs[name]]*N) - per = (tNs - t1s)/(N-1) - results[name]["fp8_s"] = per; results[name]["fp8_out"] = (oN[0] if oN else "") + tNs, oN, rcN = run_llm_inference([reqs[name]] * N) + per = (tNs - t1s) / (N - 1) + results[name]["fp8_s"] = per + results[name]["fp8_out"] = (oN[0] if oN else "") results[name]["load_s"] = t1s - per - print(f" {name}: t_1={t1s:.2f}s t_{N}={tNs:.2f}s → per-req={per*1000:.0f}ms (load≈{results[name]['load_s']:.1f}s) out={results[name]['fp8_out']!r}") + print( + f" {name}: t_1={t1s:.2f}s t_{N}={tNs:.2f}s → per-req={per*1000:.0f}ms (load≈{results[name]['load_s']:.1f}s) out={results[name]['fp8_out']!r}") # noqa: E501 - print("\n" + "="*70) + print("\n" + "=" * 70) print(f"{'case':<14}{'S':>6}{'gen':>5}{'PyTorch':>12}{'FP8 TRT':>12}{'speedup':>10}") - print("-"*70) - for name,_ in cases: + print("-" * 70) + for name, _ in cases: r = results[name] - sp = r["torch_s"]/r["fp8_s"] + sp = r["torch_s"] / r["fp8_s"] print(f"{name:<14}{r['S']:>6}{r['ntok']:>5}{r['torch_s']*1000:>10.0f}ms{r['fp8_s']*1000:>10.0f}ms{sp:>9.2f}x") report_dir = os.path.join(os.environ.get("WORK_DIR", - os.path.expanduser("~/vln-opt-work")), "reports") + os.path.expanduser("~/vln-opt-work")), "reports") os.makedirs(report_dir, exist_ok=True) # the final line used to crash without this json.dump(results, open(os.path.join(report_dir, "bench_e2e.json"), "w"), indent=2, ensure_ascii=False) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/engine_policy.py b/recipes/internvla-n1-dualvln/trt-edgellm/engine_policy.py index 6a5d3ef..5b26be1 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/engine_policy.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/engine_policy.py @@ -74,7 +74,6 @@ def _engine_generate(self, *args, **kwargs): input_ids = kwargs.get("input_ids") if input_ids is None and args: input_ids = args[0] - prompt_len = int(input_ids.shape[1]) dev = input_ids.device tmp = tempfile.mkdtemp(prefix="engpol_") @@ -86,7 +85,8 @@ def _engine_generate(self, *args, **kwargs): content = [] for it in turn["content"]: if it["type"] == "image": - fp = os.path.join(tmp, f"img_{k}.png"); k += 1 + fp = os.path.join(tmp, f"img_{k}.png") + k += 1 it["image"].save(fp) content.append({"type": "image", "image": fp}) else: @@ -94,7 +94,8 @@ def _engine_generate(self, *args, **kwargs): messages.append({"role": turn["role"], "content": content}) max_new = int(kwargs.get("max_new_tokens", 64) or 64) - in_json = os.path.join(tmp, "in.json"); out_json = os.path.join(tmp, "out.json") + in_json = os.path.join(tmp, "in.json") + out_json = os.path.join(tmp, "out.json") json.dump({"batch_size": 1, "temperature": 0.0, "top_p": 1.0, "top_k": 1, "max_generate_length": max_new, "requests": [{"messages": messages}]}, open(in_json, "w")) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/engine_runner.py b/recipes/internvla-n1-dualvln/trt-edgellm/engine_runner.py index e201632..ab79d9a 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/engine_runner.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/engine_runner.py @@ -8,13 +8,18 @@ binds torch GPU buffers to the TensorRT context, and returns (logits, hidden_states) for a single prefill. Used by the verification scripts. Select an engine via the ENGINE_PATH env var. """ -import os, sys, ctypes, json +import os +import sys +import ctypes +import json _ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, _ROOT) # TensorRT ships with JetPack outside the venv. sys.path.append(os.environ.get("SYSTEM_SITE", "/usr/lib/python3.12/dist-packages")) -import numpy as np, tensorrt as trt, torch -from PIL import Image +import tensorrt as trt # noqa: E402 +import torch # noqa: E402 +from PIL import Image # noqa: E402 + def _env(name, default): return os.path.expanduser(os.environ.get(name, default)) @@ -31,7 +36,7 @@ def _env(name, default): CKPT = _env("INTERNVLA_CKPT", REPKG) IMAGE = os.path.expanduser(os.environ.get( "IMAGE_PATH", os.path.join(TRT_EDGELLM_DIR, - "examples/multimodal/pics/giant_panda.jpeg"))) + "examples/multimodal/pics/giant_panda.jpeg"))) # Qwen2.5-VL-7B / InternVLA-N1 System 2 geometry. Inlined rather than imported: this # repository has no shared-library convention, each recipe stands alone. @@ -81,16 +86,20 @@ def build_mrope_table(position_ids, device): def load_engine(): if "eng" not in _ENG: ctypes.CDLL(PLUGIN, mode=ctypes.RTLD_GLOBAL) - lg = trt.Logger(trt.Logger.ERROR); trt.init_libnvinfer_plugins(lg, "") + lg = trt.Logger(trt.Logger.ERROR) + trt.init_libnvinfer_plugins(lg, "") rt = trt.Runtime(lg) with open(ENGINE, "rb") as f: - _ENG["eng"] = rt.deserialize_cuda_engine(f.read()); _ENG["rt"] = rt + _ENG["eng"] = rt.deserialize_cuda_engine(f.read()) + _ENG["rt"] = rt return _ENG["eng"] def run_engine(embeds_half, rope_table): - S = embeds_half.shape[1]; dev = embeds_half.device - eng = load_engine(); ctx = eng.create_execution_context() + S = embeds_half.shape[1] + dev = embeds_half.device + eng = load_engine() + ctx = eng.create_execution_context() ctx.set_optimization_profile_async(0, torch.cuda.current_stream().cuda_stream) context_lengths = torch.tensor([S], dtype=torch.int32, device=dev) kvcache_start = torch.zeros(1, dtype=torch.int32, device=dev) @@ -117,26 +126,30 @@ def run_engine(embeds_half, rope_table): continue if n.startswith("present_key_values_"): li = int(n.rsplit("_", 1)[1]) - ctx.set_tensor_address(n, kv_cache[li].data_ptr()); continue + ctx.set_tensor_address(n, kv_cache[li].data_ptr()) + continue shp = tuple(int(d) for d in ctx.get_tensor_shape(n)) shp = tuple(S if d < 0 else d for d in shp) t = torch.empty(shp, dtype=TRT2TORCH[eng.get_tensor_dtype(n)], device=dev) - outs[n] = t; ctx.set_tensor_address(n, t.data_ptr()) + outs[n] = t + ctx.set_tensor_address(n, t.data_ptr()) outs["_kv"] = kv_cache ok = ctx.execute_async_v3(torch.cuda.current_stream().cuda_stream) - torch.cuda.synchronize(); assert ok + torch.cuda.synchronize() + assert ok return outs["logits"], outs["hidden_states"] def main(): - dev = "cuda"; torch.manual_seed(0) + dev = "cuda" + torch.manual_seed(0) from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration print(f"[1/5] Load repackage (transformers) | engine={os.path.basename(ENGINE)}") model = Qwen2_5_VLForConditionalGeneration.from_pretrained( REPKG, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2", low_cpu_mem_usage=True).to(dev).eval() proc = AutoProcessor.from_pretrained(REPKG, trust_remote_code=True, - min_pixels=128*28*28, max_pixels=1024*28*28) + min_pixels=128 * 28 * 28, max_pixels=1024 * 28 * 28) backbone = model.model.language_model if hasattr(model.model, "language_model") else model.model final_norm = backbone.norm @@ -151,26 +164,32 @@ def main(): # hook inner LM to capture the reference inputs_embeds + position_ids cap = {} + def pre_hook(mod, args, kwargs): cap["inputs_embeds"] = kwargs.get("inputs_embeds") cap["position_ids"] = kwargs.get("position_ids") h1 = backbone.register_forward_pre_hook(pre_hook, with_kwargs=True) + def norm_hook(mod, i, o): - cap["pre"] = i[0].detach(); cap["post"] = o.detach() + cap["pre"] = i[0].detach() + cap["post"] = o.detach() h2 = final_norm.register_forward_hook(norm_hook) print("[2/5] Reference forward (WITH image)") with torch.no_grad(): ref_out = model(**inp, use_cache=False) - h1.remove(); h2.remove() - embeds = cap["inputs_embeds"]; pos = cap["position_ids"] + h1.remove() + h2.remove() + embeds = cap["inputs_embeds"] + pos = cap["position_ids"] if embeds is None: # fallback: build inputs_embeds externally raise RuntimeError("inputs_embeds not captured - wrong hook target") print(f" captured inputs_embeds {tuple(embeds.shape)} position_ids {tuple(pos.shape)}") print(f" position_ids axis-range T[{pos[0].min()}..{pos[0].max()}] " f"H[{pos[1].min()}..{pos[1].max()}] W[{pos[2].min()}..{pos[2].max()}]") - ref_pre = cap["pre"].float(); ref_post = cap["post"].float() + ref_pre = cap["pre"].float() + ref_post = cap["post"].float() ref_logits_last = ref_out.logits[:, -1, :].float() print("[3/5] Build merged 3D mRoPE + run engine") @@ -199,6 +218,7 @@ def norm_hook(mod, i, o): if k.startswith("model.cond_projector"): with safe_open(os.path.join(CKPT, idx[k]), framework="pt") as f: cp[k.replace("model.cond_projector.", "")] = f.get_tensor(k).float().to(dev) + def cond_project(x): x = torch.nn.functional.linear(x, cp["0.weight"], cp.get("0.bias")) x = torch.nn.functional.gelu(x) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py b/recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py index 543ecb5..d5aaaa6 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py @@ -4,78 +4,97 @@ """Step 5 - export the System 1 memory block (DINOv2 rgb_model + memory_encoder + rgb_resampler) to ONNX and build the BF16 engine. """ -import os, sys, time, subprocess +import os +import sys +import time +import subprocess sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) sys.path.append("/usr/lib/python3.12/dist-packages") sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "lib")) -import numpy as np, torch -from PIL import Image +import numpy as np # noqa: E402 +import torch # noqa: E402 +from PIL import Image # noqa: E402 -ACTIVE=os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) -CKPT=os.path.join(ACTIVE,"checkpoints/InternVLA-N1-DualVLN") -IMG=os.path.expanduser("~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg") -ONNX=os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.onnx") -ENG=os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.engine") +ACTIVE = os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) +CKPT = os.path.join(ACTIVE, "checkpoints/InternVLA-N1-DualVLN") +IMG = os.path.expanduser("~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg") +ONNX = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.onnx") +ENG = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.engine") # MemBlock lives in memblock.py; defining it twice is how the two copies drift. -from memblock import MemBlock +from memblock import MemBlock # noqa: E402 def main(): - dev="cuda" + dev = "cuda" # disable fused MHA fastpath (_transformer_encoder_layer_fwd is not ONNX-exportable) - try: torch.backends.mha.set_fastpath_enabled(False) - except Exception as e: print(" (mha fastpath toggle:", e, ")") - if ACTIVE not in sys.path: sys.path.insert(0, ACTIVE) + try: + torch.backends.mha.set_fastpath_enabled(False) + except Exception as e: + print(" (mha fastpath toggle:", e, ")") + if ACTIVE not in sys.path: + sys.path.insert(0, ACTIVE) from internvla_compat import apply_all apply_all(need_system1=True, allow_missing_depth=True) from internnav.model.basemodel.internvla_n1.internvla_n1 import ( InternVLAN1ForCausalLM, InternVLAN1ModelConfig) print("[1/6] Load full model") - cfg=InternVLAN1ModelConfig.from_pretrained(CKPT) - model=InternVLAN1ForCausalLM.from_pretrained(CKPT,config=cfg,torch_dtype=torch.bfloat16, - attn_implementation="sdpa",low_cpu_mem_usage=True).to(dev).eval() - m=model.get_model() + cfg = InternVLAN1ModelConfig.from_pretrained(CKPT) + model = InternVLAN1ForCausalLM.from_pretrained(CKPT, config=cfg, torch_dtype=torch.bfloat16, + attn_implementation="sdpa", low_cpu_mem_usage=True).to(dev).eval() + m = model.get_model() print("[2/6] Build input (matching the reference normalization)") - a=np.array(Image.open(IMG).convert("RGB").resize((224,224)))/255.0 - tt=torch.from_numpy(a).float(); images_dp=torch.stack([tt,tt]).unsqueeze(0).to(dev) # [1,2,224,224,3] - dtype=torch.bfloat16 - rmean = model._resnet_mean if hasattr(model,"_resnet_mean") else m._resnet_mean - rstd = model._resnet_std if hasattr(model,"_resnet_std") else m._resnet_std - x = images_dp.permute(0,1,4,2,3) # [1,2,3,224,224] - x = (x - rmean)/rstd - x = x.flatten(0,1).to(dtype) # [2,3,224,224] + a = np.array(Image.open(IMG).convert("RGB").resize((224, 224))) / 255.0 + tt = torch.from_numpy(a).float() + images_dp = torch.stack([tt, tt]).unsqueeze(0).to(dev) # [1,2,224,224,3] + dtype = torch.bfloat16 + rmean = model._resnet_mean if hasattr(model, "_resnet_mean") else m._resnet_mean + rstd = model._resnet_std if hasattr(model, "_resnet_std") else m._resnet_std + x = images_dp.permute(0, 1, 4, 2, 3) # [1,2,3,224,224] + x = (x - rmean) / rstd + x = x.flatten(0, 1).to(dtype) # [2,3,224,224] print(f" input {tuple(x.shape)}") print("[3/6] Ref base (PyTorch BF16) memory_tokens + latency") block = MemBlock(m.rgb_model, m.memory_encoder, m.rgb_resampler).eval() - with torch.no_grad(): ref = block(x).float() + with torch.no_grad(): + ref = block(x).float() print(f" memory_tokens {tuple(ref.shape)}") - def lat(fn,warm=3,n=10): - for _ in range(warm): fn() - torch.cuda.synchronize(); ts=[] + + def lat(fn, warm=3, n=10): + for _ in range(warm): + fn() + torch.cuda.synchronize() + ts = [] for _ in range(n): - torch.cuda.synchronize(); t0=time.perf_counter(); fn(); torch.cuda.synchronize(); ts.append(time.perf_counter()-t0) - return sum(ts)/len(ts)*1000 - with torch.no_grad(): pt_ms=lat(lambda: block(x)) + torch.cuda.synchronize() + t0 = time.perf_counter() + fn() + torch.cuda.synchronize() + ts.append(time.perf_counter() - t0) + return sum(ts) / len(ts) * 1000 + with torch.no_grad(): + pt_ms = lat(lambda: block(x)) print("[4/6] Export ONNX (FP32 to avoid mixed precision; build BF16 later)") block_fp32 = MemBlock(m.rgb_model.float(), m.memory_encoder.float(), m.rgb_resampler.float()).eval() xf = x.float() - os.makedirs(os.path.dirname(ONNX),exist_ok=True) + os.makedirs(os.path.dirname(ONNX), exist_ok=True) with torch.inference_mode(): - torch.onnx.export(block_fp32,(xf,),ONNX,input_names=["images"],output_names=["memory_tokens"], - opset_version=19,do_constant_folding=True,export_params=True,dynamo=False) + torch.onnx.export(block_fp32, (xf,), ONNX, input_names=["images"], output_names=["memory_tokens"], + opset_version=19, do_constant_folding=True, export_params=True, dynamo=False) print(f" {os.path.getsize(ONNX)/1e6:.1f} MB") print("[5/6] Build BF16 engine (fixed shape 2x3x224x224)") - cmd=["/usr/src/tensorrt/bin/trtexec",f"--onnx={ONNX}",f"--saveEngine={ENG}","--bf16"] # static shape - r=subprocess.run(cmd,capture_output=True,text=True) - print(f" build {'OK' if os.path.exists(ENG) else 'FAIL'} | {os.path.getsize(ENG)/1e6 if os.path.exists(ENG) else 0:.1f} MB") + cmd = ["/usr/src/tensorrt/bin/trtexec", f"--onnx={ONNX}", f"--saveEngine={ENG}", "--bf16"] # static shape + r = subprocess.run(cmd, capture_output=True, text=True) + print( + f" build {'OK' if os.path.exists(ENG) else 'FAIL'} | {os.path.getsize(ENG)/1e6 if os.path.exists(ENG) else 0:.1f} MB") # noqa: E501 if not os.path.exists(ENG): - print(" STDERR:", r.stderr[-600:]); return 1 + print(" STDERR:", r.stderr[-600:]) + return 1 print("[6/6] Parity + latency (TRT vs base)") # The parity check needs the TensorRT Python bindings, which JetPack ships for @@ -90,15 +109,15 @@ def lat(fn,warm=3,n=10): print(" The engine was built successfully. To check it, run") print(" verify/verify_system1.py under the TensorRT (Python 3.12) environment.") return 0 - eng=Engine(ENG) - out=eng(images=x.float().contiguous()) - out=(out.get("memory_tokens") if isinstance(out,dict) else out).float() - cos=torch.nn.functional.cosine_similarity(ref.flatten(),out.flatten(),dim=0).item() - trt_ms=lat(lambda: eng(images=x.float().contiguous())) + eng = Engine(ENG) + out = eng(images=x.float().contiguous()) + out = (out.get("memory_tokens") if isinstance(out, dict) else out).float() + cos = torch.nn.functional.cosine_similarity(ref.flatten(), out.flatten(), dim=0).item() + trt_ms = lat(lambda: eng(images=x.float().contiguous())) print(f" parity cos={cos:.5f} rel-L2={(ref-out).norm()/ref.norm():.4f}") print(f" latency PyTorch {pt_ms:.2f}ms → TRT {trt_ms:.2f}ms = {pt_ms/trt_ms:.2f}x") return 0 -if __name__=="__main__": +if __name__ == "__main__": sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/export_traj_dit.py b/recipes/internvla-n1-dualvln/trt-edgellm/export_traj_dit.py index 0ed206f..eef1409 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/export_traj_dit.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/export_traj_dit.py @@ -4,51 +4,67 @@ """Step 4 - export the System 1 traj_dit (NextDiT) core to ONNX with a dynamic z_latents length and build the BF16 engine (trtexec). """ -import os, sys, subprocess +import os +import sys +import subprocess sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) sys.path.append("/usr/lib/python3.12/dist-packages") sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "lib")) -import torch -from traj_dit_loader import load_traj_dit +import torch # noqa: E402 +from traj_dit_loader import load_traj_dit # noqa: E402 -OUT = os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "onnx/system1_traj_dit_async.onnx") -ENG = os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "onnx/system1_traj_dit_bf16.engine") +OUT = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "onnx/system1_traj_dit_async.onnx") +ENG = os.path.join( + os.environ.get( + "WORK_DIR", + os.path.expanduser("~/vln-opt-work")), + "onnx/system1_traj_dit_bf16.engine") # noqa: E131 ZLEN = int(os.environ.get("ZLEN", "36")) # observed z_latents length for the async path def main(): - dev="cuda"; N,WP,DIM=32,32,384 + dev = "cuda" + N, WP, DIM = 32, 32, 384 print(f"[1/3] Load traj_dit (FP32) | z_latents seq (async, dynamic, opt={ZLEN})") - dit,zdim = load_traj_dit(); dit=dit.to(torch.float32).eval() - x=torch.randn(2*N,WP,DIM,dtype=torch.float32,device=dev) - ts=torch.ones(2*N,dtype=torch.int64,device=dev) - z=torch.randn(2*N,ZLEN,zdim,dtype=torch.float32,device=dev) + dit, zdim = load_traj_dit() + dit = dit.to(torch.float32).eval() + x = torch.randn(2 * N, WP, DIM, dtype=torch.float32, device=dev) + ts = torch.ones(2 * N, dtype=torch.int64, device=dev) + z = torch.randn(2 * N, ZLEN, zdim, dtype=torch.float32, device=dev) + class W(torch.nn.Module): - def __init__(s,d): super().__init__(); s.d=d - def forward(s,x,timestep,z_latents): return s.d(x=x,timestep=timestep,z_latents=z_latents) - w=W(dit).eval() - with torch.no_grad(): ref=w(x,ts,z) + def __init__(s, d): + super().__init__() + s.d = d + + def forward(s, x, timestep, z_latents): return s.d(x=x, timestep=timestep, z_latents=z_latents) # noqa: E704 + w = W(dit).eval() + with torch.no_grad(): + ref = w(x, ts, z) print(f" forward OK -> {tuple(ref.shape)}") - os.makedirs(os.path.dirname(OUT),exist_ok=True) + os.makedirs(os.path.dirname(OUT), exist_ok=True) # dynamic: batch (dim0) + z_latents seq (dim1) - dyn={"x":{0:"batch"},"timestep":{0:"batch"},"z_latents":{0:"batch",1:"zlen"},"output":{0:"batch"}} + dyn = {"x": {0: "batch"}, "timestep": {0: "batch"}, "z_latents": {0: "batch", 1: "zlen"}, "output": {0: "batch"}} print(f"[2/3] Export ONNX (dynamo=False, opset 19) → {OUT}") with torch.inference_mode(): - torch.onnx.export(w,(x,ts,z),OUT,input_names=["x","timestep","z_latents"], - output_names=["output"],opset_version=19,do_constant_folding=True, - export_params=True,dynamic_axes=dyn,dynamo=False) + torch.onnx.export(w, (x, ts, z), OUT, input_names=["x", "timestep", "z_latents"], + output_names=["output"], opset_version=19, do_constant_folding=True, + export_params=True, dynamic_axes=dyn, dynamo=False) print(f" {os.path.getsize(OUT)/1e6:.1f} MB") - B=2*N - cmd=["/usr/src/tensorrt/bin/trtexec",f"--onnx={OUT}",f"--saveEngine={ENG}","--bf16", - f"--minShapes=x:{B}x{WP}x{DIM},timestep:{B},z_latents:{B}x4x{zdim}", - f"--optShapes=x:{B}x{WP}x{DIM},timestep:{B},z_latents:{B}x{ZLEN}x{zdim}", - f"--maxShapes=x:{B}x{WP}x{DIM},timestep:{B},z_latents:{B}x64x{zdim}"] - print(f"[3/3] Build engine BF16 (z dynamic 4..64, opt {ZLEN})\n {' '.join(cmd)}") - r=subprocess.run(cmd,capture_output=True,text=True) - print(" "+ "\n ".join(l for l in r.stdout.splitlines() if "successfully" in l.lower() or "PASSED" in l or "FAILED" in l)[-400:]) - print(" engine:", ENG, os.path.getsize(ENG)/1e6 if os.path.exists(ENG) else "MISSING", "MB") + B = 2 * N + cmd = ["/usr/src/tensorrt/bin/trtexec", f"--onnx={OUT}", f"--saveEngine={ENG}", "--bf16", + f"--minShapes=x:{B}x{WP}x{DIM},timestep:{B},z_latents:{B}x4x{zdim}", + f"--optShapes=x:{B}x{WP}x{DIM},timestep:{B},z_latents:{B}x{ZLEN}x{zdim}", + f"--maxShapes=x:{B}x{WP}x{DIM},timestep:{B},z_latents:{B}x64x{zdim}"] + print(f"[3/3] Build engine BF16 (z dynamic 4..64, opt {ZLEN})") + print(" " + " ".join(cmd)) + r = subprocess.run(cmd, capture_output=True, text=True) + wanted = [ln for ln in r.stdout.splitlines() + if any(w in ln for w in ("successfully", "PASSED", "FAILED"))] + print(" " + "\n ".join(wanted)[-400:]) + print(" engine:", ENG, os.path.getsize(ENG) / 1e6 if os.path.exists(ENG) else "MISSING", "MB") return 0 if os.path.exists(ENG) else 1 -if __name__=="__main__": +if __name__ == "__main__": sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_accuracy.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_accuracy.py index eb87904..2078e2f 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_accuracy.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_accuracy.py @@ -5,25 +5,37 @@ (InternVLAN1AsyncAgent) run verbatim on sample data: branch decision + pixel goal, compared frame-by-frame for both PyTorch and the FP8 TensorRT engine. """ -import os, sys, json, subprocess, glob -_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, _R); sys.path.insert(0, os.path.join(_R, "lib")) -import numpy as np, torch -from PIL import Image -from engine_runner import ENGINE +import os +import sys +import json +import subprocess +import glob +_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, _R) +sys.path.insert(0, os.path.join(_R, "lib")) +import numpy as np # noqa: E402 +import torch # noqa: E402 +from PIL import Image # noqa: E402 +from engine_runner import ENGINE # noqa: E402 ACTIVE = os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) CKPT = os.path.join(ACTIVE, "checkpoints/InternVLA-N1-DualVLN") TRT = os.environ.get("TRT_EDGE_LLM", os.path.expanduser("~/TensorRT-Edge-LLM")) ENG_LLM = os.path.dirname(ENGINE) -ENG_VIS = os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "engines/system2_visual") -SCRATCH = os.path.expanduser(os.environ.get('VLN_OPT_OUT','~/vln-opt-work/out')) +ENG_VIS = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "engines/system2_visual") +SCRATCH = os.path.expanduser(os.environ.get('VLN_OPT_OUT', '~/vln-opt-work/out')) # camera_intrinsic matches the demo (inference_only_demo cell 15) -INTR = np.array([[386.5,0,328.9,0],[0,386.5,244,0],[0,0,1,0],[0,0,0,1]]) +INTR = np.array([[386.5, 0, 328.9, 0], [0, 386.5, 244, 0], [0, 0, 1, 0], [0, 0, 0, 1]]) class Args: - device="cuda:0"; model_path=CKPT; model_path_original=CKPT - resize_w=384; resize_h=384; num_history=8; plan_step_gap=4 + device = "cuda:0" + model_path = CKPT + model_path_original = CKPT + resize_w = 384 + resize_h = 384 + num_history = 8 + plan_step_gap = 4 def main(): @@ -34,16 +46,20 @@ def main(): import importlib.util _s = importlib.util.spec_from_file_location( "iar", os.path.join(ACTIVE, "internnav/agent/internvla_n1_agent_realworld.py")) - _m = importlib.util.module_from_spec(_s); _s.loader.exec_module(_m) + _m = importlib.util.module_from_spec(_s) + _s.loader.exec_module(_m) Agent = _m.InternVLAN1AsyncAgent print("[1/4] Load the documented reference agent") agent = Agent(Args()) - dbg = os.path.join(SCRATCH, "verb_dbg"); os.makedirs(dbg, exist_ok=True) - IMG_DIR = os.path.join(SCRATCH, "verb_imgs"); os.makedirs(IMG_DIR, exist_ok=True) + dbg = os.path.join(SCRATCH, "verb_dbg") + os.makedirs(dbg, exist_ok=True) + IMG_DIR = os.path.join(SCRATCH, "verb_imgs") + os.makedirs(IMG_DIR, exist_ok=True) cap = {} orig = agent.model.generate + def hooked(*a, **kw): cap["fired"] = True cap["input_ids"] = kw.get("input_ids") @@ -59,18 +75,20 @@ def hooked(*a, **kw): # documented demo order: sorted debug_raw_*.jpg (look_down frames interleaved) rgb_paths = sorted(glob.glob(os.path.join(scene, "debug_raw_*.jpg"))) print(f"[2/4] {sname} | instr={instr!r} | {len(rgb_paths)} frames (docs loop)") - agent.reset(); agent.save_dir = dbg + agent.reset() + agent.save_dir = dbg for p in rgb_paths: look_down = ('look_down' in p) rgb = np.asarray(Image.open(p).convert('RGB')) - depth = 10*np.ones((rgb.shape[0], rgb.shape[1]), np.float32) # docs: fill depth + depth = 10 * np.ones((rgb.shape[0], rgb.shape[1]), np.float32) # docs: fill depth pose = np.eye(4) cap.clear() try: with torch.no_grad(): agent.step(rgb, depth, pose, instr, intrinsic=INTR, look_down=look_down) except Exception as e: - print(f" {os.path.basename(p)}: step err {type(e).__name__}: {e}"); continue + print(f" {os.path.basename(p)}: step err {type(e).__name__}: {e}") + continue if not cap.get("fired"): continue # S2 did not run on this frame (buffer) - skip ref_out = agent.llm_output.strip() @@ -80,49 +98,61 @@ def hooked(*a, **kw): for it in turn["content"]: if it["type"] == "image": fp = os.path.join(IMG_DIR, f"s{len(samples)}_{ii}.png") - it["image"].save(fp); content.append({"type":"image","image":fp}); ii += 1 + it["image"].save(fp) + content.append({"type": "image", "image": fp}) + ii += 1 else: - content.append({"type":"text","text":it["text"]}) - msgs.append({"role":turn["role"],"content":content}) + content.append({"type": "text", "text": it["text"]}) + msgs.append({"role": turn["role"], "content": content}) samples.append(dict(scene=sname, img=os.path.basename(p), look_down=look_down, - ref=ref_out, S=int(cap["input_ids"].shape[1]) if cap.get("input_ids") is not None else -1, + ref=ref_out, S=int(cap["input_ids"].shape[1]) if cap.get( + "input_ids") is not None else -1, messages=msgs)) print(f" {os.path.basename(p):32} look_down={look_down} S={samples[-1]['S']} ref={ref_out!r}") - in_json = os.path.join(SCRATCH, "verb_in.json"); out_json = os.path.join(SCRATCH, "verb_out.json") - json.dump({"batch_size":1,"temperature":0.0,"top_p":1.0,"top_k":1,"max_generate_length":128, - "requests":[{"messages":s["messages"]} for s in samples]}, open(in_json,"w")) - del agent; torch.cuda.empty_cache() + in_json = os.path.join(SCRATCH, "verb_in.json") + out_json = os.path.join(SCRATCH, "verb_out.json") + json.dump({"batch_size": 1, "temperature": 0.0, "top_p": 1.0, "top_k": 1, "max_generate_length": 128, + "requests": [{"messages": s["messages"]} for s in samples]}, open(in_json, "w")) + del agent + torch.cuda.empty_cache() print(f"[3/4] Engine {os.path.basename(os.path.dirname(ENGINE))} through llm_inference ({len(samples)} req)") env = dict(os.environ, EDGELLM_PLUGIN_PATH=f"{TRT}/build/libNvInfer_edgellm_plugin.so") - r = subprocess.run([f"{TRT}/build/examples/llm/llm_inference","--engineDir",ENG_LLM, - "--multimodalEngineDir",ENG_VIS,"--inputFile",in_json,"--outputFile",out_json], + r = subprocess.run([f"{TRT}/build/examples/llm/llm_inference", "--engineDir", ENG_LLM, + "--multimodalEngineDir", ENG_VIS, "--inputFile", in_json, "--outputFile", out_json], cwd=TRT, env=env, capture_output=True, text=True) - resp = json.load(open(out_json)).get("responses",[]) if os.path.exists(out_json) else [] + resp = json.load(open(out_json)).get("responses", []) if os.path.exists(out_json) else [] print(f" exit={r.returncode}, {len(resp)} responses") - print("[4/4] Compare engine against the reference verbatim\n" + "="*66) + print("[4/4] Compare engine against the reference verbatim\n" + "=" * 66) import re + def xy(t): - d=[int(c) for c in re.findall(r"\d+",t)]; return (d[0],d[1]) if len(d)>=2 else None - ex=br=0; l2=[] - for i,s in enumerate(samples): - fp = resp[i].get("output_text","").strip() if i" - s["eng"]=fp - rb="coord" if any(c.isdigit() for c in s["ref"]) else "action" - fb="coord" if any(c.isdigit() for c in fp) else "action" - ex+=s["ref"]==fp; br+=rb==fb - a,b=xy(s["ref"]),xy(fp) - if a and b: l2.append(((a[0]-b[0])**2+(a[1]-b[1])**2)**.5) + d = [int(c) for c in re.findall(r"\d+", t)] + return (d[0], d[1]) if len(d) >= 2 else None + ex = br = 0 + l2 = [] + for i, s in enumerate(samples): + fp = resp[i].get("output_text", "").strip() if i < len(resp) else "" + s["eng"] = fp + rb = "coord" if any(c.isdigit() for c in s["ref"]) else "action" + fb = "coord" if any(c.isdigit() for c in fp) else "action" + ex += s["ref"] == fp + br += rb == fb + a, b = xy(s["ref"]), xy(fp) + if a and b: + l2.append(((a[0] - b[0])**2 + (a[1] - b[1])**2)**.5) print(f" [{'OK' if s['ref']==fp else 'DIFF':4}] {s['img']:32} ref={s['ref']!r:18} eng={fp[:40]!r}") - n=len(samples) - print("\n"+"="*66) + n = len(samples) + print("\n" + "=" * 66) print(f" exact-match : {ex}/{n} = {ex/n*100:.1f}%") print(f" branch-agree : {br}/{n} = {br/n*100:.1f}%") - if l2: print(f" coord-L2 px : mean={sum(l2)/len(l2):.1f} median={sorted(l2)[len(l2)//2]:.1f} max={max(l2):.0f} (n={len(l2)})") - tag = os.path.basename(os.path.dirname(ENGINE)).replace("edgellm_engines_","") - json.dump(samples, open(os.path.join(os.environ.get("VLN_OPT_OUT",os.path.expanduser("~/vln-opt-work/out")),f"verbatim_{tag}.json"),"w"), + if l2: + print( + f" coord-L2 px : mean={sum(l2)/len(l2):.1f} median={sorted(l2)[len(l2)//2]:.1f} max={max(l2):.0f} (n={len(l2)})") # noqa: E501 + tag = os.path.basename(os.path.dirname(ENGINE)).replace("edgellm_engines_", "") + json.dump(samples, open(os.path.join(os.environ.get("VLN_OPT_OUT", os.path.expanduser("~/vln-opt-work/out")), f"verbatim_{tag}.json"), "w"), # noqa: E501 indent=2, ensure_ascii=False) return 0 diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_e2e_agent.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_e2e_agent.py index b98afa2..2b4920a 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_e2e_agent.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_e2e_agent.py @@ -5,120 +5,174 @@ (FP8 LLM for the bridge + BF16 System 1 engines) and compare its outputs, frame by frame, against the pure-PyTorch agent on the documented sample data. """ -import os, sys, glob, time -_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, _R); sys.path.insert(0, os.path.join(_R, "lib")) +import os +import sys +import glob +_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, _R) +sys.path.insert(0, os.path.join(_R, "lib")) sys.path.append("/usr/lib/python3.12/dist-packages") -sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),"lib")) -import numpy as np, torch -from PIL import Image -from engine_runner import build_mrope_table, run_engine, ENGINE # FP8 LLM engine harness -from memblock import MemBlock - -ACTIVE=os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) -CKPT=os.path.join(ACTIVE,"checkpoints/InternVLA-N1-DualVLN") -REPKG=os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "qwen25vl_system2") -TRAJDIT=os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "onnx/system1_traj_dit_async_bf16.engine") -MEM=os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.engine") -SCRATCH=os.path.expanduser(os.environ.get('VLN_OPT_OUT','~/vln-opt-work/out')) -INTR=np.array([[386.5,0,328.9,0],[0,386.5,244,0],[0,0,1,0],[0,0,0,1]]) -TRAJ_TOKEN_INDEX=151667; IMAGE_TOKEN_INDEX=151655; SEED=12345 - - -def out_of(o,k): - if isinstance(o,dict): return o.get(k, next(iter(o.values()))) - return o[0] if isinstance(o,(list,tuple)) else o +sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "lib")) +import numpy as np # noqa: E402 +import torch # noqa: E402 +from PIL import Image # noqa: E402 +from engine_runner import build_mrope_table, run_engine # LLM engine harness # noqa: E402 + +ACTIVE = os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) +CKPT = os.path.join(ACTIVE, "checkpoints/InternVLA-N1-DualVLN") +REPKG = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "qwen25vl_system2") +TRAJDIT = os.path.join( + os.environ.get( + "WORK_DIR", + os.path.expanduser("~/vln-opt-work")), + "onnx/system1_traj_dit_async_bf16.engine") # noqa: E131 +MEM = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.engine") +SCRATCH = os.path.expanduser(os.environ.get('VLN_OPT_OUT', '~/vln-opt-work/out')) +INTR = np.array([[386.5, 0, 328.9, 0], [0, 386.5, 244, 0], [0, 0, 1, 0], [0, 0, 0, 1]]) +TRAJ_TOKEN_INDEX = 151667 +IMAGE_TOKEN_INDEX = 151655 +SEED = 12345 + + +def out_of(o, k): + if isinstance(o, dict): + return o.get(k, next(iter(o.values()))) + return o[0] if isinstance(o, (list, tuple)) else o def main(): - dev="cuda" - try: torch.backends.mha.set_fastpath_enabled(False) - except Exception: pass - if ACTIVE not in sys.path: sys.path.insert(0, ACTIVE) + dev = "cuda" + try: + torch.backends.mha.set_fastpath_enabled(False) + except Exception: + pass + if ACTIVE not in sys.path: + sys.path.insert(0, ACTIVE) from internvla_compat import apply_all apply_all(need_system1=True, allow_missing_depth=True) import importlib.util - _s=importlib.util.spec_from_file_location("iar",os.path.join(ACTIVE,"internnav/agent/internvla_n1_agent_realworld.py")) - _m=importlib.util.module_from_spec(_s); _s.loader.exec_module(_m) - Agent=_m.InternVLAN1AsyncAgent + _s = importlib.util.spec_from_file_location("iar", os.path.join( + ACTIVE, "internnav/agent/internvla_n1_agent_realworld.py")) + _m = importlib.util.module_from_spec(_s) + _s.loader.exec_module(_m) + Agent = _m.InternVLAN1AsyncAgent from diffusers.schedulers import FlowMatchEulerDiscreteScheduler from diffusers.utils.torch_utils import randn_tensor from trt_torch import Engine - class A: device="cuda:0"; model_path=CKPT; model_path_original=CKPT; resize_w=384; resize_h=384; num_history=8; plan_step_gap=4 + class A: + device = "cuda:0" + model_path = CKPT + model_path_original = CKPT + resize_w = 384 + resize_h = 384 + num_history = 8 + plan_step_gap = 4 print("[1/5] Load agent + engines") - agent=Agent(A()) - agent.save_dir=os.path.join(SCRATCH,"trtagent_dbg"); os.makedirs(agent.save_dir,exist_ok=True) - model=agent.model; m=model.get_model() - mem_eng=Engine(MEM); dit_eng=Engine(TRAJDIT) - lm = m.language_model if hasattr(m,"language_model") else m - final_norm=lm.norm - NQ=model.get_n_query() - rmean=model._resnet_mean; rstd=model._resnet_std + agent = Agent(A()) + agent.save_dir = os.path.join(SCRATCH, "trtagent_dbg") + os.makedirs(agent.save_dir, exist_ok=True) + model = agent.model + m = model.get_model() + mem_eng = Engine(MEM) + dit_eng = Engine(TRAJDIT) + lm = m.language_model if hasattr(m, "language_model") else m + final_norm = lm.norm + NQ = model.get_n_query() + rmean = model._resnet_mean + rstd = model._resnet_std # ---- FP8 LLM generate_latents (latent_query + mRoPE logic, run on the engine) ---- def trt_generate_latents(input_ids, pixel_values, image_grid_thw): with torch.no_grad(): - te=m.embed_tokens(input_ids) - ie=model.visual(pixel_values.type(model.visual.dtype), grid_thw=image_grid_thw) - te[input_ids==IMAGE_TOKEN_INDEX]=ie.to(te.dtype)[:(input_ids==IMAGE_TOKEN_INDEX).sum(),:] - lq=m.latent_queries.repeat(te.shape[0],1,1) - inputs_embeds=torch.cat([te,lq],dim=1) - ids_traj=torch.cat([input_ids, torch.tensor([[TRAJ_TOKEN_INDEX]*NQ],device=dev)],dim=1) - position_ids,_=model.get_rope_index(ids_traj, image_grid_thw) - rope=build_mrope_table(position_ids, dev) - _, eng_hs=run_engine(inputs_embeds.to(torch.float16), rope) - eng_pre=eng_hs[:,-NQ:,:].float() + te = m.embed_tokens(input_ids) + ie = model.visual(pixel_values.type(model.visual.dtype), grid_thw=image_grid_thw) + te[input_ids == IMAGE_TOKEN_INDEX] = ie.to(te.dtype)[:(input_ids == IMAGE_TOKEN_INDEX).sum(), :] + lq = m.latent_queries.repeat(te.shape[0], 1, 1) + inputs_embeds = torch.cat([te, lq], dim=1) + ids_traj = torch.cat([input_ids, torch.tensor([[TRAJ_TOKEN_INDEX] * NQ], device=dev)], dim=1) + position_ids, _ = model.get_rope_index(ids_traj, image_grid_thw) + rope = build_mrope_table(position_ids, dev) + _, eng_hs = run_engine(inputs_embeds.to(torch.float16), rope) + eng_pre = eng_hs[:, -NQ:, :].float() return final_norm(eng_pre.to(torch.bfloat16)) # hidden [1,NQ,3584] (as generate_latents) # ---- S1 generate_traj through the engines (async logic matches the baseline) ---- def trt_generate_traj(traj_latents, images_dp, depths_dp=None, predict_step_nums=32, guidance_scale=1.0, num_inference_steps=10, num_sample_trajs=32): - torch.manual_seed(SEED); np.random.seed(SEED) - dtype=traj_latents.dtype + torch.manual_seed(SEED) + np.random.seed(SEED) + dtype = traj_latents.dtype with torch.no_grad(): - tl=m.cond_projector(traj_latents) - xdp=images_dp.permute(0,1,4,2,3); xdp=((xdp-rmean)/rstd).flatten(0,1) - memory_tokens=out_of(mem_eng(images=xdp.float().contiguous()),"memory_tokens").to(dtype) - hs=torch.cat([memory_tokens,tl],dim=1) - hs_in=torch.cat([torch.zeros_like(hs),hs],0) - bs=tl.shape[0] - latents=randn_tensor((bs*num_sample_trajs,predict_step_nums,3),generator=None,device=dev,dtype=dtype) - sch=FlowMatchEulerDiscreteScheduler(); sch.set_timesteps(num_inference_steps,sigmas=np.linspace(1.0,1/num_inference_steps,num_inference_steps)) - hs_in=hs_in.repeat_interleave(num_sample_trajs,dim=0) - dit_eng.set_runtime_tensor_shape("z_latents",tuple(hs_in.shape)) + tl = m.cond_projector(traj_latents) + xdp = images_dp.permute(0, 1, 4, 2, 3) + xdp = ((xdp - rmean) / rstd).flatten(0, 1) + memory_tokens = out_of(mem_eng(images=xdp.float().contiguous()), "memory_tokens").to(dtype) + hs = torch.cat([memory_tokens, tl], dim=1) + hs_in = torch.cat([torch.zeros_like(hs), hs], 0) + bs = tl.shape[0] + latents = randn_tensor((bs * num_sample_trajs, predict_step_nums, 3), + generator=None, device=dev, dtype=dtype) + sch = FlowMatchEulerDiscreteScheduler() + sch.set_timesteps( + num_inference_steps, # noqa: E122 + sigmas=np.linspace( # noqa: E122 + 1.0, + 1 / num_inference_steps, + num_inference_steps)) # noqa: E131 + hs_in = hs_in.repeat_interleave(num_sample_trajs, dim=0) + dit_eng.set_runtime_tensor_shape("z_latents", tuple(hs_in.shape)) for t in sch.timesteps: - lf=m.action_encoder(latents) - pid=torch.arange(lf.shape[1]).reshape(1,-1).repeat(bs*num_sample_trajs,1).to(dev) - lf=lf+m.pos_encoding(pid); lmi=lf.repeat(2,1,1) - if hasattr(sch,"scale_model_input"): lmi=sch.scale_model_input(lmi,t) - tt=t.unsqueeze(0).expand(lmi.shape[0]).to(dev,torch.long) - npd=out_of(dit_eng(x=lmi.float().contiguous(),timestep=tt.to(torch.int64).contiguous(),z_latents=hs_in.float().contiguous()),"output").to(dtype) - npd=m.action_decoder(npd); unc,cnd=npd.chunk(2); npd=unc+guidance_scale*(cnd-unc) - latents=sch.step(npd,t,latents).prev_sample + lf = m.action_encoder(latents) + pid = torch.arange(lf.shape[1]).reshape(1, -1).repeat(bs * num_sample_trajs, 1).to(dev) + lf = lf + m.pos_encoding(pid) + lmi = lf.repeat(2, 1, 1) + if hasattr(sch, "scale_model_input"): + lmi = sch.scale_model_input(lmi, t) + tt = t.unsqueeze(0).expand(lmi.shape[0]).to(dev, torch.long) + npd = out_of( + dit_eng( # noqa: E122 + x=lmi.float().contiguous(), + timestep=tt.to( + torch.int64).contiguous(), + z_latents=hs_in.float().contiguous()), # noqa: E131 + "output").to(dtype) # noqa: E122 + npd = m.action_decoder(npd) + unc, cnd = npd.chunk(2) + npd = unc + guidance_scale * (cnd - unc) + latents = sch.step(npd, t, latents).prev_sample return latents # PyTorch reference generate_traj: fixed seed for a fair comparison - base_traj=model.generate_traj - def seeded_base_traj(*a,**k): - torch.manual_seed(SEED); np.random.seed(SEED); return base_traj(*a,**k) + base_traj = model.generate_traj + + def seeded_base_traj(*a, **k): + torch.manual_seed(SEED) + np.random.seed(SEED) + return base_traj(*a, **k) - scene=sorted(g for g in glob.glob(os.path.join(ACTIVE,"assets/realworld_sample_data*")) if os.path.isdir(g))[0] - instr=open(os.path.join(scene,"instruction.txt")).read().strip() - rgbs=sorted(glob.glob(os.path.join(scene,"debug_raw_*.jpg")))[:40] + scene = sorted(g for g in glob.glob(os.path.join(ACTIVE, "assets/realworld_sample_data*")) if os.path.isdir(g))[0] + instr = open(os.path.join(scene, "instruction.txt")).read().strip() + rgbs = sorted(glob.glob(os.path.join(scene, "debug_raw_*.jpg")))[:40] print(f"[2/5] scene {os.path.basename(scene)} | {len(rgbs)} frames | instr={instr[:50]!r}...") def run(tag): - agent.reset(); agent.save_dir=os.path.join(SCRATCH,"trtagent_dbg") - outs=[] + agent.reset() + agent.save_dir = os.path.join(SCRATCH, "trtagent_dbg") + outs = [] for p in rgbs: - ld=('look_down' in p); rgb=np.asarray(Image.open(p).convert('RGB')) - depth=10*np.ones(rgb.shape[:2],np.float32); pose=np.eye(4) + ld = ('look_down' in p) + rgb = np.asarray(Image.open(p).convert('RGB')) + depth = 10 * np.ones(rgb.shape[:2], np.float32) + pose = np.eye(4) try: - with torch.no_grad(): o=agent.step(rgb,depth,pose,instr,intrinsic=INTR,look_down=ld) + with torch.no_grad(): + o = agent.step(rgb, depth, pose, instr, intrinsic=INTR, look_down=ld) except Exception as e: - print(f" {os.path.basename(p)}: {type(e).__name__}: {e}"); continue + print(f" {os.path.basename(p)}: {type(e).__name__}: {e}") + continue traj = o.output_trajectory - act = o.output_action + act = o.output_action if traj is not None: outs.append(("traj", os.path.basename(p), np.asarray(traj))) elif act is not None: @@ -126,33 +180,41 @@ def run(tag): return outs print("[3/5] Run PyTorch agent (reference)") - model.generate_traj=seeded_base_traj - ref=run("pytorch") + model.generate_traj = seeded_base_traj + ref = run("pytorch") print(f" {len(ref)} outputs") print("[4/5] Run TensorRT agent (FP8 LLM latents + S1 engines)") - model.generate_latents=trt_generate_latents - model.generate_traj=trt_generate_traj - trt=run("trt") + model.generate_latents = trt_generate_latents + model.generate_traj = trt_generate_traj + trt = run("trt") print(f" {len(trt)} outputs") - print("[5/5] Compare e2e TRT agent vs PyTorch\n"+"="*56) - n=min(len(ref),len(trt)); mact=0; nact=0; trajerr=[] + print("[5/5] Compare e2e TRT agent vs PyTorch\n" + "=" * 56) + n = min(len(ref), len(trt)) + mact = 0 + nact = 0 + trajerr = [] for i in range(n): - rk,rf,rv=ref[i]; tk,tf,tv=trt[i] - if rk=="act" and tk=="act": - nact+=1; mact+= (rv==tv) - elif rk=="traj" and tk=="traj": - e=np.linalg.norm(np.asarray(rv)-np.asarray(tv),axis=-1) + rk, rf, rv = ref[i] + tk, tf, tv = trt[i] + if rk == "act" and tk == "act": + nact += 1 + mact += (rv == tv) + elif rk == "traj" and tk == "traj": + e = np.linalg.norm(np.asarray(rv) - np.asarray(tv), axis=-1) trajerr.append(float(np.mean(e))) - print(f" outputs: pytorch={len(ref)} trt={len(trt)} (type match {sum(1 for i in range(n) if ref[i][0]==trt[i][0])}/{n})") - if nact: print(f" action match: {mact}/{nact}") + print( + f" outputs: pytorch={len(ref)} trt={len(trt)} (type match {sum(1 for i in range(n) if ref[i][0]==trt[i][0])}/{n})") # noqa: E501 + if nact: + print(f" action match: {mact}/{nact}") if trajerr: import statistics - print(f" trajectory per-wp L2 (m): mean={statistics.mean(trajerr):.4f} max={max(trajerr):.4f} (n={len(trajerr)})") - print(" Agent TRT runs end-to-end and matches PyTorch." if n>0 else " no output") + print( + f" trajectory per-wp L2 (m): mean={statistics.mean(trajerr):.4f} max={max(trajerr):.4f} (n={len(trajerr)})") # noqa: E501 + print(" Agent TRT runs end-to-end and matches PyTorch." if n > 0 else " no output") return 0 -if __name__=="__main__": +if __name__ == "__main__": sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_engine_policy.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_engine_policy.py index 0aa9616..d7d6ca5 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_engine_policy.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_engine_policy.py @@ -24,8 +24,8 @@ sys.path.insert(0, _R) sys.path.insert(0, os.path.join(_R, "lib")) -import numpy as np -from PIL import Image +import numpy as np # noqa: E402 +from PIL import Image # noqa: E402 # prompt_builder is the single source of truth for the VLN prompt and lives in the # quantize path, so calibration and verification cannot drift apart. Walk up to the # recipe root rather than counting directory levels -- these scripts sit at two @@ -34,7 +34,7 @@ while _d != "/" and not os.path.isdir(os.path.join(_d, "quantize")): _d = os.path.dirname(_d) sys.path.insert(0, os.path.join(_d, "quantize")) -import prompt_builder as pb +import prompt_builder as pb # noqa: E402 TRT = os.path.expanduser(os.environ.get("TRT_EDGE_LLM", "~/modelopt/TensorRT-Edge-LLM")) LLM_DIR = os.path.expanduser(os.environ.get( @@ -42,7 +42,7 @@ VIS_DIR = os.path.expanduser(os.environ.get( "VLN_VIS_ENGINE_DIR", os.path.join(os.environ.get("ENGINE_DIR", - os.path.expanduser("~/vln-opt-work/engines")), "s1_fp8/visual"))) + os.path.expanduser("~/vln-opt-work/engines")), "s1_fp8/visual"))) REPKG = os.path.expanduser(os.environ.get("REPKG", "~/vln-opt-work/repro/qwen25vl_system2")) DATA = os.path.expanduser(os.environ.get("VLN_CALIB_DATA", "~/vln-opt-work/probe_heldout")) env = dict(os.environ, EDGELLM_PLUGIN_PATH=os.path.join(TRT, "build/libNvInfer_edgellm_plugin.so")) @@ -69,7 +69,7 @@ def one_lookdown_messages(tmp): messages + save images, using the same prompt_builder the policy uses.""" meta = sorted(glob.glob(os.path.join(DATA, "**", "meta", "episodes.jsonl"), recursive=True))[0] scene = os.path.dirname(os.path.dirname(meta)) - eps = [json.loads(l) for l in open(meta) if l.strip()] + eps = [json.loads(l) for l in open(meta) if l.strip()] # noqa: E741 ep = [e for e in eps if e.get("length", 0) > 20 and e.get("tasks")][0] t = ep["length"] // 2 lvl = f"{scene}/videos/chunk-000/observation.images.rgb.125cm_0deg" @@ -87,14 +87,16 @@ def one_lookdown_messages(tmp): # the real policy resizes to (resize_w, resize_h)=384 before inference Image.open(srcs[k]).convert("RGB").resize( (pb.RESIZE_W, pb.RESIZE_H)).save(d) - it["image"] = d; k += 1 + it["image"] = d + k += 1 paths.append(d) # to llm_inference messages (images already file paths) return conv def run_llm_inference(messages, tmp): - in_json = os.path.join(tmp, "in.json"); out_json = os.path.join(tmp, "out.json") + in_json = os.path.join(tmp, "in.json") + out_json = os.path.join(tmp, "out.json") json.dump({"batch_size": 1, "temperature": 0.0, "top_p": 1.0, "top_k": 1, "max_generate_length": 32, "requests": [{"messages": messages}]}, open(in_json, "w")) subprocess.run([os.path.join(TRT, "build/examples/llm/llm_inference"), diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents.py index d9d0035..5779602 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents.py @@ -4,12 +4,16 @@ """Verify System 2 numeric fidelity: z_latents cosine of the FP8 LLM engine vs the PyTorch reference, using the real latent-query bridge path. """ -import os, sys, json -_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, _R); sys.path.insert(0, _R) -import torch -from PIL import Image -from engine_runner import (REPKG, ENGINE, CKPT, build_mrope_table, run_engine, - cos, per_tok_cos) +import os +import sys +import json +_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, _R) +sys.path.insert(0, _R) +import torch # noqa: E402 +from PIL import Image # noqa: E402 +from engine_runner import (REPKG, ENGINE, CKPT, build_mrope_table, run_engine, # noqa: E402 + cos, per_tok_cos) TRAJ_TOKEN_INDEX = 151667 IMAGE_TOKEN_INDEX = 151655 @@ -117,12 +121,12 @@ def main(): from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration print(f"[1/6] Load repackage + processor | engine={os.path.basename(ENGINE)}") model = Qwen2_5_VLForConditionalGeneration.from_pretrained( - REPKG, torch_dtype=torch.bfloat16, # flash-attn is not available on Jetson; sdpa is the supported path and + REPKG, torch_dtype=torch.bfloat16, # flash-attn is not available on Jetson; sdpa is the supported path and # is what the deployed agent uses too. attn_implementation=os.environ.get("ATTN_IMPL", "sdpa"), low_cpu_mem_usage=True).to(dev).eval() proc = AutoProcessor.from_pretrained(REPKG, trust_remote_code=True, - min_pixels=128*28*28, max_pixels=1024*28*28) + min_pixels=128 * 28 * 28, max_pixels=1024 * 28 * 28) inner = model.model # Qwen2_5_VLModel (takes inputs_embeds) lm = inner.language_model if hasattr(inner, "language_model") else inner final_norm = lm.norm @@ -173,6 +177,7 @@ def main(): eng_post = final_norm(eng_pre.to(torch.bfloat16)).float() print("[6/6] Compare z_latents (System1 traj_dit input)\n" + "=" * 60) + def cond_project(x): x = torch.nn.functional.linear(x, cp["0.weight"], cp.get("0.bias")) x = torch.nn.functional.gelu(x) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents_vln.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents_vln.py index 482ee21..f928de0 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents_vln.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents_vln.py @@ -14,6 +14,7 @@ VLN_CALIB_DATA. Uses a fixed seed so every engine sees identical inputs. """ import os +import json import sys _R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) @@ -21,9 +22,9 @@ sys.path.insert(0, _R) sys.path.insert(0, os.path.join(_R, "build", "quantize")) -import torch +import torch # noqa: E402 -from engine_runner import (REPKG, ENGINE, CKPT, build_mrope_table, run_engine, +from engine_runner import (REPKG, ENGINE, CKPT, build_mrope_table, run_engine, # noqa: E402 cos, per_tok_cos) TRAJ_TOKEN_INDEX = 151667 @@ -136,7 +137,7 @@ def main(): print(f"[1/4] Load repackage + processor | engine={os.path.basename(ENGINE)}") model = Qwen2_5_VLForConditionalGeneration.from_pretrained( - REPKG, torch_dtype=torch.bfloat16, # flash-attn is not available on Jetson; sdpa is the supported path and + REPKG, torch_dtype=torch.bfloat16, # flash-attn is not available on Jetson; sdpa is the supported path and # is what the deployed agent uses too. attn_implementation=os.environ.get("ATTN_IMPL", "sdpa"), low_cpu_mem_usage=True).to(dev).eval() @@ -196,7 +197,10 @@ def cond_project(x): hpo = per_tok_cos(eng_post, ref_post) zc = per_tok_cos(z_eng, z_ref) zf = cos(z_eng, z_ref) - hid_pre.append(hp); hid_post.append(hpo); zc_list.append(zc); zflat_list.append(zf) + hid_pre.append(hp) + hid_post.append(hpo) + zc_list.append(zc) + zflat_list.append(zf) print(f" #{bi:2d} imgs={grid.shape[0]:2d} seq={input_ids.shape[1]:4d} " f"| hidPRE={hp:.5f} hidPOST={hpo:.5f} z={zc:.5f}") diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_pixelgoal_gt.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_pixelgoal_gt.py index 025a09c..8c30f37 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_pixelgoal_gt.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_pixelgoal_gt.py @@ -32,8 +32,8 @@ sys.path.insert(0, _R) sys.path.insert(0, os.path.join(_R, "lib")) -import numpy as np -from PIL import Image +import numpy as np # noqa: E402 +from PIL import Image # noqa: E402 # prompt_builder is the single source of truth for the VLN prompt and lives in the # quantize path, so calibration and verification cannot drift apart. Walk up to the # recipe root rather than counting directory levels -- these scripts sit at two @@ -42,14 +42,14 @@ while _d != "/" and not os.path.isdir(os.path.join(_d, "quantize")): _d = os.path.dirname(_d) sys.path.insert(0, os.path.join(_d, "quantize")) -import prompt_builder as pb +import prompt_builder as pb # noqa: E402 TRT = os.path.expanduser(os.environ.get("TRT_EDGE_LLM", "~/modelopt/TensorRT-Edge-LLM")) REPKG = os.path.expanduser(os.environ.get("REPKG", "~/vln-opt-work/repro/qwen25vl_system2")) VIS = os.path.expanduser(os.environ.get( "VIS_ENG", os.path.join(os.environ.get("ENGINE_DIR", - os.path.expanduser("~/vln-opt-work/engines")), "s1_fp8/visual"))) + os.path.expanduser("~/vln-opt-work/engines")), "s1_fp8/visual"))) DATA = os.path.expanduser(os.environ.get("VLN_CALIB_DATA", "~/vln-opt-work/probe_heldout")) # Look-down config: level history (0deg) + tilted goal view. GT lives at the tilted pitch. LEVEL_KEY = "observation.images.rgb.125cm_0deg" @@ -74,7 +74,7 @@ def collect_goal_frames(): for meta in glob.glob(os.path.join(DATA, "**", "meta", "episodes.jsonl"), recursive=True): scene = os.path.dirname(os.path.dirname(meta)) eps = {e["episode_index"]: e for e in - (json.loads(l) for l in open(meta) if l.strip())} + (json.loads(l) for l in open(meta) if l.strip())} # noqa: E741 for pq in sorted(glob.glob(os.path.join(scene, "data", "chunk-000", "*.parquet"))): df = pd.read_parquet(pq) ep_idx = int(df["episode_index"].iloc[0]) @@ -116,7 +116,8 @@ def pred_uv(text): def run_engine(conv, tmp): js = {"batch_size": 1, "temperature": 0.0, "top_p": 1.0, "top_k": 1, "max_generate_length": 16, "requests": [{"messages": conv}]} - inp = os.path.join(tmp, "in.json"); out = os.path.join(tmp, "out.json") + inp = os.path.join(tmp, "in.json") + out = os.path.join(tmp, "out.json") json.dump(js, open(inp, "w")) subprocess.run([f"{TRT}/build/examples/llm/llm_inference", "--engineDir", os.path.dirname(ENGINE), "--multimodalEngineDir", VIS, @@ -145,7 +146,7 @@ def main(): REPKG, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2", low_cpu_mem_usage=True).to("cuda").eval() proc = AutoProcessor.from_pretrained(REPKG, trust_remote_code=True, - min_pixels=128*28*28, max_pixels=1024*28*28) + min_pixels=128 * 28 * 28, max_pixels=1024 * 28 * 28) pt_l2, eng_l2 = [], [] tmp = tempfile.mkdtemp(prefix="pgoal_") @@ -187,13 +188,15 @@ def main(): eng_l2.append(np.hypot(euv[0] - gu, euv[1] - gv)) def ci(a, stat=np.median, n=1000, seed=0): - a = np.asarray(a); rng = np.random.default_rng(seed) + a = np.asarray(a) + rng = np.random.default_rng(seed) bs = [stat(rng.choice(a, len(a), replace=True)) for _ in range(n)] return np.percentile(bs, 2.5), np.percentile(bs, 97.5) def report(a, tag): a = np.asarray(a) - mlo, mhi = ci(a, np.median); alo, ahi = ci(a, np.mean) + mlo, mhi = ci(a, np.median) + alo, ahi = ci(a, np.mean) print(f" {tag:13}: n={len(a)} median={np.median(a):.1f} [{mlo:.1f},{mhi:.1f}] " f"mean={a.mean():.1f} [{alo:.1f},{ahi:.1f}] max={a.max():.0f}") diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_system1.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_system1.py index 7ee3695..381d623 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_system1.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_system1.py @@ -4,32 +4,44 @@ """Verify System 1: wire the traj_dit + memory TensorRT engines into generate_traj and compare the resulting trajectory and latency against the PyTorch base (same noise seed). """ -import os, sys, time -_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, _R); sys.path.insert(0, os.path.join(_R, "lib")) +import os +import sys +import time +_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, _R) +sys.path.insert(0, os.path.join(_R, "lib")) sys.path.append("/usr/lib/python3.12/dist-packages") -sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),"lib")) -import numpy as np, torch -from PIL import Image -from memblock import MemBlock +sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "lib")) +import numpy as np # noqa: E402 +import torch # noqa: E402 +from PIL import Image # noqa: E402 -ACTIVE=os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) -CKPT=os.path.join(ACTIVE,"checkpoints/InternVLA-N1-DualVLN") -IMG=os.path.expanduser("~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg") -TRAJDIT=os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "onnx/system1_traj_dit_bf16.engine") -MEM=os.path.join(os.environ.get("WORK_DIR",os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.engine") -SEED=12345 +ACTIVE = os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) +CKPT = os.path.join(ACTIVE, "checkpoints/InternVLA-N1-DualVLN") +IMG = os.path.expanduser("~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg") +TRAJDIT = os.path.join( + os.environ.get( + "WORK_DIR", + os.path.expanduser("~/vln-opt-work")), + "onnx/system1_traj_dit_bf16.engine") # noqa: E131 +MEM = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.engine") +SEED = 12345 def out_of(o, key): - if isinstance(o, dict): return o.get(key, next(iter(o.values()))) - return o[0] if isinstance(o,(list,tuple)) else o + if isinstance(o, dict): + return o.get(key, next(iter(o.values()))) + return o[0] if isinstance(o, (list, tuple)) else o def main(): - dev="cuda" - try: torch.backends.mha.set_fastpath_enabled(False) - except Exception: pass - if ACTIVE not in sys.path: sys.path.insert(0, ACTIVE) + dev = "cuda" + try: + torch.backends.mha.set_fastpath_enabled(False) + except Exception: + pass + if ACTIVE not in sys.path: + sys.path.insert(0, ACTIVE) from internvla_compat import apply_all apply_all(need_system1=True, allow_missing_depth=True) from internnav.model.basemodel.internvla_n1.internvla_n1 import ( @@ -38,77 +50,93 @@ def main(): from diffusers.utils.torch_utils import randn_tensor from trt_torch import Engine print("[1/4] Load model + engines") - cfg=InternVLAN1ModelConfig.from_pretrained(CKPT) - model=InternVLAN1ForCausalLM.from_pretrained(CKPT,config=cfg,torch_dtype=torch.bfloat16, - attn_implementation="sdpa",low_cpu_mem_usage=True).to(dev).eval() - m=model.get_model() - mem_eng=Engine(MEM); dit_eng=Engine(TRAJDIT) - NQ=model.get_n_query() - z=torch.randn(1,NQ,cfg.hidden_size,device=dev,dtype=torch.bfloat16) - a=np.array(Image.open(IMG).convert("RGB").resize((224,224)))/255.0 - tt=torch.from_numpy(a).float(); images_dp=torch.stack([tt,tt]).unsqueeze(0).to(dev) - rmean=model._resnet_mean; rstd=model._resnet_std + cfg = InternVLAN1ModelConfig.from_pretrained(CKPT) + model = InternVLAN1ForCausalLM.from_pretrained(CKPT, config=cfg, torch_dtype=torch.bfloat16, + attn_implementation="sdpa", low_cpu_mem_usage=True).to(dev).eval() + m = model.get_model() + mem_eng = Engine(MEM) + dit_eng = Engine(TRAJDIT) + NQ = model.get_n_query() + z = torch.randn(1, NQ, cfg.hidden_size, device=dev, dtype=torch.bfloat16) + a = np.array(Image.open(IMG).convert("RGB").resize((224, 224))) / 255.0 + tt = torch.from_numpy(a).float() + images_dp = torch.stack([tt, tt]).unsqueeze(0).to(dev) + rmean = model._resnet_mean + rstd = model._resnet_std def base_gen(): - torch.manual_seed(SEED); np.random.seed(SEED) + torch.manual_seed(SEED) + np.random.seed(SEED) with torch.no_grad(): - return model.generate_traj(z,images_dp,num_sample_trajs=32,num_inference_steps=10).float().cpu() + return model.generate_traj(z, images_dp, num_sample_trajs=32, num_inference_steps=10).float().cpu() def trt_gen(steps=10, ns=32, guidance_scale=1.0, predict_step_nums=32): """Replicate generate_traj (async) with the memory + traj_dit engines.""" - torch.manual_seed(SEED); np.random.seed(SEED) - dtype=z.dtype + torch.manual_seed(SEED) + np.random.seed(SEED) + dtype = z.dtype with torch.no_grad(): traj_latents = m.cond_projector(z) # [1,4,768] # --- memory block through the engine --- - xdp = images_dp.permute(0,1,4,2,3) - xdp = ((xdp - rmean)/rstd).flatten(0,1) # [2,3,224,224] + xdp = images_dp.permute(0, 1, 4, 2, 3) + xdp = ((xdp - rmean) / rstd).flatten(0, 1) # [2,3,224,224] memory_tokens = out_of(mem_eng(images=xdp.float().contiguous()), "memory_tokens").to(dtype) # [1,32,768] hidden_states = torch.cat([memory_tokens, traj_latents], dim=1) # [1,36,768] hs_null = torch.zeros_like(hidden_states) hs_input = torch.cat([hs_null, hidden_states], 0) # [2,36,768] bs = traj_latents.shape[0] - latents = randn_tensor((bs*ns, predict_step_nums, 3), generator=None, device=dev, dtype=dtype) + latents = randn_tensor((bs * ns, predict_step_nums, 3), generator=None, device=dev, dtype=dtype) sch = FlowMatchEulerDiscreteScheduler() - sigmas = np.linspace(1.0, 1/steps, steps) + sigmas = np.linspace(1.0, 1 / steps, steps) sch.set_timesteps(steps, sigmas=sigmas) hs_input = hs_input.repeat_interleave(ns, dim=0) # [2*ns,36,768] dit_eng.set_runtime_tensor_shape("z_latents", tuple(hs_input.shape)) for t in sch.timesteps: lf = m.action_encoder(latents) - pos_ids = torch.arange(lf.shape[1]).reshape(1,-1).repeat(bs*ns,1).to(dev) + pos_ids = torch.arange(lf.shape[1]).reshape(1, -1).repeat(bs * ns, 1).to(dev) lf = lf + m.pos_encoding(pos_ids) - lmi = lf.repeat(2,1,1) - if hasattr(sch,"scale_model_input"): lmi = sch.scale_model_input(lmi, t) + lmi = lf.repeat(2, 1, 1) + if hasattr(sch, "scale_model_input"): + lmi = sch.scale_model_input(lmi, t) tt_ = t.unsqueeze(0).expand(lmi.shape[0]).to(dev, torch.long) np_ = out_of(dit_eng(x=lmi.float().contiguous(), timestep=tt_.to(torch.int64).contiguous(), z_latents=hs_input.float().contiguous()), "output").to(dtype) np_ = m.action_decoder(np_) unc, cnd = np_.chunk(2) - np_ = unc + guidance_scale*(cnd - unc) + np_ = unc + guidance_scale * (cnd - unc) latents = sch.step(np_, t, latents).prev_sample return latents.float().cpu() - def lat(fn,warm=2,n=5): - for _ in range(warm): fn() - torch.cuda.synchronize(); ts=[] + def lat(fn, warm=2, n=5): + for _ in range(warm): + fn() + torch.cuda.synchronize() + ts = [] for _ in range(n): - torch.cuda.synchronize(); t0=time.perf_counter(); fn(); torch.cuda.synchronize(); ts.append(time.perf_counter()-t0) - return sum(ts)/len(ts)*1000 + torch.cuda.synchronize() + t0 = time.perf_counter() + fn() + torch.cuda.synchronize() + ts.append(time.perf_counter() - t0) + return sum(ts) / len(ts) * 1000 print("[2/4] Base PyTorch generate_traj") - tref = base_gen(); pt_ms = lat(base_gen) + tref = base_gen() + pt_ms = lat(base_gen) print(f" {tuple(tref.shape)} | {pt_ms:.1f}ms = {1000/pt_ms:.1f}Hz") print("[3/4] Full-TRT S1 (memory+traj_dit engines)") - ttrt = trt_gen(); trt_ms = lat(trt_gen) + ttrt = trt_gen() + trt_ms = lat(trt_gen) print(f" {tuple(ttrt.shape)} | {trt_ms:.1f}ms = {1000/trt_ms:.1f}Hz") - print("[4/4] Compare\n"+"="*56) - d=(tref-ttrt).norm(dim=-1); endp=(tref[:,-1]-ttrt[:,-1]).norm(dim=-1) - cos=torch.nn.functional.cosine_similarity(tref.flatten(),ttrt.flatten(),dim=0).item() + print("[4/4] Compare\n" + "=" * 56) + d = (tref - ttrt).norm(dim=-1) + endp = (tref[:, -1] - ttrt[:, -1]).norm(dim=-1) + cos = torch.nn.functional.cosine_similarity(tref.flatten(), ttrt.flatten(), dim=0).item() print(f" Parity: per-wp L2 mean={d.mean():.4f} max={d.max():.4f} | endpoint={endp.mean():.4f} | cos={cos:.5f}") - print(f" Latency S1: PyTorch {pt_ms:.1f}ms ({1000/pt_ms:.1f}Hz) → full-TRT {trt_ms:.1f}ms ({1000/trt_ms:.1f}Hz) = {pt_ms/trt_ms:.2f}x") + print( + f" Latency S1: PyTorch {pt_ms:.1f}ms ({1000/pt_ms:.1f}Hz) → full-TRT {trt_ms:.1f}ms ({1000/trt_ms:.1f}Hz) = {pt_ms/trt_ms:.2f}x") # noqa: E501 return 0 -if __name__=="__main__": +if __name__ == "__main__": sys.exit(main()) From 0612fa881ae6c5c64ff4913adac16600ff326f51 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Thu, 13 Aug 2026 19:23:18 +0700 Subject: [PATCH 15/30] feat(internvla-n1-dualvln): build the NVFP4 engine and locate the rest of its collapse Rebuilding the NVFP4 engine with the fork's CASK workaround settles most of the question the investigation opened. Three measurements, each isolating a layer: FP8, weights only (PyTorch) 0.998020 NVFP4, weights only (PyTorch) 0.987986 NVFP4 engine, with the CASK workaround 0.931005 NVFP4 engine, no workaround (source figure) 0.647 So the 0.647 was mostly a compiler artifact. -cask_fusion:max_num_epilogues=1 -- applied automatically by the fork, gated to NVFP4 graphs at batch 1 -- moves it to 0.931, and the build log confirms all three flags fired. Weight quantization was never the problem: NVFP4 weights cost 0.988, an error 3.5x FP8's with the bridge degrading in proportion. A real gap remains between 0.988 and 0.931, and it now has a named suspect rather than being a mystery: NVFP4 is W4A4, and 4-bit activations through a 3584-wide hidden state in blocks of 16 are the part the PyTorch weights-only path cannot model. The 0.647 is quoted from the source project and was not reproduced here. What is measured here is that the same checkpoint reaches 0.931 once the workaround is applied. NVFP4 stays experimental: 0.931 is below the 0.99 gate. But it is far from broken, and the 5.5 GB engine is the smallest of the three. Adds 06_measure.sh to collect prefill/decode latency for every built engine into the JSON run_matrix.py reads, and run_matrix.py itself, which assembles sizes, latency, bridge fidelity and task accuracy into one table without recomputing anything. (cherry picked from commit bd71e87dc203256d94b52787d9d62fc60b8b345e) --- recipes/internvla-n1-dualvln/README.md | 62 +++++---- recipes/internvla-n1-dualvln/run_matrix.py | 124 ++++++++++++++++++ .../trt-edgellm/scripts/06_measure.sh | 55 ++++++++ 3 files changed, 212 insertions(+), 29 deletions(-) create mode 100644 recipes/internvla-n1-dualvln/run_matrix.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/scripts/06_measure.sh diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index d1dd272..ecbfb7b 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -127,37 +127,41 @@ engine and skip their in-script parity check with a message rather than failing; project) generates; it now fails with an explanation of the expected file shape rather than a bare `FileNotFoundError` at the first read. -### NVFP4 — the 0.647 collapse is not weight quantization +### NVFP4 — the collapse had two causes, and one of them was the compiler -`trt-edgellm/investigate_nvfp4.py` was written to find out why NVFP4 keeps text fluent while -the bridge collapses. Measured here, in PyTorch, weights only: +`trt-edgellm/investigate_nvfp4.py` was written to explain why NVFP4 keeps text fluent while +the bridge collapses. Three measurements, each isolating a different layer: -| | weight rel-err | weight cos | **z_latents** | -|---|---|---|---| -| FP8 | 2.67 % | 0.999644 | **0.998020** | -| NVFP4 | 9.45 % | 0.995534 | **0.987986** | - -**NVFP4 weight quantization costs 0.988, not 0.647.** The weight error is 3.5x FP8's and the -bridge degrades roughly in proportion — nothing anomalous. So whatever produces 0.647 is -*not* the weights, and the two remaining candidates are the parts this measurement does not -model: 4-bit **activation** quantization, and the engine itself. - -The engine hypothesis deserves weight here rather than dismissal. This platform has already -produced one "quantization is broken" conclusion that turned out to be a TensorRT miscompile -(`fc_h_fusion`, see below), and there is a second known one specific to NVFP4: CASK -miscompiles when it fuses two or more epilogues into one NVFP4 GEMM at batch 1, fixed by -`-cask_fusion:max_num_epilogues=1`. A 0.647 measured on an engine built without that -workaround would be measuring the miscompile, not the format. - -Channel analysis rules out the obvious remedy. At the final layer the top 128 channels by -magnitude carry only 29.5 % of the squared error, and masking them *lowers* cosine rather -than restoring it — the error is spread, not concentrated in outliers. So AWQ scaling or a -targeted exclusion cannot recover the weight-side loss; that part is a capacity limit. - -**Where this leaves NVFP4:** experimental, and the next step is not more weight analysis. It -is to rebuild the NVFP4 engine *with* the CASK workaround and re-measure z_latents end to -end. If that lands near 0.988 the format is usable and the old number was a compiler -artifact; if it stays at 0.647 the cause is activation quantization. +| | weight rel-err | z_latents | +|---|---|---| +| FP8, weights only (PyTorch) | 2.67 % | 0.998020 | +| **NVFP4, weights only (PyTorch)** | 9.45 % | **0.987986** | +| **NVFP4 engine, with the CASK workaround** | — | **0.931005** | +| NVFP4 engine, no workaround (source project's figure) | — | 0.647 | + +Reading them together: + +* **Weight quantization is not the problem.** NVFP4 weights cost 0.988 — error 3.5x FP8's, + with the bridge degrading roughly in proportion. Nothing anomalous. +* **Most of the old 0.647 was a compiler artifact.** Rebuilding the engine with + `-cask_fusion:max_num_epilogues=1` (which the fork applies automatically, gated to NVFP4 + graphs at batch 1) moves it to 0.931. The build log confirms all three flags fired: + `-peep:fc_h_fusion=off -peep:match_dual_gemm=off -cask_fusion:max_num_epilogues=1`. +* **A real gap remains.** 0.988 weights-only versus 0.931 through the engine is the part + this platform's PyTorch path cannot model: NVFP4 is W4A4, and 4-bit *activations* through + a 3584-wide hidden state in blocks of 16 are the remaining suspect. + +The 0.647 figure is quoted from the source project and was not reproduced here; what is +measured here is that the same checkpoint reaches 0.931 once the workaround is applied. + +Channel analysis rules out the obvious remedy for the weight-side loss: at the final layer +the top 128 channels by magnitude carry only 29.5 % of the squared error, and masking them +*lowers* cosine rather than restoring it. The error is spread, not concentrated in outliers, +so AWQ scaling or a targeted exclusion has nothing to grip. + +**Verdict: NVFP4 stays experimental.** 0.931 is below the 0.99 gate, so it is not +recommended for navigation — but it is far from the broken 0.647 it appeared to be, and the +remaining gap now has a named suspect rather than a mystery. ### Earlier full-pipeline figures diff --git a/recipes/internvla-n1-dualvln/run_matrix.py b/recipes/internvla-n1-dualvln/run_matrix.py new file mode 100644 index 0000000..baa009b --- /dev/null +++ b/recipes/internvla-n1-dualvln/run_matrix.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Collect every measurement for every built variant into one comparison table. + +Reads what is already on disk rather than recomputing: checkpoint sizes, engine sizes, +the accuracy JSON from benchmark_accuracy.py, the NVFP4 investigation JSON, and the +z_latents numbers passed in. Anything missing prints as "-" instead of failing, so the +table can be produced at any point in the pipeline. + +Emits Markdown, ready to paste into the recipe README. +""" +import argparse +import json +import os + + +def dir_size_gb(path: str) -> float | None: + if not path or not os.path.isdir(path): + return None + total = 0 + for root, _dirs, files in os.walk(path): + for name in files: + try: + total += os.path.getsize(os.path.join(root, name)) + except OSError: + pass + return total / 1e9 + + +def file_size_gb(path: str) -> float | None: + return os.path.getsize(path) / 1e9 if path and os.path.isfile(path) else None + + +def fmt(value, spec: str = ".2f", suffix: str = "") -> str: + return "-" if value is None else f"{value:{spec}}{suffix}" + + +def load_json(path: str) -> dict | list | None: + if path and os.path.isfile(path): + with open(path) as f: + return json.load(f) + return None + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--work_dir", default=os.path.expanduser("~/vln-opt-work")) + p.add_argument("--output_path", default=None, help="Write the Markdown here") + args = p.parse_args() + w = args.work_dir + + # variant -> (checkpoint dir, engine dir, label) + variants = [ + ("BF16 (unquantized)", "qwen25vl_system2", "engines/base_fp16"), + ("FP8 s1", "qwen25vl_s1_fp8", "engines/s1_fp8"), + ("NVFP4 s1 (experimental)", "qwen25vl_s1_nvfp4", "engines/s1_nvfp4"), + ] + + acc = load_json(os.path.join(w, "out/accuracy_bf16_vs_fp8.json")) or [] + acc_by_path = {os.path.basename(r["model_path"]): r for r in acc} + inv = load_json(os.path.join(w, "nvfp4_investigation.json")) or {} + lat = load_json(os.path.join(w, "out/latency.json")) or {} + zlat = load_json(os.path.join(w, "out/z_latents.json")) or {} + + rows = [] + for label, ckpt_name, eng_name in variants: + ckpt = os.path.join(w, ckpt_name) + eng_llm = os.path.join(w, eng_name, "llm/llm.engine") + eng_vis = os.path.join(w, eng_name, "visual/visual.engine") + a = acc_by_path.get(ckpt_name) + rows.append({ + "label": label, + "ckpt_gb": dir_size_gb(ckpt), + "llm_gb": file_size_gb(eng_llm), + "vis_gb": file_size_gb(eng_vis), + "prefill_ms": lat.get(ckpt_name, {}).get("prefill_ms"), + "decode_ms": lat.get(ckpt_name, {}).get("decode_ms"), + "z_engine": zlat.get(ckpt_name), + "z_weights": inv.get("bridge", {}).get( + {"qwen25vl_s1_fp8": "fp8", "qwen25vl_s1_nvfp4": "nvfp4"}.get(ckpt_name, "")), + "l2_mean": a and a.get("pixel_goal_l2_mean"), + "l2_median": a and a.get("pixel_goal_l2_median"), + "w_err": inv.get("weights", {}).get( + {"qwen25vl_s1_fp8": "fp8", "qwen25vl_s1_nvfp4": "nvfp4"}.get(ckpt_name, ""), + {}).get("rel_err_mean"), + }) + + out = [] + out.append("### Benchmark matrix\n") + out.append("Every column is a real measurement on this machine; `-` means not measured.\n") + out.append("| Variant | checkpoint | LLM engine | visual | prefill | decode " + "| z_latents (engine) | z_latents (weights) | pixel L2 mean / median |") + out.append("|---|---|---|---|---|---|---|---|---|") + for r in rows: + l2 = ("-" if r["l2_mean"] is None + else f"{r['l2_mean']:.2f} / {r['l2_median']:.2f} px") + out.append( + f"| {r['label']} | {fmt(r['ckpt_gb'], '.1f', ' GB')} | " + f"{fmt(r['llm_gb'], '.2f', ' GB')} | {fmt(r['vis_gb'], '.2f', ' GB')} | " + f"{fmt(r['prefill_ms'], '.1f', ' ms')} | {fmt(r['decode_ms'], '.1f', ' ms')} | " + f"{fmt(r['z_engine'], '.6f')} | {fmt(r['z_weights'], '.6f')} | {l2} |") + + out.append("\n**Weight quantization error** (mean relative, 21 projections across " + "layers 0/13/27):\n") + out.append("| Variant | rel-err |") + out.append("|---|---|") + for r in rows: + if r["w_err"] is not None: + out.append(f"| {r['label']} | {100 * r['w_err']:.2f} % |") + + text = "\n".join(out) + print(text) + if args.output_path: + os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True) + with open(args.output_path, "w") as f: + f.write(text + "\n") + print(f"\nWrote {args.output_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/scripts/06_measure.sh b/recipes/internvla-n1-dualvln/trt-edgellm/scripts/06_measure.sh new file mode 100644 index 0000000..7ff4c7f --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/scripts/06_measure.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +# +# Measure prefill/decode latency and bridge fidelity for every built engine, and write the +# numbers where run_matrix.py can find them. +# +# llm_bench reads base_config.json while llm_build writes config.json, so the copy below is +# not optional -- without it the benchmark cannot open an engine it was just handed. + +set -euo pipefail + +WORK_DIR="${WORK_DIR:-$HOME/vln-opt-work}" +ENGINE_DIR="${ENGINE_DIR:-$WORK_DIR/engines}" +TRT_EDGELLM_DIR="${TRT_EDGELLM_DIR:-$HOME/modelopt/Hung-TRT-Edge-LLM}" +INPUT_LEN="${INPUT_LEN:-1024}" +PAST_KV_LEN="${PAST_KV_LEN:-1024}" +OUT="${OUT:-$WORK_DIR/out/latency.json}" + +LLM_BENCH="$TRT_EDGELLM_DIR/build/examples/llm/llm_bench" +[[ -x "$LLM_BENCH" ]] || { echo "[ERROR] not found: $LLM_BENCH" >&2; exit 1; } +export EDGELLM_PLUGIN_PATH="$TRT_EDGELLM_DIR/build/libNvInfer_edgellm_plugin.so" + +# Map engine directory -> the checkpoint name run_matrix.py keys on. +declare -A CKPT_OF=( + [base_fp16]=qwen25vl_system2 + [s1_fp8]=qwen25vl_s1_fp8 + [s1_nvfp4]=qwen25vl_s1_nvfp4 +) + +mkdir -p "$(dirname "$OUT")" +echo "{" > "$OUT" +first=1 + +for name in "${!CKPT_OF[@]}"; do + engine="$ENGINE_DIR/$name/llm" + [[ -f "$engine/llm.engine" ]] || { echo " [skip] $name: no engine"; continue; } + [[ -f "$engine/base_config.json" ]] || cp "$engine/config.json" "$engine/base_config.json" + + echo " measuring $name ..." + prefill=$("$LLM_BENCH" --engineDir "$engine" --mode prefill --inputLen "$INPUT_LEN" 2>&1 \ + | grep -oE "E2E Time \(actual performance\): [0-9.]+" | tail -1 | grep -oE "[0-9.]+$" || echo "") + decode=$("$LLM_BENCH" --engineDir "$engine" --mode decode --pastKVLen "$PAST_KV_LEN" 2>&1 \ + | grep -oE "E2E Time \(actual performance\): [0-9.]+" | tail -1 | grep -oE "[0-9.]+$" || echo "") + + [[ $first -eq 1 ]] || echo "," >> "$OUT" + first=0 + printf ' "%s": {"prefill_ms": %s, "decode_ms": %s}' \ + "${CKPT_OF[$name]}" "${prefill:-null}" "${decode:-null}" >> "$OUT" + echo " prefill ${prefill:-n/a} ms | decode ${decode:-n/a} ms" +done + +echo "" >> "$OUT" +echo "}" >> "$OUT" +echo "Wrote $OUT" From cc5cc15f459c579fc8dbc7eb0fe65ffe5882a7cf Mon Sep 17 00:00:00 2001 From: hungho77 Date: Thu, 13 Aug 2026 20:43:15 +0700 Subject: [PATCH 16/30] docs(internvla-n1-dualvln): the full three-variant benchmark matrix All three variants built from one repackaged System 2 and measured on an idle GPU: variant ckpt LLM engine prefill decode z_latents pixel L2 mean/median BF16 16.6 GB 14.15 GB 135.8ms 56.4ms 0.999471 47.24 / 27.05 px FP8 s1 10.1 GB 7.62 GB 82.1ms 31.5ms 0.991861 46.26 / 22.51 px NVFP4 s1 7.2 GB 4.77 GB 73.2ms 20.2ms 0.931005 40.69 / 23.54 px FP8 is the recommendation: 1.86x smaller and 1.65x/1.79x faster than BF16, bridge held at 0.9919, and the median waypoint error does not worsen. It reads slightly better (27.05 -> 22.51 px), which for 42 samples should be read as unchanged rather than as a gain from quantization. NVFP4 is smaller and faster still but its bridge sits at 0.931, below the gate, so it stays experimental. Two caveats are documented alongside the table because both cost time to learn here. Latency must be measured on an idle GPU: sharing the device with another job inflated FP8 prefill from 82 ms to 117 ms, and nothing in the scripts enforces this. The 'identical replies' count is a weak indicator and is now labelled as such. It moved from 31/42 to 11/42 for FP8 between two revisions of the checkpoint loader while pixel_goal_l2 barely moved. I initially attributed that to GPU contention; a clean re-run reproduced the contended numbers exactly, so the cause was the loader change, not contention. The current loader is verified correct -- all 729 weights in the loaded model match the dequantized state dict -- so the current numbers stand. Greedy decoding over a 152k vocabulary flips on tiny logit differences; the L2 median is the number to trust. (cherry picked from commit b60e7d3069229d8aae659c0e464c608c943447d1) --- recipes/internvla-n1-dualvln/README.md | 34 ++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index ecbfb7b..539d4dc 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -35,6 +35,40 @@ this recipe is gated on `z_latents` cosine against the FP32 reference, not on ge ## Results +### Benchmark matrix + +All three variants built from the same repackaged System 2 and measured on Jetson Thor, +batch 1, on an idle GPU. + +| Variant | checkpoint | LLM engine | visual | prefill (1024) | decode (pastKV 1024) | z_latents (engine) | z_latents (weights) | pixel L2 mean / median | +|---|---|---|---|---|---|---|---|---| +| BF16 (unquantized) | 16.6 GB | 14.15 GB | 1.36 GB | 135.8 ms | 56.4 ms | **0.999471** | — | 47.24 / 27.05 px | +| **FP8 s1** | 10.1 GB | **7.62 GB** | 1.36 GB | **82.1 ms** | **31.5 ms** | **0.991861** | 0.998020 | 46.26 / **22.51 px** | +| NVFP4 s1 (experimental) | 7.2 GB | 4.77 GB | 1.36 GB | 73.2 ms | 20.2 ms | 0.931005 ✗ | 0.987986 | 40.69 / 23.54 px | + +Weight quantization error, mean relative over 21 projections in layers 0/13/27: +FP8 **2.67 %**, NVFP4 **9.45 %**. + +**FP8 is the recommended scheme.** Against BF16 it is 1.86x smaller and 1.65x/1.79x faster, +holds the bridge at 0.9919, and the median waypoint error does not get worse — it improves +slightly (27.05 → 22.51 px), which is within the spread of a 42-sample set and should be read +as "unchanged", not as a gain from quantization. + +**NVFP4 is faster and smaller still but fails the gate.** Its bridge sits at 0.931, below the +0.99 threshold, so it is not recommended for navigation despite the attractive size and +latency. See the NVFP4 section for where that number comes from. + +### Two things to know before reading these numbers + +**Run everything on an idle GPU.** Latency measured while another job shared the device came +out 40-60 % higher (FP8 prefill 117 ms against 82 ms). The measurement script does not +enforce this. + +**The "identical replies" count is a weak indicator.** It moved from 31/42 to 11/42 for FP8 +across code revisions of the checkpoint loader, while `pixel_goal_l2` barely changed. Greedy +decoding over a 152k vocabulary flips on tiny logit differences, so treat the L2 median as the +number that matters and the identity count as colour. + ### Task accuracy — does the quantized model still pick the same waypoint? Measured with `quantize/benchmark_accuracy.py` on 42 held-out samples from two scenes, From 4451315246a091e97012a634848d27d58dc91339 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Fri, 14 Aug 2026 00:50:49 +0700 Subject: [PATCH 17/30] feat(internvla-n1-dualvln): support AWQ checkpoints, and add QAT AWQ does not help this model. End to end the bridge measures 0.986293 against plain NVFP4's 0.987986 -- the same within noise. That matches what the channel analysis predicted: the error is spread across channels rather than carried by outliers, so AWQ's per-channel rescaling has nothing to grip. Getting to that number needed two loader fixes, and the first reading was wrong in a way worth recording. AWQ-lite scales activations by a per-input-channel factor and stores the weight pre-divided by it, so a plain matmul has to multiply it back in. Ignoring the pre_quant_scale tensors entirely produced ~190% relative weight error, which reads as AWQ being catastrophically broken rather than being loaded wrong. The direction is also easy to invert: verified empirically on this checkpoint that multiplying gives cosine 0.9898, leaving it alone 0.9606, and dividing 0.8089. The second is subtler and invalidates a metric rather than a number. Only down_proj and o_proj carry an explicit pre_quant_scale, because q/k/v share one layernorm and gate/up share another -- for those, AWQ folds the scale into the layernorm weight instead. Measured here, AWQ moves the layernorm weights by 0.012x to 91x while plain NVFP4 leaves them at exactly 1.000. So comparing AWQ's linear weights against the reference's linear weights is not a comparison at all; the invariant is the composition. The weights stage now detects AWQ checkpoints and skips with that explanation rather than reporting a meaningless 180%. qat.py adds quantization-aware fine-tuning, which is the remaining lever: the gap between NVFP4 weights-only (0.988) and the engine (0.931) is 4-bit activations, and no weight-side method reaches those after the fact. ModelOpt has no separate QAT entry point -- it is mtq.quantize followed by ordinary fine-tuning with the fake-quantizers left in so gradients pass through them. It excludes scene YmJkqBEsHnH from training by default: the same MP3D building appears as calib_scenes/r2r and probe_heldout/rxr, and the source project has already reported one calibration gain that turned out to be exactly this leak. --allow_overlap disables the guard for anyone who wants to measure the effect deliberately. Note that success rate cannot be measured on this machine -- SR, SPL and NE are all closed-loop. QAT is evaluated on the proxies it can move, z_latents and pixel-goal L2, and the docstring says so. (cherry picked from commit 332bf0bc8cc40103bf2c92b44d0623414664eee3) --- .../quantize/load_quantized.py | 21 +- recipes/internvla-n1-dualvln/quantize/qat.py | 248 ++++++++++++++++++ .../trt-edgellm/investigate_nvfp4.py | 21 ++ 3 files changed, 287 insertions(+), 3 deletions(-) create mode 100644 recipes/internvla-n1-dualvln/quantize/qat.py diff --git a/recipes/internvla-n1-dualvln/quantize/load_quantized.py b/recipes/internvla-n1-dualvln/quantize/load_quantized.py index af0ba6b..fcc2542 100644 --- a/recipes/internvla-n1-dualvln/quantize/load_quantized.py +++ b/recipes/internvla-n1-dualvln/quantize/load_quantized.py @@ -91,19 +91,33 @@ def dequantize_state_dict(model_path: str, with safe_open(shard, framework="pt") as f: for key in f.keys(): tensor = f.get_tensor(key) - if key.endswith(("weight_scale", "input_scale", "weight_scale_2")): + if key.endswith(("weight_scale", "input_scale", "weight_scale_2", + "pre_quant_scale")): scales[key] = tensor else: raw[key] = tensor out: dict[str, torch.Tensor] = {} n_dequant = 0 + n_awq = 0 for key, tensor in raw.items(): # NVFP4: packed uint8 plus a block scale and a global scale. block = scales.get(key + "_scale") glob = scales.get(key + "_scale_2") if tensor.dtype == torch.uint8 and block is not None and glob is not None: - out[key] = unpack_nvfp4(tensor, block, glob, dtype) + w = unpack_nvfp4(tensor, block, glob, torch.float32) + # AWQ-lite scales the activations by a per-input-channel s and stores the + # weight pre-divided by it, so that y = (x * s) @ (W / s) reproduces x @ W. + # A plain matmul needs s multiplied back in. Verified empirically on this + # checkpoint against the unquantized weights, since the direction is easy to + # get backwards: multiplying gives cosine 0.9898, leaving it alone 0.9606, + # dividing 0.8089. Skipping this entirely reads as ~190% relative error and + # looks like AWQ being catastrophically bad rather than loaded wrong. + pqs = scales.get(key.replace(".weight", ".pre_quant_scale")) + if pqs is not None: + w = w * pqs.to(torch.float32) + n_awq += 1 + out[key] = w.to(dtype) n_dequant += 1 continue scale = scales.get(key + "_scale") @@ -116,7 +130,8 @@ def dequantize_state_dict(model_path: str, else: out[key] = tensor.to(dtype) if tensor.is_floating_point() else tensor - print(f" [load] dequantized {n_dequant} tensors, " + awq_note = f", {n_awq} with an AWQ pre-quant scale folded out" if n_awq else "" + print(f" [load] dequantized {n_dequant} tensors{awq_note}, " f"{len(out) - n_dequant} passed through unchanged") return out diff --git a/recipes/internvla-n1-dualvln/quantize/qat.py b/recipes/internvla-n1-dualvln/quantize/qat.py new file mode 100644 index 0000000..d60f381 --- /dev/null +++ b/recipes/internvla-n1-dualvln/quantize/qat.py @@ -0,0 +1,248 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Quantization-aware fine-tuning for the InternVLA-N1 System 2 planner. + +Why this exists. Post-training NVFP4 leaves the System 2 -> System 1 bridge at z_latents +0.931 through the engine, under the 0.99 gate that FP8 clears. Weight quantization is not +what costs it -- NVFP4 weights alone measure 0.988 -- so the remaining loss sits in the +4-bit *activations*, which no amount of weight-side scaling can fix after the fact. QAT is +the standard answer: let the model see the quantization noise during training and adapt to +it. + +What this does. ModelOpt has no separate QAT entry point; QAT is ``mtq.quantize`` followed +by ordinary fine-tuning, with the fake-quantizers left in place so gradients flow through +them (straight-through estimation). So this: + + 1. loads the repackaged System 2, + 2. calibrates and inserts quantizers exactly as ``quantize.py`` would, + 3. fine-tunes on real VLN episodes with the deployed prompt, + 4. exports through the same path, producing a checkpoint the existing engine build and + verification scripts accept unchanged. + +Read the result honestly. **Success rate is not measurable here** -- SR, SPL and NE are all +closed-loop and need Habitat or InternUtopia, neither of which runs on a Jetson. What this +optimises and what you can check is the proxy: z_latents cosine and pixel-goal L2 on +held-out episodes. A proxy that improves is a necessary condition for SR to improve, not +evidence that it did. + +Train/eval separation. The calibration and probe sets share one MP3D scene +(``YmJkqBEsHnH`` appears as ``calib_scenes/r2r/`` and ``probe_heldout/rxr/`` -- same +building, different split), and the source project has already been bitten by exactly this +overlap once, reporting a calibration gain that turned out to be leakage. That scene is +excluded from training by default; ``--allow_overlap`` turns the guard off if you want to +measure the effect of the leak deliberately. +""" +import argparse +import json +import os +import sys +import time + +import numpy as np +import torch + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import prompt_builder as pb # noqa: E402 +from model_loader import export_quantized_model, load_model # noqa: E402 +from quant_schemes import build_quant_config, calib_batch_size, validate # noqa: E402 + +# Same building as probe_heldout/rxr/YmJkqBEsHnH. Training on it contaminates the +# held-out evaluation even though the episodes and instructions differ. +OVERLAPPING_SCENES = ("YmJkqBEsHnH",) + + +def discover_training_samples(data_root: str, max_samples: int, camera: str, + allow_overlap: bool, seed: int = 0) -> list[dict]: + """Collect prompt/target pairs from LeRobot episodes, skipping leaked scenes.""" + import glob + import pyarrow.parquet as pq + + rng = np.random.default_rng(seed) + rgb_key = f"observation.images.rgb.{camera}" + level_key = "observation.images.rgb.125cm_0deg" + goal_col = f"goal.{camera}" + samples: list[dict] = [] + skipped_scenes: set[str] = set() + + for meta in sorted(glob.glob(os.path.join(data_root, "**", "meta", "episodes.jsonl"), + recursive=True)): + scene_dir = os.path.dirname(os.path.dirname(meta)) + scene = os.path.basename(scene_dir) + if not allow_overlap and scene in OVERLAPPING_SCENES: + skipped_scenes.add(scene) + continue + + for ep in (json.loads(line) for line in open(meta)): + idx, length = ep["episode_index"], ep["length"] + table = None + for parquet in glob.glob(os.path.join(scene_dir, "data", "**", + f"episode_{idx:06d}.parquet"), + recursive=True): + table = pq.read_table(parquet) + break + if table is None or goal_col not in table.schema.names: + continue + + goals = table.column(goal_col).to_pylist() + usable = [i for i in range(1, min(length, len(goals))) + if goals[i] is not None and int(goals[i][0]) >= 0] + if not usable: + continue + t = int(usable[rng.integers(0, len(usable))]) + + history = np.unique(np.linspace(0, t - 1, pb.NUM_HISTORY, dtype=np.int32)).tolist() + frames = [os.path.join(scene_dir, "videos", "chunk-000", level_key, + f"episode_{idx:06d}_{i}.jpg") for i in history + [t]] + lookdown = os.path.join(scene_dir, "videos", "chunk-000", rgb_key, + f"episode_{idx:06d}_{t}.jpg") + if not all(os.path.isfile(p) for p in frames + [lookdown]): + continue + + gt = goals[t] + samples.append({ + "episode": scene, + "episode_idx": t, + "instruction": (ep.get("tasks") or [""])[0], + "images": frames + [lookdown], + "turn": 2, + "assistant_turn1": "↓", + # The supervision target is the deployed answer format: "row col". + "target": f"{int(gt[0])} {int(gt[1])}", + }) + if len(samples) >= max_samples: + break + if len(samples) >= max_samples: + break + + if skipped_scenes: + print(f" [data] excluded {sorted(skipped_scenes)} -- also present in the " + f"held-out probe set") + return samples + + +def build_batch(sample: dict, processor, device: str): + """Prompt plus target, with the prompt masked out of the loss.""" + inputs = pb.build_sample_inputs(sample, processor) + prompt_len = inputs["input_ids"].shape[1] + + target_ids = processor.tokenizer(sample["target"], add_special_tokens=False, + return_tensors="pt")["input_ids"] + input_ids = torch.cat([inputs["input_ids"], target_ids], dim=1) + labels = input_ids.clone() + labels[:, :prompt_len] = -100 # supervise the answer only + + batch = {k: v for k, v in inputs.items() if k != "input_ids"} + batch["input_ids"] = input_ids + batch["labels"] = labels + if "attention_mask" in batch: + pad = torch.ones((1, target_ids.shape[1]), dtype=batch["attention_mask"].dtype) + batch["attention_mask"] = torch.cat([batch["attention_mask"], pad], dim=1) + return {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in batch.items()} + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--model_path", required=True, help="Repackaged System 2 checkpoint") + p.add_argument("--output_path", required=True) + p.add_argument("--data_root", required=True, help="LeRobot episodes for training") + p.add_argument("--scheme", default="nvfp4_default") + p.add_argument("--strategy", default="s1") + p.add_argument("--num_train_samples", type=int, default=64) + p.add_argument("--num_calib_samples", type=int, default=64) + p.add_argument("--epochs", type=int, default=1) + p.add_argument("--lr", type=float, default=1e-5, + help="Small on purpose: QAT adapts to quantization noise, it does not " + "re-learn the task, and a large step undoes the pretrained planner") + p.add_argument("--grad_accum", type=int, default=4) + p.add_argument("--camera", default="125cm_30deg") + p.add_argument("--allow_overlap", action="store_true", + help="Permit training on scenes that also appear in the probe set") + p.add_argument("--device", default="cuda") + p.add_argument("--seed", type=int, default=0) + return p.parse_args() + + +def main() -> int: + import modelopt.torch.quantization as mtq + + args = parse_args() + torch.manual_seed(args.seed) + + try: + validate(args.scheme, args.strategy, model_path=args.model_path, + allow_experimental=True) + except ValueError as exc: + print(f"[ERROR] {exc}") + return 1 + + print(f"[1/5] Loading {args.model_path}") + model, tokenizer, processor = load_model(args.model_path, dtype="bf16", + device=args.device) + if processor is None: + print("[ERROR] no processor; VLN prompts cannot be built") + return 1 + + print(f"[2/5] Collecting training samples from {args.data_root}") + train = discover_training_samples(args.data_root, args.num_train_samples, + args.camera, args.allow_overlap, args.seed) + if not train: + print("[ERROR] no usable training samples") + return 1 + print(f" {len(train)} samples from " + f"{len({s['episode'] for s in train})} scene(s)") + + print(f"[3/5] Inserting quantizers ({args.scheme} / {args.strategy})") + quant_cfg = build_quant_config(args.scheme, args.strategy) + calib = train[:min(args.num_calib_samples, len(train))] + + def forward_loop(m): + for s in calib: + with torch.no_grad(): + m(**build_batch(s, processor, args.device)) + + model = mtq.quantize(model, quant_cfg, forward_loop=forward_loop) + _ = calib_batch_size(args.scheme, is_image_calib=True) + + print(f"[4/5] Fine-tuning: {args.epochs} epoch(s), lr={args.lr}, " + f"grad_accum={args.grad_accum}") + # Gradients flow through the fake-quantizers by straight-through estimation, which is + # what lets the weights adapt to 4-bit activation noise rather than merely to 4-bit + # weights. + model.train() + model.config.use_cache = False + optim = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], + lr=args.lr, weight_decay=0.0) + step, t0, losses = 0, time.time(), [] + for epoch in range(args.epochs): + order = np.random.default_rng(args.seed + epoch).permutation(len(train)) + optim.zero_grad(set_to_none=True) + for n, i in enumerate(order, 1): + out = model(**build_batch(train[i], processor, args.device)) + loss = out.loss / args.grad_accum + loss.backward() + losses.append(float(out.loss)) + if n % args.grad_accum == 0: + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + optim.step() + optim.zero_grad(set_to_none=True) + step += 1 + if step % 4 == 0: + recent = float(np.mean(losses[-4 * args.grad_accum:])) + print(f" epoch {epoch} step {step:4d} loss {recent:.4f} " + f"{time.time() - t0:.0f}s", flush=True) + + print(f" trained {step} optimizer steps, final loss " + f"{float(np.mean(losses[-8:])):.4f}") + + print(f"[5/5] Exporting to {args.output_path}") + model.eval() + export_quantized_model(model, tokenizer, processor, + model_dir=args.model_path, output_dir=args.output_path) + print(f"Saved to {args.output_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/investigate_nvfp4.py b/recipes/internvla-n1-dualvln/trt-edgellm/investigate_nvfp4.py index c1c0f61..44520bd 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/investigate_nvfp4.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/investigate_nvfp4.py @@ -66,6 +66,17 @@ def rel_err(ref: torch.Tensor, other: torch.Tensor) -> float: return float((other.double() - ref).norm() / ref.norm()) +def _has_awq_scales(model_path: str) -> bool: + """True if the checkpoint carries AWQ pre-quant scales.""" + import glob + from safetensors import safe_open + for shard in sorted(glob.glob(os.path.join(model_path, "*.safetensors"))): + with safe_open(shard, framework="pt") as f: + if any(k.endswith("pre_quant_scale") for k in f.keys()): + return True + return False + + def stage_weights(args, report: dict) -> None: """Compare per-projection weight error, NVFP4 vs FP8, against the unquantized weights.""" from load_quantized import dequantize_state_dict @@ -88,6 +99,16 @@ def stage_weights(args, report: dict) -> None: for label, path in (("fp8", args.fp8_ckpt), ("nvfp4", args.nvfp4_ckpt)): if not path: continue + # AWQ redistributes scale between each layernorm and the projections that read it + # (measured on this checkpoint: layernorm weights move by 0.012x to 91x while + # plain NVFP4 leaves them at exactly 1.0). Comparing its linear weights against + # the reference's linear weights is therefore meaningless -- the invariant is the + # composition, not either factor. Only the end-to-end bridge stage is valid there. + if _has_awq_scales(path): + print(f" {label:6s} [skip] AWQ checkpoint: per-layer weight comparison is not " + f"meaningful, see the bridge stage") + rows[label] = {"skipped": "awq_rescales_layernorms"} + continue state = dequantize_state_dict(path) errs, coss = [], [] for k, ref in base.items(): From 76abd0aa2af9b53cce98c4364a5948dfeb0e436a Mon Sep 17 00:00:00 2001 From: hungho77 Date: Fri, 14 Aug 2026 00:58:57 +0700 Subject: [PATCH 18/30] docs(internvla-n1-dualvln): no post-training NVFP4 preset closes the weight-side gap All three NVFP4 presets measure the same end to end against the unquantized reference: nvfp4_default 0.987986 baseline nvfp4_awq_full 0.986293 equal within noise nvfp4_local_hessian 0.987986 byte-identical to default nvfp4_local_hessian is a no-op on this model, and it is worth stating plainly because it fails silently. The preset genuinely differs -- algorithm={'method': 'local_hessian', 'fp8_scale_sweep': True} against 'max' -- and the run exits 0, but the exported weights match nvfp4_default bit for bit: 0 of 6,422,528 bytes differ, scales included. Two runs with different algorithms cannot produce identical output unless the algorithm never ran. It stays selectable via --scheme, so a user would reasonably believe they had tried it. AWQ does run, its weights genuinely differ, and it still does not help -- which is what the channel analysis predicted: the error is spread across channels rather than carried by outliers, so per-channel rescaling has nothing to grip. Three weight-side methods stopping at the same 0.988 is the argument for QAT rather than a fourth. The gap from there to the engine's 0.931 is 4-bit activations, and no post-training method reaches those. (cherry picked from commit 1d1e8be52200c38b7924ce5499410fea83a05388) --- recipes/internvla-n1-dualvln/README.md | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index 539d4dc..a8a87da 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -193,6 +193,30 @@ the top 128 channels by magnitude carry only 29.5 % of the squared error, and ma *lowers* cosine rather than restoring it. The error is spread, not concentrated in outliers, so AWQ scaling or a targeted exclusion has nothing to grip. +**No post-training method closes the weight-side gap.** All three NVFP4 presets land in the +same place, measured end to end against the unquantized reference: + +| Preset | z_latents (weights) | note | +|---|---|---| +| `nvfp4_default` | 0.987986 | baseline | +| `nvfp4_awq_full` | 0.986293 | equal within noise | +| `nvfp4_local_hessian` | 0.987986 | **byte-identical to default** | + +`nvfp4_local_hessian` is a no-op on this model. Its preset genuinely differs +(`algorithm={'method': 'local_hessian', 'fp8_scale_sweep': True}` against `'max'`) and the +run exits 0, but the exported weights match `nvfp4_default` bit for bit — 0 of 6,422,528 +bytes differ, scales included. Two quantization runs with different algorithms cannot +produce identical output unless the algorithm did not run. It remains selectable, so a user +would reasonably believe they had tried it. + +AWQ does run — its weights genuinely differ — but does not help, which is what the channel +analysis predicted: the error is spread rather than carried by outliers, so per-channel +rescaling has nothing to grip. + +That leaves quantization-aware training as the only remaining lever, since the gap between +0.988 (weights) and 0.931 (engine) is 4-bit activations and no post-training method reaches +those. See `quantize/qat.py`. + **Verdict: NVFP4 stays experimental.** 0.931 is below the 0.99 gate, so it is not recommended for navigation — but it is far from the broken 0.647 it appeared to be, and the remaining gap now has a named suspect rather than a mystery. From d185f0a3a46114984346318814ca8c61249ec057 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Fri, 14 Aug 2026 02:39:36 +0700 Subject: [PATCH 19/30] feat(internvla-n1-dualvln): QAT, and an honest negative result from one run Adds quantization-aware fine-tuning and reports what it actually did, which was make things worse: z_latents (engine) pixel L2 mean / median NVFP4 PTQ 0.931005 40.69 / 23.54 px NVFP4 + QAT 0.891583 41.86 / 23.16 px The training loss rose from 0.78 to 1.13 over 16 steps, so the run pushed the weights the wrong way rather than converging. That makes this a failed training run, not evidence that QAT cannot help: the experiment never reached the question it was meant to answer. Saying 'QAT does not work for this model' from sixteen steps would be a conclusion the data does not support, so the README says what would change next instead -- a much smaller learning rate, hundreds of steps rather than sixteen, and a warmup. Two obstacles are worth recording because both cost a run to find. Full fine-tuning does not fit on this hardware. Weights plus gradients plus AdamW moments for 8.29 B parameters are ~100 GB before any activations, against a 122 GB pool shared with the host; the first attempt was killed by the OOM killer. --train_last_n_layers (default 4) makes it fit at 932 M trainable parameters, and it targets the right place anyway, since the bridge reads the last layer's hidden states. Freezing combined with gradient checkpointing needs two fixes applied together. The activations entering the first trainable layer carry no grad_fn, and reentrant checkpointing then discards the graph, so loss.backward() fails with 'element 0 of tensors does not require grad'. Both use_reentrant=False and enable_input_require_grads() are required; either alone still fails. Also documents a metric trap I walked into first: z_latents against the unquantized reference is the wrong measure for QAT. PTQ approximates the original model, so similarity to it is meaningful; QAT deliberately moves away from the original to absorb quantization noise, so a good QAT run can lower that number while improving behaviour. QAT is judged through the engine and on task accuracy instead. (cherry picked from commit 2f0b80709f477a8f5fc065d8498ce279d7742397) --- recipes/internvla-n1-dualvln/README.md | 40 ++++++++++++++ recipes/internvla-n1-dualvln/quantize/qat.py | 58 +++++++++++++++++++- 2 files changed, 96 insertions(+), 2 deletions(-) diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index a8a87da..2dffaed 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -217,6 +217,46 @@ That leaves quantization-aware training as the only remaining lever, since the g 0.988 (weights) and 0.931 (engine) is 4-bit activations and no post-training method reaches those. See `quantize/qat.py`. +### QAT was tried and made it worse — but the run did not converge + +One exploratory run: 64 samples, 16 optimizer steps, lr 1e-5, the last 4 decoder layers +trainable (932 M parameters), the rest frozen. + +| | z_latents (engine) | pixel L2 mean / median | +|---|---|---| +| NVFP4 PTQ | **0.931005** | 40.69 / 23.54 px | +| NVFP4 + QAT | **0.891583** | 41.86 / 23.16 px | + +Worse on the bridge, unchanged on the task within noise. That is consistent with the +training loss, which *rose* from 0.78 to 1.13 across the 16 steps — the run pushed the +weights in the wrong direction rather than converging. + +**Read this as a failed training run, not as evidence that QAT cannot work here.** The loss +never fell, so the experiment never reached the question it was meant to answer. What would +change next: a much smaller learning rate (1e-6 or below — QAT adapts to quantization noise, +it does not relearn the task, and 1e-5 over 932 M parameters is too large a step), several +hundred steps rather than sixteen, and a warmup instead of a flat schedule. + +Two practical notes for anyone repeating this: + +Full fine-tuning does not fit. At 8.29 B parameters, weights plus gradients plus AdamW +moments come to ~100 GB before any activations, against a 122 GB pool shared with the host — +the first attempt was killed by the OOM killer. `--train_last_n_layers` (default 4) is what +makes it fit, and it also targets the right place: the bridge reads the last layer's hidden +states. + +Freezing plus gradient checkpointing needs both fixes at once. The activations entering the +first trainable layer carry no `grad_fn`, and reentrant checkpointing then discards the +graph — `loss.backward()` fails with "element 0 of tensors does not require grad". Setting +`use_reentrant=False` **and** calling `enable_input_require_grads()` is required; either +alone still fails. + +Finally, note that `z_latents` against the unquantized reference is the wrong metric for +QAT and is only reported here through the engine. PTQ tries to approximate the original +model, so similarity to it is meaningful. QAT deliberately moves the weights away from the +original to compensate for quantization noise, so a successful QAT run can *lower* that +similarity while improving real behaviour. Judge QAT on the engine and on task accuracy. + **Verdict: NVFP4 stays experimental.** 0.931 is below the 0.99 gate, so it is not recommended for navigation — but it is far from the broken 0.647 it appeared to be, and the remaining gap now has a named suspect rather than a mystery. diff --git a/recipes/internvla-n1-dualvln/quantize/qat.py b/recipes/internvla-n1-dualvln/quantize/qat.py index d60f381..47e7208 100644 --- a/recipes/internvla-n1-dualvln/quantize/qat.py +++ b/recipes/internvla-n1-dualvln/quantize/qat.py @@ -141,6 +141,33 @@ def build_batch(sample: dict, processor, device: str): return {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in batch.items()} +def freeze_all_but_last_layers(model, n_last: int) -> int: + """Freeze everything except the final ``n_last`` decoder layers. + + Returns the number of frozen parameters. Full fine-tuning does not fit: at 8.29B + parameters, weights plus gradients plus AdamW moments come to about 100 GB before any + activation memory, against a 122 GB pool shared with the host. Restricting to the last + layers is also where the quantity being repaired lives -- the System 1 bridge reads the + last layer's hidden states. + """ + layers = None + for owner in (getattr(model, "model", None), model): + inner = getattr(owner, "language_model", owner) + if inner is not None and hasattr(inner, "layers"): + layers = inner.layers + break + if layers is None: + raise AttributeError("could not locate the decoder layer list on this model") + + keep = {id(p) for layer in layers[-n_last:] for p in layer.parameters()} + frozen = 0 + for param in model.parameters(): + if id(param) not in keep: + param.requires_grad_(False) + frozen += param.numel() + return frozen + + def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) @@ -156,6 +183,14 @@ def parse_args() -> argparse.Namespace: help="Small on purpose: QAT adapts to quantization noise, it does not " "re-learn the task, and a large step undoes the pretrained planner") p.add_argument("--grad_accum", type=int, default=4) + p.add_argument("--train_last_n_layers", type=int, default=4, + help="Train only the last N decoder layers; freeze the rest. Full " + "fine-tuning of the 8.29B planner needs ~100 GB for weights, " + "gradients and AdamW state alone, which the 122 GB unified pool " + "cannot hold alongside a 10-image prompt's activations. The bridge " + "reads the last layer's hidden states, so the last layers are also " + "where the signal it depends on is formed. 0 trains everything.") + p.add_argument("--gradient_checkpointing", action="store_true", default=True) p.add_argument("--camera", default="125cm_30deg") p.add_argument("--allow_overlap", action="store_true", help="Permit training on scenes that also appear in the probe set") @@ -212,8 +247,27 @@ def forward_loop(m): # weights. model.train() model.config.use_cache = False - optim = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], - lr=args.lr, weight_decay=0.0) + + if args.train_last_n_layers > 0: + n_frozen = freeze_all_but_last_layers(model, args.train_last_n_layers) + print(f" froze {n_frozen} parameters; training the last " + f"{args.train_last_n_layers} decoder layers") + if args.gradient_checkpointing and hasattr(model, "gradient_checkpointing_enable"): + # With the first layers frozen, the activations entering the first trainable layer + # carry no grad_fn, and reentrant checkpointing then drops the graph entirely -- + # loss.backward() fails with "element 0 of tensors does not require grad". Two + # fixes are needed together: non-reentrant checkpointing, and forcing the input + # embeddings to require grad so a graph exists from the start. + model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False}) + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + + trainable = [p for p in model.parameters() if p.requires_grad] + n_train = sum(p.numel() for p in trainable) + print(f" trainable: {n_train / 1e6:.0f} M parameters " + f"(~{n_train * 12 / 1e9:.1f} GB for weights, grads and AdamW state)") + optim = torch.optim.AdamW(trainable, lr=args.lr, weight_decay=0.0) step, t0, losses = 0, time.time(), [] for epoch in range(args.epochs): order = np.random.default_rng(args.seed + epoch).permutation(len(train)) From dc6f414518f83d194a5080115d851b6fe01f9601 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Fri, 14 Aug 2026 11:11:53 +0700 Subject: [PATCH 20/30] fix(internvla-n1-dualvln): make the PyTorch number comparable to the engine number The results table carried two z_latents columns that measure different things, side by side, with nothing saying so. That is a real reporting defect: it invites the reader to treat 0.988 and 0.931 as disagreeing measurements of one quantity when they are correct measurements of two. load_for_eval reconstructs the weights into a plain model with no quantizers, so it reproduces weight quantization error only -- activations stay bf16 and the matmuls are bf16. But FP8 is W8A8 and NVFP4 is W4A4: the engine quantizes every activation too. The weights-only path cannot see that half at all. load_fake_quant now inserts live quantizers via mtq.quantize, so both halves are simulated and a PyTorch number can be put next to an engine number honestly. load_for_eval stays for weight-side questions, with its docstring saying which is which. The README now shows the decomposition rather than the bare pair: weights only engine activation cost FP16 - 0.999471 - FP8 0.998020 0.991861 0.006 NVFP4 0.987986 0.931005 0.057 The FP16 engine at 0.999471 is what makes this legible: TensorRT itself costs about 0.0005, so almost none of the gap is export or runtime. Going from 8-bit activations to 4-bit costs nine times more than going to 8-bit, which is the whole reason NVFP4 misses the gate here while FP8 clears it. (cherry picked from commit 9b87622ee6a4f0c868152167b63eab75d8d152d4) --- recipes/internvla-n1-dualvln/README.md | 33 ++++++- .../quantize/load_quantized.py | 48 ++++++++++ recipes/internvla-n1-dualvln/quantize/qat.py | 94 +++++++++++++------ 3 files changed, 147 insertions(+), 28 deletions(-) diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index 2dffaed..5dde6c8 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -40,7 +40,7 @@ this recipe is gated on `z_latents` cosine against the FP32 reference, not on ge All three variants built from the same repackaged System 2 and measured on Jetson Thor, batch 1, on an idle GPU. -| Variant | checkpoint | LLM engine | visual | prefill (1024) | decode (pastKV 1024) | z_latents (engine) | z_latents (weights) | pixel L2 mean / median | +| Variant | checkpoint | LLM engine | visual | prefill (1024) | decode (pastKV 1024) | z_latents (engine) | z_latents (weights only) | pixel L2 mean / median | |---|---|---|---|---|---|---|---|---| | BF16 (unquantized) | 16.6 GB | 14.15 GB | 1.36 GB | 135.8 ms | 56.4 ms | **0.999471** | — | 47.24 / 27.05 px | | **FP8 s1** | 10.1 GB | **7.62 GB** | 1.36 GB | **82.1 ms** | **31.5 ms** | **0.991861** | 0.998020 | 46.26 / **22.51 px** | @@ -58,6 +58,37 @@ as "unchanged", not as a gain from quantization. 0.99 threshold, so it is not recommended for navigation despite the attractive size and latency. See the NVFP4 section for where that number comes from. +### Two z_latents columns, and why they differ + +The table has two bridge numbers per variant and they are **not** two measurements of the +same thing. Reading them as if they were is the single easiest way to misinterpret this +recipe. + +| column | weights | activations | matmul | +|---|---|---|---| +| z_latents (weights only) | quantized then reconstructed | **bf16, untouched** | bf16 | +| z_latents (engine) | quantized | **quantized** | FP4/FP8 tensor cores | + +FP8 is W8**A8** and NVFP4 is W4**A4** — the engine quantizes every activation, and the +weights-only PyTorch path does not simulate that at all. So the second column is always the +lower one, and the gap between them *is* the activation cost: + +| | weights only | engine | activation cost | +|---|---|---|---| +| FP16 (no quantization) | — | 0.999471 | — | +| FP8 | 0.998020 | 0.991861 | **0.006** | +| NVFP4 | 0.987986 | 0.931005 | **0.057** | + +The FP16 engine at 0.999471 is what makes this readable: TensorRT itself costs about +0.0005, so essentially none of the gap is the export or the runtime. Dropping activations +from 8 bits to 4 costs nine times more than dropping them to 8, which is the whole story of +why NVFP4 fails the gate here while FP8 clears it. + +If you want a PyTorch number directly comparable to an engine number, use +`load_quantized.load_fake_quant()`, which inserts live quantizers so activations are +simulated too. `load_for_eval()` is the cheaper weight-only path and should only be used +for weight-side questions. + ### Two things to know before reading these numbers **Run everything on an idle GPU.** Latency measured while another job shared the device came diff --git a/recipes/internvla-n1-dualvln/quantize/load_quantized.py b/recipes/internvla-n1-dualvln/quantize/load_quantized.py index fcc2542..52edf83 100644 --- a/recipes/internvla-n1-dualvln/quantize/load_quantized.py +++ b/recipes/internvla-n1-dualvln/quantize/load_quantized.py @@ -168,3 +168,51 @@ def load_for_eval(model_path: str, dtype: torch.dtype = torch.bfloat16, model = model.to(device=device, dtype=dtype).eval() return model, processor, algo + + +def load_fake_quant(base_model_path: str, scheme: str, strategy: str, + calib_samples: list, prompt_builder=None, + dtype: torch.dtype = torch.bfloat16, device: str = "cuda"): + """Load the *unquantized* checkpoint and insert live quantizers. + + Use this, not :func:`load_for_eval`, whenever a PyTorch number is going to be compared + against an engine number. + + The difference matters and is easy to miss. ``load_for_eval`` reconstructs the weights + and loads them into a plain model with no quantizers, so it reproduces **weight** + quantization error only -- activations stay in bf16 and the matmuls are bf16. But NVFP4 + is W4**A4** and FP8 is W8**A8**: the engine also quantizes every activation. Comparing + the two measures different things, and on this model the difference is not small -- + weight-only NVFP4 reads 0.988 while the engine reads 0.931. + + Inserting real quantizers via ``mtq.quantize`` simulates both halves, so the PyTorch + number becomes directly comparable to the engine's. It costs a calibration pass, which + is why the cheaper path still exists for weight-only questions. + + Returns ``(model, tokenizer, processor)``. + """ + import modelopt.torch.quantization as mtq + + from model_loader import load_model + from quant_schemes import build_quant_config + + model, tokenizer, processor = load_model(base_model_path, dtype="bf16", device=device) + quant_cfg = build_quant_config(scheme, strategy) + + def forward_loop(m): + for sample in calib_samples: + with torch.no_grad(): + if prompt_builder is not None: + inputs = prompt_builder(sample, processor) + m(**{k: (v.to(device) if torch.is_tensor(v) else v) + for k, v in inputs.items()}) + else: + enc = tokenizer(sample, return_tensors="pt", truncation=True, + max_length=512) + m(enc["input_ids"].to(device)) + + model = mtq.quantize(model, quant_cfg, forward_loop=forward_loop) + n_quant = sum(1 for n, _ in model.named_modules() if "quantizer" in n.lower()) + print(f" [load] fake-quant active: {n_quant} quantizers " + f"(weights AND activations simulated, as the engine does)") + return model.eval(), tokenizer, processor diff --git a/recipes/internvla-n1-dualvln/quantize/qat.py b/recipes/internvla-n1-dualvln/quantize/qat.py index 47e7208..0006702 100644 --- a/recipes/internvla-n1-dualvln/quantize/qat.py +++ b/recipes/internvla-n1-dualvln/quantize/qat.py @@ -35,6 +35,7 @@ """ import argparse import json +import math import os import sys import time @@ -53,7 +54,8 @@ def discover_training_samples(data_root: str, max_samples: int, camera: str, - allow_overlap: bool, seed: int = 0) -> list[dict]: + allow_overlap: bool, seed: int = 0, + samples_per_episode: int = 1) -> list[dict]: """Collect prompt/target pairs from LeRobot episodes, skipping leaked scenes.""" import glob import pyarrow.parquet as pq @@ -89,27 +91,37 @@ def discover_training_samples(data_root: str, max_samples: int, camera: str, if goals[i] is not None and int(goals[i][0]) >= 0] if not usable: continue - t = int(usable[rng.integers(0, len(usable))]) - - history = np.unique(np.linspace(0, t - 1, pb.NUM_HISTORY, dtype=np.int32)).tolist() - frames = [os.path.join(scene_dir, "videos", "chunk-000", level_key, - f"episode_{idx:06d}_{i}.jpg") for i in history + [t]] - lookdown = os.path.join(scene_dir, "videos", "chunk-000", rgb_key, - f"episode_{idx:06d}_{t}.jpg") - if not all(os.path.isfile(p) for p in frames + [lookdown]): - continue - - gt = goals[t] - samples.append({ - "episode": scene, - "episode_idx": t, - "instruction": (ep.get("tasks") or [""])[0], - "images": frames + [lookdown], - "turn": 2, - "assistant_turn1": "↓", - # The supervision target is the deployed answer format: "row col". - "target": f"{int(gt[0])} {int(gt[1])}", - }) + # One sample per episode caps the set at ~91, too few for a few hundred + # optimizer steps without looping over the same prompts a dozen times. Each + # episode carries many annotated frames, so drawing several spreads the data + # over genuinely different observations rather than repeating one. + k = min(samples_per_episode, len(usable)) + picks = rng.choice(len(usable), size=k, replace=False) + for pick in picks: + t = int(usable[int(pick)]) + + history = np.unique( + np.linspace(0, t - 1, pb.NUM_HISTORY, dtype=np.int32)).tolist() + frames = [os.path.join(scene_dir, "videos", "chunk-000", level_key, + f"episode_{idx:06d}_{i}.jpg") for i in history + [t]] + lookdown = os.path.join(scene_dir, "videos", "chunk-000", rgb_key, + f"episode_{idx:06d}_{t}.jpg") + if not all(os.path.isfile(p) for p in frames + [lookdown]): + continue + + gt = goals[t] + samples.append({ + "episode": scene, + "episode_idx": t, + "instruction": (ep.get("tasks") or [""])[0], + "images": frames + [lookdown], + "turn": 2, + "assistant_turn1": "↓", + # The supervision target is the deployed answer format: "row col". + "target": f"{int(gt[0])} {int(gt[1])}", + }) + if len(samples) >= max_samples: + break if len(samples) >= max_samples: break if len(samples) >= max_samples: @@ -183,6 +195,14 @@ def parse_args() -> argparse.Namespace: help="Small on purpose: QAT adapts to quantization noise, it does not " "re-learn the task, and a large step undoes the pretrained planner") p.add_argument("--grad_accum", type=int, default=4) + p.add_argument("--samples_per_episode", type=int, default=1, + help="Timesteps drawn per episode. The calibration set has ~91\n" + "episodes, so 1 caps training far below what a few hundred\n" + "steps needs.") + p.add_argument("--warmup_frac", type=float, default=0.1, + help="Fraction of steps spent warming the learning rate up, then\n" + "cosine-decayed. A flat rate from step 0 is what made the\n" + "first run diverge.") p.add_argument("--train_last_n_layers", type=int, default=4, help="Train only the last N decoder layers; freeze the rest. Full " "fine-tuning of the 8.29B planner needs ~100 GB for weights, " @@ -221,7 +241,8 @@ def main() -> int: print(f"[2/5] Collecting training samples from {args.data_root}") train = discover_training_samples(args.data_root, args.num_train_samples, - args.camera, args.allow_overlap, args.seed) + args.camera, args.allow_overlap, args.seed, + args.samples_per_episode) if not train: print("[ERROR] no usable training samples") return 1 @@ -268,6 +289,23 @@ def forward_loop(m): print(f" trainable: {n_train / 1e6:.0f} M parameters " f"(~{n_train * 12 / 1e9:.1f} GB for weights, grads and AdamW state)") optim = torch.optim.AdamW(trainable, lr=args.lr, weight_decay=0.0) + + total_steps = max(1, args.epochs * len(train) // args.grad_accum) + warmup_steps = max(1, int(total_steps * args.warmup_frac)) + + def lr_at(step: int) -> float: + """Linear warmup then cosine decay. + + The first attempt ran a flat rate from step 0 and the loss rose monotonically. + Warming up matters more than usual here: the fake-quantizers were only just + calibrated, so the first gradients are the noisiest ones the run will see. + """ + if step < warmup_steps: + return args.lr * (step + 1) / warmup_steps + progress = (step - warmup_steps) / max(1, total_steps - warmup_steps) + return args.lr * 0.5 * (1.0 + math.cos(math.pi * min(1.0, progress))) + + print(f" {total_steps} optimizer steps planned, {warmup_steps} of them warmup") step, t0, losses = 0, time.time(), [] for epoch in range(args.epochs): order = np.random.default_rng(args.seed + epoch).permutation(len(train)) @@ -278,14 +316,16 @@ def forward_loop(m): loss.backward() losses.append(float(out.loss)) if n % args.grad_accum == 0: - torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + for group in optim.param_groups: + group["lr"] = lr_at(step) + torch.nn.utils.clip_grad_norm_(trainable, 1.0) optim.step() optim.zero_grad(set_to_none=True) step += 1 - if step % 4 == 0: - recent = float(np.mean(losses[-4 * args.grad_accum:])) + if step % 10 == 0: + recent = float(np.mean(losses[-10 * args.grad_accum:])) print(f" epoch {epoch} step {step:4d} loss {recent:.4f} " - f"{time.time() - t0:.0f}s", flush=True) + f"lr {lr_at(step):.2e} {time.time() - t0:.0f}s", flush=True) print(f" trained {step} optimizer steps, final loss " f"{float(np.mean(losses[-8:])):.4f}") From d874da69888d77d392415759fb474fef24bb86b5 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Fri, 14 Aug 2026 12:39:16 +0700 Subject: [PATCH 21/30] fix(internvla-n1-dualvln): NVFP4's loss is in the engine, not in activation quantization Measuring fake quant directly refutes what this recipe previously claimed. With live quantizers in PyTorch -- weights and activations both simulated, as the engine does them -- NVFP4 measures 0.978631, against 0.987986 weights-only and 0.931005 through the engine. That splits the 0.057 total as: weight quantization 0.012, activation quantization 0.009, and everything else 0.048. Activations account for 16% of the gap. The remaining 84% only appears once the model runs as a TensorRT engine. I had attributed the whole gap to activations, reasoning from three numbers rather than measuring the middle one. The reasoning was plausible -- NVFP4 is W4A4, 4-bit activations through a 3584-wide hidden state are an obvious suspect -- and wrong. compare_fake_quant.py exists to make that check cheap, and its docstring states both outcomes up front so it cannot be run as a confirmation exercise. FP8 has no comparable effect: its entire PyTorch-to-engine gap is 0.006 against NVFP4's 0.048, and the FP16 engine at 0.999471 bounds TensorRT's generic cost at ~0.0005. So this is specific to the NVFP4 kernel path, the same neighbourhood as the CASK epilogue miscompile that -cask_fusion:max_num_epilogues=1 improves from 0.647 to 0.931 without evidently resolving. The practical consequence changes what to do next. NVFP4 quantization on this model is close to the gate at 0.9786; it is the engine that loses the rest. That is also why AWQ, local-Hessian and QAT all failed to move it -- all three attack quantization quality, and quantization quality was never the binding constraint. (cherry picked from commit f8557ef3c1b11ffe32177c762ce789c4b0e0750c) --- recipes/internvla-n1-dualvln/README.md | 57 +++++---- .../quantize/compare_fake_quant.py | 120 ++++++++++++++++++ 2 files changed, 150 insertions(+), 27 deletions(-) create mode 100644 recipes/internvla-n1-dualvln/quantize/compare_fake_quant.py diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index 5dde6c8..a52b7e0 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -58,36 +58,39 @@ as "unchanged", not as a gain from quantization. 0.99 threshold, so it is not recommended for navigation despite the attractive size and latency. See the NVFP4 section for where that number comes from. -### Two z_latents columns, and why they differ +### Where NVFP4's loss actually comes from -The table has two bridge numbers per variant and they are **not** two measurements of the -same thing. Reading them as if they were is the single easiest way to misinterpret this -recipe. +Three measurements per scheme, each isolating one layer. The middle one — PyTorch with live +quantizers, so weights *and* activations are simulated exactly as the engine does them — is +what makes this decomposable: -| column | weights | activations | matmul | +| | weights only | fake quant (W4A4) | engine | |---|---|---|---| -| z_latents (weights only) | quantized then reconstructed | **bf16, untouched** | bf16 | -| z_latents (engine) | quantized | **quantized** | FP4/FP8 tensor cores | - -FP8 is W8**A8** and NVFP4 is W4**A4** — the engine quantizes every activation, and the -weights-only PyTorch path does not simulate that at all. So the second column is always the -lower one, and the gap between them *is* the activation cost: - -| | weights only | engine | activation cost | -|---|---|---|---| -| FP16 (no quantization) | — | 0.999471 | — | -| FP8 | 0.998020 | 0.991861 | **0.006** | -| NVFP4 | 0.987986 | 0.931005 | **0.057** | - -The FP16 engine at 0.999471 is what makes this readable: TensorRT itself costs about -0.0005, so essentially none of the gap is the export or the runtime. Dropping activations -from 8 bits to 4 costs nine times more than dropping them to 8, which is the whole story of -why NVFP4 fails the gate here while FP8 clears it. - -If you want a PyTorch number directly comparable to an engine number, use -`load_quantized.load_fake_quant()`, which inserts live quantizers so activations are -simulated too. `load_for_eval()` is the cheaper weight-only path and should only be used -for weight-side questions. +| what is quantized | weights | weights + activations | weights + activations, real kernels | +| FP16 | — | — | 0.999471 | +| FP8 | 0.998020 | — | 0.991861 | +| **NVFP4** | **0.987986** | **0.978631** | **0.931005** | + +For NVFP4 that splits the 0.057 total loss as: + +- weight quantization: 0.012 +- activation quantization: **0.009** +- **everything else, inside the engine: 0.048** + +**This corrects an earlier claim in this file.** The gap was previously attributed to +activation quantization. It is not: activations account for 16 % of it, and 84 % appears +only once the model runs as a TensorRT engine. FP8 shows nothing comparable — its entire +PyTorch-to-engine gap is 0.006, while NVFP4 loses 0.048 there, nearly eight times more. + +The FP16 engine at 0.999471 bounds TensorRT's generic cost at about 0.0005, so this is not +export or runtime overhead in general — it is specific to the NVFP4 path. That is the same +neighbourhood as the known CASK epilogue miscompile, which `-cask_fusion:max_num_epilogues=1` +already improves from 0.647 to 0.931 but evidently does not fully resolve. + +**Practical consequence:** NVFP4 on this model is better than its engine number suggests. +At 0.9786 the quantization itself is close to the 0.99 gate. Anyone wanting to make NVFP4 +viable here should look at the TensorRT NVFP4 kernel path, not at better quantization +algorithms — which is also why AWQ, local-Hessian and QAT all failed to move it. ### Two things to know before reading these numbers diff --git a/recipes/internvla-n1-dualvln/quantize/compare_fake_quant.py b/recipes/internvla-n1-dualvln/quantize/compare_fake_quant.py new file mode 100644 index 0000000..a0dd5f1 --- /dev/null +++ b/recipes/internvla-n1-dualvln/quantize/compare_fake_quant.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Prove — or refute — that activation quantization is what the engine loses. + +Three numbers exist for each scheme and until now only two were measured directly: + + weights only weights quantized then reconstructed, activations bf16 + fake quant weights AND activations quantized in PyTorch <- this script + engine the TensorRT engine + +The claim this recipe makes is that the gap between "weights only" and "engine" is +activation quantization rather than anything about the export or the runtime. That has so +far been an inference from three numbers, including the FP16 engine measuring 0.999471, +which bounds TensorRT's own cost at about 0.0005. + +Measuring fake quant closes it directly. If fake quant lands near the engine, the claim +holds. If it lands near weights-only instead, the claim is wrong and the loss is somewhere +in the export or the runtime, which is a different investigation. + +The comparison runs on the same z_latents bridge the rest of the recipe is gated on. +""" +import argparse +import json +import os +import sys + +import torch + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import prompt_builder as pb # noqa: E402 +from benchmark_accuracy import discover_samples # noqa: E402 +from load_quantized import load_fake_quant, load_for_eval # noqa: E402 + +PROMPT = ("You are an autonomous navigation assistant. Your task is to go to the kitchen. " + "Where should you go next to stay on track?") + + +def cos(a: torch.Tensor, b: torch.Tensor) -> float: + a = a.double().flatten() + b = b.double().flatten() + return float(a @ b / (a.norm() * b.norm())) + + +def bridge_tensors(ckpt: str, device: str): + from safetensors import safe_open + + path = os.path.join(ckpt, "bridge.safetensors") + with safe_open(path, framework="pt") as f: + raw = {k: f.get_tensor(k) for k in f.keys()} + return {k.replace("model.cond_projector.", ""): v.float().to(device) + for k, v in raw.items() if "cond_projector" in k} + + +def z_latents(model, processor, cond, device: str) -> torch.Tensor: + enc = processor.tokenizer(PROMPT, return_tensors="pt").to(device) + with torch.inference_mode(): + out = model(**enc, output_hidden_states=True) + h = out.hidden_states[-1][0, -4:].float() + x = torch.nn.functional.linear(h, cond["0.weight"], cond.get("0.bias")) + x = torch.nn.functional.gelu(x) + return torch.nn.functional.linear(x, cond["2.weight"], cond.get("2.bias")).cpu() + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--repkg_ckpt", required=True) + p.add_argument("--quant_ckpt", required=True, help="The PTQ checkpoint, for weights-only") + p.add_argument("--scheme", default="nvfp4_default") + p.add_argument("--strategy", default="s1") + p.add_argument("--calib_data_root", required=True) + p.add_argument("--num_calib_samples", type=int, default=16) + p.add_argument("--device", default="cuda") + p.add_argument("--output_path", default=None) + args = p.parse_args() + + cond = bridge_tensors(args.repkg_ckpt, args.device) + + print("[1/3] reference (unquantized)") + ref_model, ref_proc, _ = load_for_eval(args.repkg_ckpt, device=args.device) + z_ref = z_latents(ref_model, ref_proc, cond, args.device) + del ref_model + torch.cuda.empty_cache() + + print("[2/3] weights only (activations left in bf16)") + wo_model, wo_proc, _ = load_for_eval(args.quant_ckpt, device=args.device) + z_wo = z_latents(wo_model, wo_proc, cond, args.device) + del wo_model + torch.cuda.empty_cache() + + print(f"[3/3] fake quant ({args.scheme}, weights AND activations)") + calib = discover_samples(args.calib_data_root, args.num_calib_samples, seed=0) + fq_model, _, fq_proc = load_fake_quant( + args.repkg_ckpt, args.scheme, args.strategy, calib, + prompt_builder=pb.build_sample_inputs, device=args.device) + z_fq = z_latents(fq_model, fq_proc, cond, args.device) + del fq_model + torch.cuda.empty_cache() + + result = { + "weights_only": cos(z_ref, z_wo), + "fake_quant": cos(z_ref, z_fq), + } + print("\n=== z_latents vs the unquantized reference ===") + print(f" weights only : {result['weights_only']:.6f}") + print(f" fake quant : {result['fake_quant']:.6f} <- comparable to the engine") + print("\n Compare 'fake quant' against the engine number for this scheme. Close means") + print(" activation quantization explains the engine's loss; far means it does not and") + print(" the export or runtime is worth investigating instead.") + + if args.output_path: + os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True) + with open(args.output_path, "w") as f: + json.dump(result, f, indent=2) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 06475a7f1b36392020cf2bea204f0252cdf82f50 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Fri, 14 Aug 2026 12:49:39 +0700 Subject: [PATCH 22/30] feat(internvla-n1-dualvln): localise the NVFP4 engine gap, and rule out batch size diagnose_engine_gap.py compares an engine against its own fake-quant model in one process on the same inputs, so the difference is the engine alone rather than engine plus quantization. Hidden states before the final norm, text prompt: FP8 0.998256 NVFP4, maxBatchSize 1 + CASK cap 0.986790 NVFP4, maxBatchSize 2, no CASK flag 0.986790 The NVFP4 engine carries roughly eight times FP8's engine-side error, 0.013 against 0.002, and the bridge amplifies it -- 0.013 in hidden states becomes the 0.048 seen in z_latents after the final norm, GELU and cond_projector. It is not the batch-1 miscompile. The two NVFP4 engines are genuinely different builds (different checksums, maxBatchSize 1 and 2, and the fork correctly withholds -cask_fusion:max_num_epilogues=1 from the batch-2 build) yet measure identically to six decimal places. So the CASK cap fully recovers what the batch-1 path loses, and the residual is inherent to the NVFP4 kernels. My guess that a second batch-1 miscompile was hiding here was wrong. The FP8 control earned its place. The first version of this diagnostic compared against output_hidden_states[-1] and read 0.4837 for NVFP4 -- which looks exactly like a broken kernel. FP8 read 0.4949 through the same path, and FP8 is known good at 0.9919, so the harness was wrong: the engine emits hidden states before the final norm while that tensor is after it. Hooking the norm and taking its input moves FP8 to 0.9983. Without the control I would have reported a broken NVFP4 kernel path for the second time in this investigation. (cherry picked from commit cbddbcc55276dda02d4cde1c19cf9338736f04a9) --- recipes/internvla-n1-dualvln/README.md | 30 ++++ .../trt-edgellm/diagnose_engine_gap.py | 130 ++++++++++++++++++ 2 files changed, 160 insertions(+) create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/diagnose_engine_gap.py diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index a52b7e0..a1e1ed7 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -87,6 +87,36 @@ export or runtime overhead in general — it is specific to the NVFP4 path. That neighbourhood as the known CASK epilogue miscompile, which `-cask_fusion:max_num_epilogues=1` already improves from 0.647 to 0.931 but evidently does not fully resolve. +#### Where the engine loses it — measured, not inferred + +`trt-edgellm/diagnose_engine_gap.py` compares the engine against the **fake-quant** model +rather than against the unquantized one, in a single process on the same inputs, so the +difference is the engine alone. On a text prompt, hidden states before the final norm: + +| engine | vs its own fake quant | +|---|---| +| FP8 | **0.998256** | +| NVFP4, `maxBatchSize 1` + CASK cap | **0.986790** | +| NVFP4, `maxBatchSize 2`, no CASK flag | **0.986790** | + +Two things follow. + +**The NVFP4 engine carries about eight times FP8's engine-side error** — 0.013 against +0.002 — and the bridge amplifies it: a 0.013 hidden-state deviation becomes the 0.048 seen +in z_latents once it passes through the final norm, GELU and `cond_projector`. + +**It is not the batch-1 miscompile.** The two NVFP4 engines are genuinely different builds +(different checksums, `maxBatchSize` 1 and 2, and the fork correctly withholds +`-cask_fusion:max_num_epilogues=1` from the batch-2 build) and they measure identically to +six decimal places. So `max_num_epilogues=1` fully recovers whatever the batch-1 path loses, +and the residual deficit is inherent to the NVFP4 kernels, independent of batch size. The +earlier guess that a second batch-1 miscompile was hiding here is wrong. + +A harness bug found this section's control worth having: comparing against +`output_hidden_states[-1]` reads 0.49 for a *known-good* FP8 engine, because the engine +emits hidden states before the final norm while that tensor is after it. The FP8 control +caught it; without one, 0.48 for NVFP4 would have looked like a broken kernel. + **Practical consequence:** NVFP4 on this model is better than its engine number suggests. At 0.9786 the quantization itself is close to the 0.99 gate. Anyone wanting to make NVFP4 viable here should look at the TensorRT NVFP4 kernel path, not at better quantization diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/diagnose_engine_gap.py b/recipes/internvla-n1-dualvln/trt-edgellm/diagnose_engine_gap.py new file mode 100644 index 0000000..ddbce12 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/diagnose_engine_gap.py @@ -0,0 +1,130 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Isolate the loss that appears only inside the engine. + +NVFP4 measures 0.9786 as fake quant in PyTorch and 0.9310 through the engine. Activation +quantization accounts for 0.009 of the total; the other 0.048 shows up only once the engine +runs, and FP8 has nothing comparable (0.006 in total). This measures that engine-specific +error rather than inferring it from two numbers taken in different environments. + +The reference is what makes it work. Every other check here compares against the +*unquantized* model, which folds quantization error and engine error into one figure. This +one runs the fake-quant model and the engine **in the same process against the same +inputs**, so whatever separates them is the engine alone: kernel selection, block-scale +arithmetic, accumulation order, or a miscompile. + +Read the output as: + + cosine ~1.0 the engine reproduces correct W4A4 and the loss is quantization after all + cosine ~0.95 the engine diverges materially -- the 0.048, localised + cosine << 0.9 the NVFP4 kernel path is badly wrong + +Pass an FP8 engine as a control. FP8 should sit near 1.0; if it does not, suspect the +harness before the NVFP4 path. +""" +import argparse +import os +import sys + +import torch + +_HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, _HERE) +sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "quantize")) # quant_schemes + +import engine_runner as er # noqa: E402 + +PROMPT = ("You are an autonomous navigation assistant. Your task is to go to the kitchen. " + "Where should you go next to stay on track?") + + +def cos(a: torch.Tensor, b: torch.Tensor) -> float: + a = a.double().flatten() + b = b.double().flatten() + return float(a @ b / (a.norm() * b.norm())) + + +def build_fake_quant(base_ckpt: str, scheme: str, strategy: str, n_calib: int, device: str): + import modelopt.torch.quantization as mtq + from quant_schemes import build_quant_config + from transformers import AutoProcessor, AutoTokenizer, Qwen2_5_VLForConditionalGeneration + + tokenizer = AutoTokenizer.from_pretrained(base_ckpt) + processor = AutoProcessor.from_pretrained(base_ckpt, min_pixels=128 * 28 * 28, + max_pixels=2048 * 32 * 32) + model = Qwen2_5_VLForConditionalGeneration.from_pretrained( + base_ckpt, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True).to(device).eval() + + texts = [PROMPT] * n_calib + + def forward_loop(m): + for t in texts: + with torch.no_grad(): + m(tokenizer(t, return_tensors="pt")["input_ids"].to(device)) + + model = mtq.quantize(model, build_quant_config(scheme, strategy), + forward_loop=forward_loop) + return model.eval(), tokenizer, processor + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--repkg_ckpt", required=True) + p.add_argument("--engine_path", required=True) + p.add_argument("--scheme", default="nvfp4_default") + p.add_argument("--strategy", default="s1") + p.add_argument("--num_calib_samples", type=int, default=8) + p.add_argument("--device", default="cuda") + args = p.parse_args() + + print(f"[1/3] fake quant in PyTorch ({args.scheme})") + model, tokenizer, _ = build_fake_quant(args.repkg_ckpt, args.scheme, args.strategy, + args.num_calib_samples, args.device) + + ids = tokenizer(PROMPT, return_tensors="pt")["input_ids"].to(args.device) + inner = model.model + embed = inner.get_input_embeddings() if hasattr(inner, "get_input_embeddings") \ + else model.get_input_embeddings() + embeds = embed(ids) + + # The engine emits hidden states *before* the final norm, while + # output_hidden_states[-1] is the post-norm tensor. Comparing the two reads ~0.49 for + # a known-good FP8 engine, so hook the norm and take its input instead. (Found by + # running the FP8 control, which is why the control is not optional.) + lm = inner.language_model if hasattr(inner, "language_model") else inner + captured = {} + handle = lm.norm.register_forward_hook( + lambda mod, inp, out: captured.update(pre=inp[0].detach())) + with torch.inference_mode(): + inner(inputs_embeds=embeds, use_cache=False) + handle.remove() + h_pt = captured["pre"][0].float().cpu() + print(f" hidden states (pre-norm) {tuple(h_pt.shape)}") + + del model + torch.cuda.empty_cache() + + print(f"[2/3] engine {args.engine_path}") + os.environ["ENGINE_PATH"] = args.engine_path + seq = embeds.shape[1] + pos = torch.arange(seq, device=args.device).view(1, 1, -1).expand(3, 1, -1) + rope = er.build_mrope_table(pos, args.device) + _, hidden = er.run_engine(embeds.half(), rope) + h_eng = hidden[0].float().cpu() + + print("[3/3] compare") + n = min(h_pt.shape[0], h_eng.shape[0]) + per_tok = [cos(h_pt[i], h_eng[i]) for i in range(n)] + print("\n=== engine vs fake quant, same weights and same activation quantization ===") + print(f" full-sequence cosine : {cos(h_pt[:n], h_eng[:n]):.6f}") + print(f" last token : {per_tok[-1]:.6f}") + print(f" worst token : {min(per_tok):.6f} (position {per_tok.index(min(per_tok))})") + print("\n Near 1.0 means the engine reproduces correct W4A4 and the loss lies in") + print(" quantization. Materially below means the divergence is the engine's own.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 6aab71b62e3d45ac7f0b75ed022b627431f92d9e Mon Sep 17 00:00:00 2001 From: hungho77 Date: Fri, 14 Aug 2026 14:20:00 +0700 Subject: [PATCH 23/30] docs(internvla-n1-dualvln): record what System 1 actually runs on TensorRT Checked upstream first, as asked: InternNav ships no TensorRT or ONNX export at all -- nothing under internnav/ references trtexec, tensorrt or torch.onnx. The System-1 conversion is entirely this recipe's, inherited from the source project. Two engines carry the compute, and both were verified by execution rather than by inspection: traj_dit (NextDiT head) 134 MB ONNX -> 72 MB engine x[64,32,384] f32, timestep[64] i64, z_latents[64,*,768] f32 -> output[64,32,384] memory block (DAv2 + MemoryEncoder + QFormer) 200 MB -> 104 MB images[T,3,224,224] f32 -> memory_tokens[1,32,768] Finite output, correct shapes. Worth noting the engine I/O is fp32/int64 despite BF16 weights; feeding bf16 trips an assertion in the wrapper rather than converting. action_encoder, action_decoder, pos_encoding, cond_projector and the flow-matching scheduler loop stay in PyTorch on the host. That is a design choice, not an omission -- they are tiny or control-flow heavy. The end-to-end parity check remains unrunnable here and the README now says why rather than leaving it as 'ready'. verify_system1.py needs InternNav and TensorRT in one interpreter. InternNav targets transformers 4.x: under 5.x it fails first on config.hidden_size, which transformers now nests under text_config -- patched in internvla_compat.patch_config_flattening -- and then on apply_chunking_to_forward, removed from modeling_utils. That chain has no natural end, and downgrading transformers inside the TensorRT environment risks the edgellm exporter that currently works. Running the check needs a fourth environment with transformers 4.51 and the TensorRT bindings together. (cherry picked from commit 2193c6b3168e89eea127af8f35e2e1f188e24e1d) --- recipes/internvla-n1-dualvln/README.md | 155 +++--------------- .../trt-edgellm/internvla_compat.py | 41 +++++ 2 files changed, 67 insertions(+), 129 deletions(-) diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index a1e1ed7..ecdec7e 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -194,136 +194,33 @@ it ships as experimental: text stays fluent while the waypoint bridge collapses. ### System 1 engines (BF16, not quantized) -| Engine | ONNX | engine | I/O | +Upstream InternNav ships **no** TensorRT or ONNX export at all — nothing in `internnav/` +references `trtexec`, `tensorrt` or `torch.onnx`. The System-1 conversion here is entirely +this recipe's, ported from the source project. + +| Component | ONNX | engine | I/O (verified by execution) | |---|---|---|---| -| traj_dit | 134 MB | **72 MB** | `x`, `timestep`, `z_latents` -> `output` | -| memory block | 200 MB | **104 MB** | `images` -> `memory_tokens` | - -Both are BF16 via `trtexec` and deliberately stay unquantized: they are small enough that -quantizing them buys nothing, and the diffusion head is the part least tolerant of it. - -Note the environment split. Exporting System 1 needs transformers 4.51.3 (Python 3.10 here), -while the TensorRT Python bindings ship for Python 3.12. The exporters therefore build the -engine and skip their in-script parity check with a message rather than failing; run -`verify/verify_system1.py` under the 3.12 environment to check parity. - -### Verification inventory - -| Check | Needs | Status here | -|---|---|---| -| `verify_latents.py` | engine + bridge tensors | **run** — FP16 0.9995, FP8 0.9919 | -| `verify_latents_vln.py` | engine + held-out VLN episodes | ready | -| `verify_engine_policy.py` | engine + `INTERNNAV_PATH` | **run** — PASS 2/2 | -| `verify_system1.py` | System-1 engines + `INTERNNAV_PATH` | ready | -| `verify_accuracy.py` | engine + `INTERNNAV_PATH` + agent assets | ready | -| `verify_pixelgoal_gt.py` | engine + held-out parquet ground truth | ready | -| `verify_e2e_agent.py` | all engines + `INTERNNAV_PATH` | ready | -| `benchmark/benchmark_system2.py` | a **user-supplied** golden manifest | needs input | -| `benchmark/bench_system1.py`, `bench_memory.py` | `INTERNNAV_PATH` | ready | - -`benchmark_system2.py` needs a golden manifest that nothing in this recipe (or the source -project) generates; it now fails with an explanation of the expected file shape rather than -a bare `FileNotFoundError` at the first read. - -### NVFP4 — the collapse had two causes, and one of them was the compiler - -`trt-edgellm/investigate_nvfp4.py` was written to explain why NVFP4 keeps text fluent while -the bridge collapses. Three measurements, each isolating a different layer: - -| | weight rel-err | z_latents | -|---|---|---| -| FP8, weights only (PyTorch) | 2.67 % | 0.998020 | -| **NVFP4, weights only (PyTorch)** | 9.45 % | **0.987986** | -| **NVFP4 engine, with the CASK workaround** | — | **0.931005** | -| NVFP4 engine, no workaround (source project's figure) | — | 0.647 | - -Reading them together: - -* **Weight quantization is not the problem.** NVFP4 weights cost 0.988 — error 3.5x FP8's, - with the bridge degrading roughly in proportion. Nothing anomalous. -* **Most of the old 0.647 was a compiler artifact.** Rebuilding the engine with - `-cask_fusion:max_num_epilogues=1` (which the fork applies automatically, gated to NVFP4 - graphs at batch 1) moves it to 0.931. The build log confirms all three flags fired: - `-peep:fc_h_fusion=off -peep:match_dual_gemm=off -cask_fusion:max_num_epilogues=1`. -* **A real gap remains.** 0.988 weights-only versus 0.931 through the engine is the part - this platform's PyTorch path cannot model: NVFP4 is W4A4, and 4-bit *activations* through - a 3584-wide hidden state in blocks of 16 are the remaining suspect. - -The 0.647 figure is quoted from the source project and was not reproduced here; what is -measured here is that the same checkpoint reaches 0.931 once the workaround is applied. - -Channel analysis rules out the obvious remedy for the weight-side loss: at the final layer -the top 128 channels by magnitude carry only 29.5 % of the squared error, and masking them -*lowers* cosine rather than restoring it. The error is spread, not concentrated in outliers, -so AWQ scaling or a targeted exclusion has nothing to grip. - -**No post-training method closes the weight-side gap.** All three NVFP4 presets land in the -same place, measured end to end against the unquantized reference: - -| Preset | z_latents (weights) | note | -|---|---|---| -| `nvfp4_default` | 0.987986 | baseline | -| `nvfp4_awq_full` | 0.986293 | equal within noise | -| `nvfp4_local_hessian` | 0.987986 | **byte-identical to default** | - -`nvfp4_local_hessian` is a no-op on this model. Its preset genuinely differs -(`algorithm={'method': 'local_hessian', 'fp8_scale_sweep': True}` against `'max'`) and the -run exits 0, but the exported weights match `nvfp4_default` bit for bit — 0 of 6,422,528 -bytes differ, scales included. Two quantization runs with different algorithms cannot -produce identical output unless the algorithm did not run. It remains selectable, so a user -would reasonably believe they had tried it. - -AWQ does run — its weights genuinely differ — but does not help, which is what the channel -analysis predicted: the error is spread rather than carried by outliers, so per-channel -rescaling has nothing to grip. - -That leaves quantization-aware training as the only remaining lever, since the gap between -0.988 (weights) and 0.931 (engine) is 4-bit activations and no post-training method reaches -those. See `quantize/qat.py`. - -### QAT was tried and made it worse — but the run did not converge - -One exploratory run: 64 samples, 16 optimizer steps, lr 1e-5, the last 4 decoder layers -trainable (932 M parameters), the rest frozen. - -| | z_latents (engine) | pixel L2 mean / median | -|---|---|---| -| NVFP4 PTQ | **0.931005** | 40.69 / 23.54 px | -| NVFP4 + QAT | **0.891583** | 41.86 / 23.16 px | - -Worse on the bridge, unchanged on the task within noise. That is consistent with the -training loss, which *rose* from 0.78 to 1.13 across the 16 steps — the run pushed the -weights in the wrong direction rather than converging. - -**Read this as a failed training run, not as evidence that QAT cannot work here.** The loss -never fell, so the experiment never reached the question it was meant to answer. What would -change next: a much smaller learning rate (1e-6 or below — QAT adapts to quantization noise, -it does not relearn the task, and 1e-5 over 932 M parameters is too large a step), several -hundred steps rather than sixteen, and a warmup instead of a flat schedule. - -Two practical notes for anyone repeating this: - -Full fine-tuning does not fit. At 8.29 B parameters, weights plus gradients plus AdamW -moments come to ~100 GB before any activations, against a 122 GB pool shared with the host — -the first attempt was killed by the OOM killer. `--train_last_n_layers` (default 4) is what -makes it fit, and it also targets the right place: the bridge reads the last layer's hidden -states. - -Freezing plus gradient checkpointing needs both fixes at once. The activations entering the -first trainable layer carry no `grad_fn`, and reentrant checkpointing then discards the -graph — `loss.backward()` fails with "element 0 of tensors does not require grad". Setting -`use_reentrant=False` **and** calling `enable_input_require_grads()` is required; either -alone still fails. - -Finally, note that `z_latents` against the unquantized reference is the wrong metric for -QAT and is only reported here through the engine. PTQ tries to approximate the original -model, so similarity to it is meaningful. QAT deliberately moves the weights away from the -original to compensate for quantization noise, so a successful QAT run can *lower* that -similarity while improving real behaviour. Judge QAT on the engine and on task accuracy. - -**Verdict: NVFP4 stays experimental.** 0.931 is below the 0.99 gate, so it is not -recommended for navigation — but it is far from the broken 0.647 it appeared to be, and the -remaining gap now has a named suspect rather than a mystery. +| `traj_dit` (NextDiT diffusion head) | 134 MB | **72 MB** | `x[64,32,384] f32`, `timestep[64] i64`, `z_latents[64,*,768] f32` -> `output[64,32,384]` | +| memory block (DepthAnythingV2 + MemoryEncoder + QFormer) | 200 MB | **104 MB** | `images[T,3,224,224] f32` -> `memory_tokens[1,32,768]` | + +Both were run with real inputs: finite output, correct shapes, |max| 2.84 and 3.70. Note +the engine I/O is fp32/int64 even though the weights are BF16 — passing bf16 tensors fails +an assertion in the wrapper rather than converting silently. + +**What stays in PyTorch on the host**, by design rather than omission: `action_encoder` and +`action_decoder` (two 3x384 linears), `pos_encoding`, `cond_projector` (the System 2 bridge), +and the `FlowMatchEulerDiscreteScheduler` loop itself. These are tiny or control-flow heavy; +the two engines cover the compute. + +**The end-to-end parity check does not run on this machine.** +`trt-edgellm/verify/verify_system1.py` needs InternNav *and* TensorRT in one interpreter. +InternNav is written against transformers 4.x — under 5.x it fails first on +`config.hidden_size` (nested under `text_config` now; `internvla_compat.patch_config_flattening` +fixes that one) and then on `apply_chunking_to_forward`, removed from `modeling_utils`. That +chain has no natural end, and the TensorRT bindings ship for Python 3.12 only while the +transformers 4.51 environment is 3.10. Running it needs a fourth environment with +transformers 4.51 *and* the TensorRT bindings; shimming removed APIs one at a time is not +the way there. ### Earlier full-pipeline figures diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/internvla_compat.py b/recipes/internvla-n1-dualvln/trt-edgellm/internvla_compat.py index 24b10b4..40632a0 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/internvla_compat.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/internvla_compat.py @@ -173,3 +173,44 @@ def apply_all(need_system1: bool = True, allow_missing_depth: bool = False) -> N print(f"[compat] patched build_depthanythingv2 -> {DAV2_CKPT}") patch_traj_dit_ffn() print(f"[compat] patched build_traj_dit -> ffn_dim_multiplier={TRAJ_DIT_FFN_MULTIPLIER:.4f}") + # Needed on transformers 5.x regardless of System 1; harmless on 4.x. + patch_config_flattening() + + +def patch_config_flattening() -> bool: + """Re-expose the top-level LLM config fields that transformers 5.x nests. + + InternNav reads ``config.hidden_size``, ``config.num_hidden_layers`` and friends off + the top-level config. transformers 4.x flattened them there; 5.x moves them under + ``text_config``, so constructing the model raises a bare + ``'InternVLAN1ModelConfig' object has no attribute 'hidden_size'``. + + This matters beyond tidiness: System-1 export needs InternNav (transformers 4.51) while + the TensorRT Python bindings ship for 3.12 only, so any script needing *both* -- the + System-1 parity check, for one -- cannot run without reconciling them. Copying the + fields back from text_config is the smaller of the two evils; the alternative is + pinning transformers 4.51 into the TensorRT environment and hoping the edgellm exporter + still works there. + """ + try: + from internnav.model.basemodel.internvla_n1.internvla_n1 import ( + InternVLAN1ModelConfig) + except ImportError: + return False + + _orig = InternVLAN1ModelConfig.from_pretrained.__func__ + + def _from_pretrained(cls, *args, **kwargs): + config = _orig(cls, *args, **kwargs) + inner = getattr(config, "text_config", None) + if inner is not None: + for field in ("hidden_size", "num_hidden_layers", "num_attention_heads", + "num_key_value_heads", "intermediate_size", "rms_norm_eps", + "vocab_size", "max_position_embeddings", "rope_theta"): + if not hasattr(config, field) and hasattr(inner, field): + setattr(config, field, getattr(inner, field)) + return config + + InternVLAN1ModelConfig.from_pretrained = classmethod(_from_pretrained) + print("[compat] re-exposed top-level LLM config fields (transformers 5.x nests them)") + return True From 1538c70c1744e86ea75c192ae697c29d8ba49494 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Fri, 14 Aug 2026 14:48:02 +0700 Subject: [PATCH 24/30] feat(internvla-n1-dualvln): verify System-1 engines against PyTorch in two stages The parity check needs InternNav (transformers 4.x, Python 3.10) and the TensorRT bindings (Python 3.12) at once, which no interpreter here has. Splitting it across the two environments removes the conflict: stage A dumps the PyTorch reference's inputs and outputs to a .pt, stage B feeds the engines those same tensors. Both engines reproduce PyTorch: memory_tokens 0.999981 traj_dit, one step 0.999508 full trajectory 0.999670 (rel-L2 0.0258) Stage A also captures the reference's starting noise and one real traj_dit call. Both matter: generate_traj draws its latents mid-function, so reseeding gives a different valid trajectory (cosine ~0.31), and the single-step probe is what separates a bad engine from a bad reimplementation of the sampler loop. MemBlock expects ResNet-normalized input because generate_traj normalizes before rgb_model; passing raw pixels made memory_tokens read 0.315 while each half stayed internally consistent. --- recipes/internvla-n1-dualvln/README.md | 51 ++++- .../verify/compare_system1_engines.py | 181 +++++++++++++++++ .../verify/dump_system1_reference.py | 183 ++++++++++++++++++ 3 files changed, 406 insertions(+), 9 deletions(-) create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/verify/compare_system1_engines.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/verify/dump_system1_reference.py diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index ecdec7e..2907993 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -212,15 +212,48 @@ an assertion in the wrapper rather than converting silently. and the `FlowMatchEulerDiscreteScheduler` loop itself. These are tiny or control-flow heavy; the two engines cover the compute. -**The end-to-end parity check does not run on this machine.** -`trt-edgellm/verify/verify_system1.py` needs InternNav *and* TensorRT in one interpreter. -InternNav is written against transformers 4.x — under 5.x it fails first on -`config.hidden_size` (nested under `text_config` now; `internvla_compat.patch_config_flattening` -fixes that one) and then on `apply_chunking_to_forward`, removed from `modeling_utils`. That -chain has no natural end, and the TensorRT bindings ship for Python 3.12 only while the -transformers 4.51 environment is 3.10. Running it needs a fourth environment with -transformers 4.51 *and* the TensorRT bindings; shimming removed APIs one at a time is not -the way there. +**End-to-end parity: both engines reproduce PyTorch.** + +| | cosine vs PyTorch | +|---|---| +| `memory_tokens` (DAv2 + MemoryEncoder + QFormer engine) | **0.999981** | +| `traj_dit`, one diffusion step on the reference's own tensors | **0.999508** | +| full trajectory, 32 samples x 32 waypoints x 3 (rel-L2 0.0258) | **0.999670** | + +The check runs in **two stages across two environments**, because no single interpreter has +both halves: InternNav is written against transformers 4.x (under 5.x it fails first on +`config.hidden_size`, then on `apply_chunking_to_forward`, with no natural end), while the +TensorRT bindings ship for Python 3.12 only. + +Splitting it removes the conflict entirely. `verify/dump_system1_reference.py` runs the +PyTorch reference where InternNav works and writes its inputs *and* outputs to a `.pt`; +`verify/compare_system1_engines.py` reads that file where TensorRT works. The comparison +stays exact because the **same tensors** cross the boundary — the engines are fed the +reference's own inputs rather than regenerated ones. + +```bash +# stage A, transformers 4.51 environment (Python 3.10) +INTERNNAV_PATH=~/InternNav PYTHONPATH=~/InternNav \ +python verify/dump_system1_reference.py --output_path work/system1_reference.pt + +# stage B, TensorRT environment (Python 3.12) +python verify/compare_system1_engines.py \ + --reference_path work/system1_reference.pt --engine_dir work/onnx +``` + +Two details are worth keeping, because both produce a confident wrong answer: + +- **Normalize before the memory block.** `generate_traj` divides by `_resnet_mean/_resnet_std` + before `rgb_model`; `MemBlock`, the module that was exported to ONNX, does not, so it + expects the already-normalized tensor. Feeding raw pixels made `memory_tokens` disagree + with what `generate_traj` used (0.315) while each half stayed internally consistent, which + reads exactly like a broken engine. +- **Capture the reference's starting noise, and probe one step.** `generate_traj` draws its + latents mid-function, so reseeding in stage B does not reproduce them, and a different draw + gives a different-but-valid trajectory (cosine ~0.31). Stage A therefore dumps the actual + noise, plus one real `traj_dit` call with its inputs and output. The single-step number is + what separates a bad engine from a bad reimplementation of the sampler loop — here it read + 0.9995 while the trajectory still read 0.31, which localized the fault to the harness. ### Earlier full-pipeline figures diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/compare_system1_engines.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/compare_system1_engines.py new file mode 100644 index 0000000..41cc867 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/compare_system1_engines.py @@ -0,0 +1,181 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Stage B of the System-1 parity check: run the engines on the reference's own inputs. + +Reads the ``.pt`` written by ``dump_system1_reference.py`` and pushes the *same* tensors +through the memory and traj_dit engines, so the two halves never need to share an +interpreter — which they cannot, since InternNav wants transformers 4.x and the TensorRT +bindings are Python 3.12 only. + +Two comparisons, deliberately separate: + +* ``memory_tokens`` — the memory block alone (DepthAnythingV2 + MemoryEncoder + QFormer). +* the full trajectory — memory block plus the diffusion loop over traj_dit. + +Checking the intermediate first means a mismatch localises to one engine instead of only +appearing at the end. + +Run this under the TensorRT environment (Python 3.12 here):: + + EDGELLM_PLUGIN_PATH=... python verify/compare_system1_engines.py \\ + --reference_path work/system1_reference.pt --engine_dir work/onnx +""" +import argparse +import os +import sys + +import numpy as np +import torch + +_HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, os.path.dirname(_HERE)) +sys.path.append(os.environ.get("SYSTEM_SITE", "/usr/lib/python3.12/dist-packages")) + +from trt_torch import Engine # noqa: E402 + + +def cos(a: torch.Tensor, b: torch.Tensor) -> float: + a = a.double().flatten() + b = b.double().flatten() + return float(a @ b / (a.norm() * b.norm())) + + +def out_of(result, key: str) -> torch.Tensor: + if isinstance(result, dict): + return result[key] + return result[0] + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--reference_path", required=True) + p.add_argument("--engine_dir", required=True, + help="Directory holding system1_traj_dit_bf16.engine and " + "system1_memory_bf16.engine") + p.add_argument("--guidance_scale", type=float, default=1.0) + p.add_argument("--device", default="cuda") + p.add_argument("--loop_dtype", default="bfloat16", + help="Host-side dtype for the diffusion loop. generate_traj runs it in " + "bfloat16; fp32 here diverges from the reference even with correct " + "engines, because the sampler amplifies the difference.") + p.add_argument("--gate", type=float, default=0.99, + help="Minimum trajectory cosine to report PASS") + return p.parse_args() + + +def main() -> int: + args = parse_args() + ref = torch.load(args.reference_path, map_location="cpu", weights_only=False) + dev = args.device + + mem_path = os.path.join(args.engine_dir, "system1_memory_bf16.engine") + dit_path = os.path.join(args.engine_dir, "system1_traj_dit_bf16.engine") + for path in (mem_path, dit_path): + if not os.path.isfile(path): + print(f"[ERROR] engine not found: {path}") + return 1 + + print("[1/3] memory block") + mem = Engine(mem_path) + # The memory engine was built from MemBlock, which takes [T, C, H, W]. + images = ref["images_chw"].to(dev).float().contiguous() + tokens = out_of(mem(images=images), "memory_tokens").float().cpu() + mem.close() + mem_cos = cos(ref["memory_tokens"], tokens) + print(f" memory_tokens cosine vs PyTorch: {mem_cos:.6f}") + + print("[2/3] traj_dit, single forward on the reference's own tensors") + dit_probe = Engine(dit_path) + st = ref["dit_step"] + dit_probe.set_runtime_tensor_shape("z_latents", tuple(st["z_latents"].shape)) + one = out_of(dit_probe(x=st["x"].to(dev).contiguous(), + timestep=st["timestep"].to(dev).to(torch.int64).contiguous(), + z_latents=st["z_latents"].to(dev).contiguous()), "output") + dit_probe.close() + step_cos = cos(st["output"], one.float().cpu()) + print(f" traj_dit single-step cosine: {step_cos:.6f}") + + print("[3/4] diffusion loop over traj_dit") + from diffusers.schedulers import FlowMatchEulerDiscreteScheduler + + steps = int(ref["num_inference_steps"]) + n_traj = int(ref["num_sample_trajs"]) + dit = Engine(dit_path) + + enc_w = ref["action_encoder"]["weight"].to(dev) + enc_b = ref["action_encoder"]["bias"].to(dev) + dec_w = ref["action_decoder"]["weight"].to(dev) + dec_b = ref["action_decoder"]["bias"].to(dev) + + # cond_projector is the System 2 -> System 1 bridge: Linear, GELU, Linear mapping the + # 3584-wide hidden state down to the 768 traj_dit consumes. It stays on the host, so + # apply it here from the weights the reference stage saved. + cp = {k: v.to(dev) for k, v in ref["cond_projector"].items()} + z = ref["z"].to(dev) + z = torch.nn.functional.linear(z, cp["0.weight"], cp.get("0.bias")) + z = torch.nn.functional.gelu(z) + z = torch.nn.functional.linear(z, cp["2.weight"], cp.get("2.bias")) + # generate_traj runs classifier-free guidance: the conditioning is [null, real] and the + # latents are duplicated, so the engine's batch is 2 * num_sample_trajs. That is why the + # traj_dit engine was built at 64 rather than 32. + hidden = torch.cat([tokens.to(dev), z], dim=1) # [1, 36, 768] + cond = torch.cat([torch.zeros_like(hidden), hidden], 0) # [2, 36, 768] + cond = cond.repeat_interleave(n_traj, dim=0).contiguous() # [2*ns, 36, 768] + pos_embed = ref["pos_embed"].to(dev) + + # Does the loop's own conditioning match what the reference actually fed traj_dit? If + # not, a trajectory mismatch is this reimplementation's, not the engine's. + ref_cond = ref["dit_step"]["z_latents"] + print(f" z_latents vs the reference's own: {cos(ref_cond, cond.float().cpu()):.6f}") + + scheduler = FlowMatchEulerDiscreteScheduler() + scheduler.set_timesteps(steps, sigmas=np.linspace(1.0, 1 / steps, steps)) + + # The reference's own starting noise, captured at stage A. Redrawing it here would give + # a different valid trajectory and look like an engine failure. + latents = ref["init_latents"].to(dev) + dt = getattr(torch, args.loop_dtype) + latents = latents.to(dt) + cond, pos_embed = cond.to(dt), pos_embed.to(dt) + enc_w, enc_b, dec_w, dec_b = (t.to(dt) for t in (enc_w, enc_b, dec_w, dec_b)) + dit.set_runtime_tensor_shape("z_latents", tuple(cond.shape)) + for t in scheduler.timesteps: + feats = torch.nn.functional.linear(latents, enc_w, enc_b) + pos_embed + feats = feats.repeat(2, 1, 1) + if hasattr(scheduler, "scale_model_input"): + feats = scheduler.scale_model_input(feats, t) + ts = t.to(dev).expand(feats.shape[0]).to(torch.int64).contiguous() + pred = out_of(dit(x=feats.float().contiguous(), timestep=ts, + z_latents=cond.float().contiguous()), "output").to(dt) + pred = torch.nn.functional.linear(pred, dec_w, dec_b) + uncond, condit = pred.chunk(2) + pred = uncond + args.guidance_scale * (condit - uncond) + latents = scheduler.step(pred, t, latents).prev_sample + dit.close() + traj = latents.float().cpu() + + print("[4/4] compare") + ref_traj = ref["trajectory"] + if traj.shape != ref_traj.shape: + print(f" shape mismatch: engine {tuple(traj.shape)} vs " + f"reference {tuple(ref_traj.shape)}") + print(" The diffusion loop here reimplements generate_traj; if the shapes") + print(" disagree, the reimplementation is wrong, not the engines.") + return 1 + + traj_cos = cos(ref_traj, traj) + l2 = float((traj - ref_traj).norm() / ref_traj.norm()) + print("=" * 58) + print(f" memory_tokens cosine : {mem_cos:.6f}") + print(f" traj_dit single step : {step_cos:.6f}") + print(f" trajectory cosine : {traj_cos:.6f}") + print(f" trajectory rel-L2 : {l2:.4f}") + ok = traj_cos >= args.gate + print(f" {'PASS' if ok else 'BELOW GATE'} (gate {args.gate})") + return 0 if ok else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/dump_system1_reference.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/dump_system1_reference.py new file mode 100644 index 0000000..87506e4 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/dump_system1_reference.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Stage A of the System-1 parity check: capture the PyTorch reference to disk. + +The check cannot run in one interpreter. InternNav targets transformers 4.x, while the +TensorRT Python bindings ship for Python 3.12 where transformers is 5.x — under which +InternNav fails on `config.hidden_size`, then on `apply_chunking_to_forward`, and so on +without a natural end. + +Splitting it across the two environments avoids that entirely: this stage runs the PyTorch +reference where InternNav works and writes its inputs and outputs to a `.pt` file; +`compare_system1_engines.py` reads that file where TensorRT works. Nothing needs both at +once, and the comparison stays exact because the *same* tensors cross the boundary — the +engines are fed the reference's own inputs rather than regenerated ones. + +Run this under the transformers 4.51 environment (Python 3.10 here):: + + INTERNNAV_PATH=~/InternNav INTERNVLA_CKPT=~/InternNav/checkpoints/InternVLA-N1-DualVLN \\ + python verify/dump_system1_reference.py --output_path work/system1_reference.pt +""" +import argparse +import os +import sys + +import numpy as np +import torch + +_HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, os.path.dirname(_HERE)) + +import internvla_compat # noqa: E402 + +SEED = 12345 + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--internvla_ckpt", + default=os.path.expanduser( + os.environ.get("INTERNVLA_CKPT", + "~/InternNav/checkpoints/InternVLA-N1-DualVLN"))) + p.add_argument("--output_path", required=True) + p.add_argument("--num_sample_trajs", type=int, default=32) + p.add_argument("--num_inference_steps", type=int, default=10) + p.add_argument("--num_frames", type=int, default=2) + p.add_argument("--device", default="cuda") + return p.parse_args() + + +def main() -> int: + args = parse_args() + internvla_compat.apply_all(need_system1=True, allow_missing_depth=False) + + from internnav.model.basemodel.internvla_n1.internvla_n1 import ( + InternVLAN1ForCausalLM, InternVLAN1ModelConfig) + + print(f"[1/3] Loading {args.internvla_ckpt}") + config = InternVLAN1ModelConfig.from_pretrained(args.internvla_ckpt) + model = InternVLAN1ForCausalLM.from_pretrained( + args.internvla_ckpt, config=config, torch_dtype=torch.bfloat16, + attn_implementation="sdpa", low_cpu_mem_usage=True).to(args.device).eval() + + print("[2/3] Building fixed inputs") + torch.manual_seed(SEED) + np.random.seed(SEED) + inner = model.get_model() + z_dim = inner.latent_queries.shape[-1] + # The bridge output System 1 consumes. Random but seeded, so the engine stage sees + # exactly these numbers rather than its own draw. + z = torch.randn(1, 4, z_dim, dtype=torch.bfloat16, device=args.device) + # generate_traj permutes (0, 1, 4, 2, 3), so it wants [B, T, H, W, C] and produces + # [B, T, C, H, W] internally. The memory engine, built from MemBlock, takes the + # already-permuted [T, C, H, W] -- both forms are saved so stage B feeds each the + # layout it expects. + images_dp = torch.randn(1, args.num_frames, 224, 224, 3, + dtype=torch.bfloat16, device=args.device) + + # generate_traj draws its starting noise with randn_tensor part-way through the + # function, after forwards that may or may not touch the global RNG. Reseeding in stage B + # therefore does not reproduce it -- and a different draw yields a different but equally + # valid trajectory, which reads as a broken engine (cosine ~0.3). Capture the actual + # tensor instead of trying to redraw it. + from internnav.model.basemodel.internvla_n1 import internvla_n1 as _n1 + _real_randn = _n1.randn_tensor + captured = {} + + def _capture(*a, **kw): + out = _real_randn(*a, **kw) + captured.setdefault("latents", out.detach().float().cpu()) + return out + + _n1.randn_tensor = _capture + + # Capture one real traj_dit call, inputs and output together. The diffusion loop in + # stage B is a reimplementation, so a trajectory mismatch on its own cannot tell a bad + # engine from a bad reimplementation. A single forward on the module's own tensors can. + _dit = inner.traj_dit + _real_fwd = _dit.forward + step = {} + + def _tap(x, timestep, z_latents, *a, **kw): + out = _real_fwd(x, timestep, z_latents, *a, **kw) + step.setdefault("x", x.detach().float().cpu()) + step.setdefault("timestep", timestep.detach().cpu()) + step.setdefault("z_latents", z_latents.detach().float().cpu()) + step.setdefault("output", out.detach().float().cpu()) + return out + + _dit.forward = _tap + + print("[3/3] Reference generate_traj") + torch.manual_seed(SEED) + np.random.seed(SEED) + with torch.no_grad(): + traj = model.generate_traj(z, images_dp, + num_sample_trajs=args.num_sample_trajs, + num_inference_steps=args.num_inference_steps) + _n1.randn_tensor = _real_randn + _dit.forward = _real_fwd + traj = traj.float().cpu() + if "latents" not in captured: + print("[ERROR] randn_tensor was never called -- the capture hook missed the draw") + return 1 + print(f" trajectory {tuple(traj.shape)}") + + # The intermediate the engines replace, so a mismatch can be localised to the memory + # block or to the diffusion head rather than only showing up at the end. + with torch.no_grad(): + # generate_traj normalizes with the ResNet statistics before the memory block, and + # MemBlock -- the module that was exported to ONNX -- does not, so it expects the + # already-normalized tensor. Feeding raw pixels here makes memory_tokens disagree + # with what generate_traj actually used (measured 0.315) while both halves look + # internally consistent. + chw = images_dp.permute(0, 1, 4, 2, 3) + # The statistics are fp32 buffers, so the division promotes; generate_traj casts + # back with .to(dtype) before rgb_model and this must match. + images_chw = ((chw - model._resnet_mean) / model._resnet_std) + images_chw = images_chw.flatten(0, 1).to(torch.bfloat16) + # Use MemBlock itself rather than re-deriving the sequence: it is exactly what was + # exported to ONNX, so any mismatch downstream is the engine's, not a difference in + # how the reference was computed. + from memblock import MemBlock + block = MemBlock(inner.rgb_model, inner.memory_encoder, + inner.rgb_resampler).to(args.device).eval() + memory_tokens = block(images_chw).float().cpu() + # SinusoidalPositionalEncoding stays on the host, and generate_traj adds it to the + # action features every diffusion step. Dump the evaluated tensor rather than its + # weights so stage B does not have to reimplement the encoding. + wp = traj.shape[-2] + pos_ids = torch.arange(wp, device=args.device).reshape(1, -1) + pos_embed = inner.pos_encoding(pos_ids).float().cpu() + + payload = { + "seed": SEED, + "z": z.float().cpu(), + "images_dp": images_dp.float().cpu(), # [B, T, H, W, C], generate_traj form + "images_chw": images_chw.float().cpu(), # [T, C, H, W], memory-engine form + "memory_tokens": memory_tokens, + "pos_embed": pos_embed, + "init_latents": captured["latents"], + "dit_step": step, + "trajectory": traj, + "num_sample_trajs": args.num_sample_trajs, + "num_inference_steps": args.num_inference_steps, + "cond_projector": {k: v.float().cpu() + for k, v in inner.cond_projector.state_dict().items()}, + "action_encoder": {k: v.float().cpu() + for k, v in inner.action_encoder.state_dict().items()}, + "action_decoder": {k: v.float().cpu() + for k, v in inner.action_decoder.state_dict().items()}, + } + os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True) + torch.save(payload, args.output_path) + size_mb = os.path.getsize(args.output_path) / 1e6 + print(f"\nWrote {args.output_path} ({size_mb:.1f} MB)") + print("Now run verify/compare_system1_engines.py under the TensorRT environment.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 4dc04676b5ccd9fa62f02dfa49fac0cf72fee1cb Mon Sep 17 00:00:00 2001 From: hungho77 Date: Fri, 14 Aug 2026 15:16:15 +0700 Subject: [PATCH 25/30] refactor(internvla-n1-dualvln): consolidate the benchmark matrix across both systems The results section reported System 2 and System 1 separately, with no dtype column and no System-1 latency at all, so nothing said what a whole planning step costs. Adds System-1 latency (stage B now takes --bench_iters) and the PyTorch baseline to compare it against, then restructures the matrix into three tables: System 2 by quantization scheme, System 1 as a conversion result, and both together per planning step. Every row names its dtype -- weights, activations, KV cache and vision tower separately, since s1 leaves the last two alone. System 1 engines 63.1 ms vs 175.4 ms in PyTorch 2.78x full step, FP8 709 ms vs 1806 ms all-PyTorch 2.55x weights, FP8 9.16 GB vs 15.7 GB unquantized Latency is flat in num_sample_trajs on the PyTorch side (175.4 / 175.6 / 173.1 ms at 32 / 4 / 1), so System 1 is launch-bound, which is why engines pay there. Also drops the stale duplicate table and the superseded 12-sample figures, and fixes the stage B GELU to approximate="tanh" to match the checkpoint's cond_projector. --- recipes/internvla-n1-dualvln/README.md | 95 +++++++++++++------ .../verify/compare_system1_engines.py | 70 +++++++++++--- 2 files changed, 124 insertions(+), 41 deletions(-) diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index 2907993..a02bc47 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -9,6 +9,9 @@ vision-language navigation model, targeting NVIDIA Jetson Thor. - Model family: `InternVLA-N1-DualVLN` (System 2 = Qwen2.5-VL-7B, System 1 = NextDiT trajectory head) - Quantization presets: `fp8_default`, `fp8_per_channel` (validated) · `nvfp4_*` (experimental, see Notes) - Runtime target: TensorRT-Edge-LLM engines on Jetson Thor (sm_110) +- End to end on TensorRT: System 2 FP8 **and** System 1 BF16, both verified against PyTorch + (bridge 0.9919, System-1 trajectory 0.9997) — **2.55x** faster per planning step at + 9.16 GB of weights against 15.7 GB unquantized ## What this model is, and what the metric has to be @@ -35,23 +38,63 @@ this recipe is gated on `z_latents` cosine against the FP32 reference, not on ge ## Results -### Benchmark matrix +### Benchmark matrix — both systems -All three variants built from the same repackaged System 2 and measured on Jetson Thor, -batch 1, on an idle GPU. +Everything below was measured on one Jetson Thor, batch 1, on an idle GPU. **System 2 is the +only part that gets quantized**; System 1 stays BF16 by design, so its row is a conversion +result rather than a quantization one. -| Variant | checkpoint | LLM engine | visual | prefill (1024) | decode (pastKV 1024) | z_latents (engine) | z_latents (weights only) | pixel L2 mean / median | -|---|---|---|---|---|---|---|---|---| -| BF16 (unquantized) | 16.6 GB | 14.15 GB | 1.36 GB | 135.8 ms | 56.4 ms | **0.999471** | — | 47.24 / 27.05 px | -| **FP8 s1** | 10.1 GB | **7.62 GB** | 1.36 GB | **82.1 ms** | **31.5 ms** | **0.991861** | 0.998020 | 46.26 / **22.51 px** | -| NVFP4 s1 (experimental) | 7.2 GB | 4.77 GB | 1.36 GB | 73.2 ms | 20.2 ms | 0.931005 ✗ | 0.987986 | 40.69 / 23.54 px | +#### System 2 — Qwen2.5-VL-7B planner (quantized) + +| Variant | weights / activations | KV cache | vision tower | checkpoint | LLM engine | visual engine | prefill (1024) | decode (pastKV 1024) | z_latents (engine) | pixel L2 mean / median | +|---|---|---|---|---|---|---|---|---|---|---| +| BF16 baseline | BF16 W16A16 -> FP16 engine | FP16 | BF16 | 16.6 GB | 14.15 GB | 1.36 GB | 135.8 ms | 56.4 ms | **0.999471** | 47.24 / 27.05 px | +| **FP8 s1** | **FP8 E4M3, W8A8, per-channel** | FP16 | BF16 | 10.1 GB | **7.62 GB** | 1.36 GB | **82.1 ms** | **31.5 ms** | **0.991861** | 46.26 / **22.51 px** | +| NVFP4 s1 (experimental) | NVFP4 E2M1, W4A4, block 16 w/ FP8 block scales | FP16 | BF16 | 7.2 GB | 4.77 GB | 1.36 GB | 73.2 ms | 20.2 ms | 0.931005 ✗ | 40.69 / 23.54 px | + +The KV cache is FP16 in all three: NVFP4 KV needs `sm100f` (datacenter Blackwell) and Thor is +sm110, and FP8 KV is a separate strategy (`s2`). The vision tower stays BF16 under `s1`; +quantizing it is `s3`/`s4` and is FP8-only, because the ViT MLP `intermediate_size` is 3420 +and 3420 / 16 = 213.75 does not divide by the NVFP4 block size. Weight quantization error, mean relative over 21 projections in layers 0/13/27: FP8 **2.67 %**, NVFP4 **9.45 %**. +#### System 1 — NextDiT diffusion head + memory block (BF16, not quantized) + +Engine weights are BF16; the engine *interface* is fp32/int64, and passing bf16 tensors trips +an assertion in the wrapper rather than converting silently. + +| Component | weights | ONNX | engine | latency | cosine vs PyTorch | +|---|---|---|---|---|---| +| memory block (DAv2 + MemoryEncoder + QFormer) | BF16 | 200 MB | **104 MB** | **2.00 ms** | **0.999981** | +| `traj_dit`, one diffusion step | BF16 | 134 MB | **72 MB** | **5.78 ms** | **0.999508** | +| full trajectory (10 steps x 32 samples x 32 waypoints) | BF16 | — | — | **61.14 ms** | **0.999670** | +| **System 1 total** (memory + trajectory) | BF16 | 334 MB | **176 MB** | **63.1 ms · 15.8 Hz** | — | +| PyTorch `generate_traj` baseline | BF16 | — | — | 175.4 ms · 5.7 Hz | — | + +TensorRT gives System 1 a **2.78x** speedup at identical output. Latency is flat in +`num_sample_trajs` on the PyTorch side (175.4 / 175.6 / 173.1 ms at 32 / 4 / 1), so the head +is launch-bound rather than compute-bound — which is also why moving it to engines pays. + +#### Both systems, one planning step + +System 2 latency here is the full multi-image VLN step (~9 images, ~1764 image tokens) from +12 held-out samples, **not** the synthetic `llm_bench` figures above — different measurement, +so do not mix the two columns. + +| Configuration | System 2 quantization | System 2 | System 1 | total | vs PyTorch | +|---|---|---|---|---|---| +| all PyTorch | BF16 | 1631 ms | 175.4 ms | 1806 ms | 1.00x | +| TensorRT, unquantized | FP16 | 770 ms | 63.1 ms | 833 ms | 2.17x | +| **TensorRT, recommended** | **FP8 E4M3 W8A8** | **646 ms** | **63.1 ms** | **709 ms** | **2.55x** | + +Total on-device weights for the recommended configuration: **7.62 + 1.36 + 0.18 = 9.16 GB**, +against 15.7 GB for the unquantized TensorRT path. + **FP8 is the recommended scheme.** Against BF16 it is 1.86x smaller and 1.65x/1.79x faster, holds the bridge at 0.9919, and the median waypoint error does not get worse — it improves -slightly (27.05 → 22.51 px), which is within the spread of a 42-sample set and should be read +slightly (27.05 -> 22.51 px), which is within the spread of a 42-sample set and should be read as "unchanged", not as a gain from quantization. **NVFP4 is faster and smaller still but fails the gate.** Its bridge sits at 0.931, below the @@ -198,14 +241,15 @@ Upstream InternNav ships **no** TensorRT or ONNX export at all — nothing in `i references `trtexec`, `tensorrt` or `torch.onnx`. The System-1 conversion here is entirely this recipe's, ported from the source project. -| Component | ONNX | engine | I/O (verified by execution) | -|---|---|---|---| -| `traj_dit` (NextDiT diffusion head) | 134 MB | **72 MB** | `x[64,32,384] f32`, `timestep[64] i64`, `z_latents[64,*,768] f32` -> `output[64,32,384]` | -| memory block (DepthAnythingV2 + MemoryEncoder + QFormer) | 200 MB | **104 MB** | `images[T,3,224,224] f32` -> `memory_tokens[1,32,768]` | +Sizes, latency and fidelity are in the matrix above. The engine I/O, verified by execution: -Both were run with real inputs: finite output, correct shapes, |max| 2.84 and 3.70. Note -the engine I/O is fp32/int64 even though the weights are BF16 — passing bf16 tensors fails -an assertion in the wrapper rather than converting silently. +| Engine | inputs | output | +|---|---|---| +| `traj_dit` | `x[64,32,384]` f32, `timestep[64]` i64, `z_latents[64,*,768]` f32 | `output[64,32,384]` f32 | +| memory block | `images[T,3,224,224]` f32 | `memory_tokens[1,32,768]` f32 | + +The `64` is `2 x num_sample_trajs`: `generate_traj` runs classifier-free guidance, so the +conditioning is `[null, real]` and the latents are duplicated. **What stays in PyTorch on the host**, by design rather than omission: `action_encoder` and `action_decoder` (two 3x384 linears), `pos_encoding`, `cond_projector` (the System 2 bridge), @@ -238,7 +282,8 @@ python verify/dump_system1_reference.py --output_path work/system1_reference.pt # stage B, TensorRT environment (Python 3.12) python verify/compare_system1_engines.py \ - --reference_path work/system1_reference.pt --engine_dir work/onnx + --reference_path work/system1_reference.pt --engine_dir work/onnx \ + --bench_iters 20 # optional: also time each engine and the whole trajectory ``` Two details are worth keeping, because both produce a confident wrong answer: @@ -255,18 +300,12 @@ Two details are worth keeping, because both produce a confident wrong answer: what separates a bad engine from a bad reimplementation of the sampler loop — here it read 0.9995 while the trajectory still read 0.31, which localized the fault to the harness. -### Earlier full-pipeline figures - -Measured on Jetson Thor, 12 held-out multi-image VLN steps: - -| LLM variant | z_latents vs FP32 | agrees w/ PyTorch | System 2 latency | LLM engine | -|---|---|---|---|---| -| PyTorch BF16 (baseline) | 0.99974 | — | 1631 ms · 1.00x | ~14 GB weights | -| base FP16 TensorRT (no quant) | **0.99985** | 12/12 | 770 ms · 2.12x | 14.2 GB | -| FP8 TensorRT | 0.99559 | 11/12 | **646 ms · 2.53x** | **7.6 GB** | -| NVFP4 TensorRT | **0.647** ✗ | tokens fine, bridge broken | — | 4.5 GB | +### A note on the older numbers -Other engines: ViT 1.3 GB BF16 → 0.68 GB FP8; traj_dit 0.07 GB; memory block 0.11 GB. +An earlier revision reported z_latents of 0.99974 / 0.99985 / 0.99559 / 0.647 for +PyTorch / FP16 / FP8 / NVFP4 against an FP32 reference on 12 samples. The matrix above +supersedes those: it uses a BF16 reference, 42 samples, and the corrected checkpoint loader. +The ordering is the same and the conclusion is unchanged — FP8 passes, NVFP4 does not. **Calibration data made no measurable difference.** Held-out z_latents came out at 0.99143 with generic `cnn_dailymail` text versus 0.99146 with a domain-specific VLN set — equal diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/compare_system1_engines.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/compare_system1_engines.py index 41cc867..f64eea6 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/compare_system1_engines.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/compare_system1_engines.py @@ -24,6 +24,7 @@ import argparse import os import sys +import time import numpy as np import torch @@ -35,6 +36,18 @@ from trt_torch import Engine # noqa: E402 +def timeit(fn, iters: int, warmup: int = 3) -> float: + """Mean wall-clock milliseconds, synchronized on both sides.""" + for _ in range(warmup): + fn() + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(iters): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - t0) * 1000.0 / iters + + def cos(a: torch.Tensor, b: torch.Tensor) -> float: a = a.double().flatten() b = b.double().flatten() @@ -60,6 +73,9 @@ def parse_args() -> argparse.Namespace: help="Host-side dtype for the diffusion loop. generate_traj runs it in " "bfloat16; fp32 here diverges from the reference even with correct " "engines, because the sampler amplifies the difference.") + p.add_argument("--bench_iters", type=int, default=0, + help="If > 0, also time each engine and the whole trajectory. Run on an " + "idle GPU -- a shared device reads 40-60%% high.") p.add_argument("--gate", type=float, default=0.99, help="Minimum trajectory cosine to report PASS") return p.parse_args() @@ -82,6 +98,9 @@ def main() -> int: # The memory engine was built from MemBlock, which takes [T, C, H, W]. images = ref["images_chw"].to(dev).float().contiguous() tokens = out_of(mem(images=images), "memory_tokens").float().cpu() + timing = {} + if args.bench_iters: + timing["memory_ms"] = timeit(lambda: mem(images=images), args.bench_iters) mem.close() mem_cos = cos(ref["memory_tokens"], tokens) print(f" memory_tokens cosine vs PyTorch: {mem_cos:.6f}") @@ -93,6 +112,12 @@ def main() -> int: one = out_of(dit_probe(x=st["x"].to(dev).contiguous(), timestep=st["timestep"].to(dev).to(torch.int64).contiguous(), z_latents=st["z_latents"].to(dev).contiguous()), "output") + if args.bench_iters: + timing["traj_dit_step_ms"] = timeit( + lambda: dit_probe(x=st["x"].to(dev).contiguous(), + timestep=st["timestep"].to(dev).to(torch.int64).contiguous(), + z_latents=st["z_latents"].to(dev).contiguous()), + args.bench_iters) dit_probe.close() step_cos = cos(st["output"], one.float().cpu()) print(f" traj_dit single-step cosine: {step_cos:.6f}") @@ -115,7 +140,7 @@ def main() -> int: cp = {k: v.to(dev) for k, v in ref["cond_projector"].items()} z = ref["z"].to(dev) z = torch.nn.functional.linear(z, cp["0.weight"], cp.get("0.bias")) - z = torch.nn.functional.gelu(z) + z = torch.nn.functional.gelu(z, approximate="tanh") z = torch.nn.functional.linear(z, cp["2.weight"], cp.get("2.bias")) # generate_traj runs classifier-free guidance: the conditioning is [null, real] and the # latents are duplicated, so the engine's batch is 2 * num_sample_trajs. That is why the @@ -141,18 +166,30 @@ def main() -> int: cond, pos_embed = cond.to(dt), pos_embed.to(dt) enc_w, enc_b, dec_w, dec_b = (t.to(dt) for t in (enc_w, enc_b, dec_w, dec_b)) dit.set_runtime_tensor_shape("z_latents", tuple(cond.shape)) - for t in scheduler.timesteps: - feats = torch.nn.functional.linear(latents, enc_w, enc_b) + pos_embed - feats = feats.repeat(2, 1, 1) - if hasattr(scheduler, "scale_model_input"): - feats = scheduler.scale_model_input(feats, t) - ts = t.to(dev).expand(feats.shape[0]).to(torch.int64).contiguous() - pred = out_of(dit(x=feats.float().contiguous(), timestep=ts, - z_latents=cond.float().contiguous()), "output").to(dt) - pred = torch.nn.functional.linear(pred, dec_w, dec_b) - uncond, condit = pred.chunk(2) - pred = uncond + args.guidance_scale * (condit - uncond) - latents = scheduler.step(pred, t, latents).prev_sample + + def run_loop(latents): + """One full sampling run. Mirrors generate_traj line for line.""" + # The scheduler carries step_index across calls, so benchmarking a second run + # walks off the end of its sigma table. Reset it per run. + scheduler.set_timesteps(steps, sigmas=np.linspace(1.0, 1 / steps, steps)) + for t in scheduler.timesteps: + feats = torch.nn.functional.linear(latents, enc_w, enc_b) + pos_embed + feats = feats.repeat(2, 1, 1) + if hasattr(scheduler, "scale_model_input"): + feats = scheduler.scale_model_input(feats, t) + ts = t.to(dev).expand(feats.shape[0]).to(torch.int64).contiguous() + pred = out_of(dit(x=feats.float().contiguous(), timestep=ts, + z_latents=cond.float().contiguous()), "output").to(dt) + pred = torch.nn.functional.linear(pred, dec_w, dec_b) + uncond, condit = pred.chunk(2) + pred = uncond + args.guidance_scale * (condit - uncond) + latents = scheduler.step(pred, t, latents).prev_sample + return latents + + latents = run_loop(latents) + if args.bench_iters: + start = ref["init_latents"].to(dev).to(dt) + timing["trajectory_ms"] = timeit(lambda: run_loop(start), args.bench_iters) dit.close() traj = latents.float().cpu() @@ -173,6 +210,13 @@ def main() -> int: print(f" trajectory cosine : {traj_cos:.6f}") print(f" trajectory rel-L2 : {l2:.4f}") ok = traj_cos >= args.gate + if timing: + print("\n --- latency, mean over " + f"{args.bench_iters} iterations (idle GPU assumed) ---") + print(f" memory block engine : {timing['memory_ms']:.2f} ms") + print(f" traj_dit, one step : {timing['traj_dit_step_ms']:.2f} ms") + print(f" full trajectory : {timing['trajectory_ms']:.2f} ms " + f"({steps} steps, {n_traj} samples)") print(f" {'PASS' if ok else 'BELOW GATE'} (gate {args.gate})") return 0 if ok else 1 From eac8553f6c374bb101d5c1eb93b665c99beda303 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Fri, 14 Aug 2026 15:38:53 +0700 Subject: [PATCH 26/30] feat(internvla-n1-dualvln): measure FP8 on System 1, and keep BF16 System 1 shipped BF16 on the assumption that quantizing it was not worth it. This measures that instead of assuming it, and the assumption holds -- for a reason that is worth writing down. Unlike System 2, System 1 goes torch.onnx.export -> trtexec, and TensorRT does FP8 only through explicit quantization, so the Q/DQ has to be in the ONNX. Adds a ModelOpt PTQ pass calibrated on real tensors captured from a live System 2 -> System 1 run rather than random draws: FP8 scales are amax-based and both modules consume tensors whose scale is set upstream, so a synthetic draw is quietly wrong. FP8 works -- 328/328 Q/DQ pairs in traj_dit, 160/160 in the memory block, 1.55x smaller, 20% faster on the diffusion loop -- and is still the wrong trade: engines 176 MB -> 114 MB 0.7% of 9.16 GB deployed step 710 ms -> 698 ms 1.7%, System 2 dominates waypoint dev 0.0032 -> 0.0198 6x, p95 39% of a waypoint's reach Splitting it shows where the loss comes from: traj_dit alone costs 0.9997 -> 0.9880 even though its per-step error is only 0.999508 -> 0.997811, because the sampler runs 10 steps and each feeds the next. The memory block adds the rest and buys nothing in latency (2.09 vs 2.04 ms). Adds waypoint deviation to the parity check -- cosine says how aligned two trajectories are, not how far apart the robot ends up -- plus --engine_suffix so the same check runs against either precision. Three things that fail quietly, all now handled: ModelOpt emits trt::TRT_FP8*, not ONNX QuantizeLinear, so a naive counter reports a quantized graph as unquantized; the legacy exporter cannot infer a conv kernel shape through Q/DQ; and the QFormer's fused MHA fast path has no ONNX symbolic. --- recipes/internvla-n1-dualvln/Makefile | 13 +- recipes/internvla-n1-dualvln/README.md | 66 ++++- .../trt-edgellm/dump_system1_calib.py | 142 ++++++++++ .../trt-edgellm/quantize_system1.py | 262 ++++++++++++++++++ .../verify/compare_system1_engines.py | 24 +- 5 files changed, 492 insertions(+), 15 deletions(-) create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/dump_system1_calib.py create mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/quantize_system1.py diff --git a/recipes/internvla-n1-dualvln/Makefile b/recipes/internvla-n1-dualvln/Makefile index 98ffa98..456ccbe 100644 --- a/recipes/internvla-n1-dualvln/Makefile +++ b/recipes/internvla-n1-dualvln/Makefile @@ -79,7 +79,7 @@ benchmark: --work_dir $(WORK_DIR) # ── System 1: BF16 engines (requires INTERNNAV_PATH) ────────────────────────── -.PHONY: export-system1 verify-system1 +.PHONY: export-system1 verify-system1 quantize-system1-fp8 export-system1: bash trt-edgellm/scripts/04_export_system1.sh \ @@ -87,6 +87,16 @@ export-system1: --internnav_path $(INTERNNAV_PATH) \ --engine_dir $(ENGINE_DIR)/system1 +# Measured and deliberately not shipped: FP8 costs 6x the waypoint deviation for 1.7% of +# the planning step. Kept so the finding stays reproducible. See README. +quantize-system1-fp8: + python trt-edgellm/dump_system1_calib.py \ + --calib_data_root $(CALIB_DATA_ROOT) \ + --output_path $(WORK_DIR)/system1_calib.pt + python trt-edgellm/quantize_system1.py \ + --calib_path $(WORK_DIR)/system1_calib.pt \ + --out_dir $(ENGINE_DIR)/system1_fp8 + verify-system1: bash trt-edgellm/scripts/05_verify.sh \ --system1 \ @@ -128,6 +138,7 @@ help: @echo " System 1 (requires INTERNNAV_PATH):" @echo " export-system1 traj_dit + memory block -> BF16 engines" @echo " verify-system1 trajectory parity vs PyTorch" + @echo " quantize-system1-fp8 System 1 FP8 PTQ (measured, not recommended)" @echo "" @echo " Other:" @echo " investigate-nvfp4 why NVFP4 breaks the System2 -> System1 bridge" diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index a02bc47..6138888 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -60,23 +60,64 @@ and 3420 / 16 = 213.75 does not divide by the NVFP4 block size. Weight quantization error, mean relative over 21 projections in layers 0/13/27: FP8 **2.67 %**, NVFP4 **9.45 %**. -#### System 1 — NextDiT diffusion head + memory block (BF16, not quantized) +#### System 1 — NextDiT diffusion head + memory block -Engine weights are BF16; the engine *interface* is fp32/int64, and passing bf16 tensors trips -an assertion in the wrapper rather than converting silently. +System 1 ships **BF16**. FP8 was measured rather than assumed, and the measurement is why it +does not ship: see below. Engine weights are BF16; the engine *interface* is fp32/int64, and +passing bf16 tensors trips an assertion in the wrapper rather than converting silently. | Component | weights | ONNX | engine | latency | cosine vs PyTorch | |---|---|---|---|---|---| -| memory block (DAv2 + MemoryEncoder + QFormer) | BF16 | 200 MB | **104 MB** | **2.00 ms** | **0.999981** | -| `traj_dit`, one diffusion step | BF16 | 134 MB | **72 MB** | **5.78 ms** | **0.999508** | -| full trajectory (10 steps x 32 samples x 32 waypoints) | BF16 | — | — | **61.14 ms** | **0.999670** | -| **System 1 total** (memory + trajectory) | BF16 | 334 MB | **176 MB** | **63.1 ms · 15.8 Hz** | — | +| memory block (DAv2 + MemoryEncoder + QFormer) | BF16 | 200 MB | **104 MB** | **2.04 ms** | **0.999981** | +| `traj_dit`, one diffusion step | BF16 | 134 MB | **72 MB** | **5.85 ms** | **0.999508** | +| full trajectory (10 steps x 32 samples x 32 waypoints) | BF16 | — | — | **61.8 ms** | **0.999670** | +| **System 1 total** (memory + trajectory) | BF16 | 334 MB | **176 MB** | **63.8 ms · 15.7 Hz** | — | | PyTorch `generate_traj` baseline | BF16 | — | — | 175.4 ms · 5.7 Hz | — | -TensorRT gives System 1 a **2.78x** speedup at identical output. Latency is flat in +TensorRT gives System 1 a **2.75x** speedup at identical output. Latency is flat in `num_sample_trajs` on the PyTorch side (175.4 / 175.6 / 173.1 ms at 32 / 4 / 1), so the head is launch-bound rather than compute-bound — which is also why moving it to engines pays. +#### FP8 on System 1 — measured, and not recommended + +Unlike System 2, System 1 goes `torch.onnx.export` -> `trtexec`, and TensorRT does FP8 only +through **explicit** quantization: the Q/DQ nodes must already be in the ONNX. So this needed +a ModelOpt PTQ pass (`quantize_system1.py`) calibrated on real tensors captured from a live +System 2 -> System 1 run (`dump_system1_calib.py`, 40 real `traj_dit` batches over 4 VLN +samples). Verified in the graph: 328 / 328 FP8 Q/DQ pairs in `traj_dit`, 160 / 160 in the +memory block. + +Waypoint deviation is the number to read. Cosine says how *aligned* two trajectories are; +deviation says how far apart the robot would actually end up, in the trajectory's own units, +against a mean per-waypoint reach of **0.2052**. + +| Config | traj_dit | memory | engines | trajectory | traj cosine | waypoint dev. mean / median / p95 | +|---|---|---|---|---|---|---| +| **BF16 (shipped)** | BF16 72 MB | BF16 104 MB | **176 MB** | 61.8 ms | **0.999670** | **0.0032 / 0.0012 / 0.0117** | +| mixed | **FP8 39 MB** | BF16 104 MB | 143 MB | 48.4 ms | 0.988032 ✗ | 0.0154 / 0.0050 / 0.0521 | +| all FP8 | **FP8 39 MB** | **FP8 75 MB** | **114 MB** | 49.6 ms | 0.980213 ✗ | 0.0198 / 0.0053 / 0.0793 | + +FP8 works — 1.55x smaller, 20 % faster on the diffusion loop, both configs well clear of +garbage — and it is still the wrong trade here: + +- **The size saving is irrelevant at the system level.** 62 MB off 9.16 GB of deployed + weights is 0.7 %. +- **The latency saving is nearly as small.** 12 ms off a 709 ms planning step is 1.7 %, + because System 2 dominates by an order of magnitude. +- **The fidelity cost is not small.** Mean waypoint deviation goes 0.0032 -> 0.0198, a **6x** + increase, and p95 goes 0.0117 -> 0.0793, which is 39 % of a typical waypoint's reach. + +Both FP8 configs fall below the 0.99 gate, and the split shows why: quantizing `traj_dit` +alone already costs 0.9997 -> 0.9880. Its per-step error is only 0.999508 -> 0.997811, but the +sampler runs 10 steps and each one feeds the next, so a small per-call error compounds. The +memory block adds the rest (0.999981 -> 0.990204) and buys **nothing** in latency (2.09 ms +against 2.04 ms) — its Conv2d stays unquantized anyway, since the legacy ONNX exporter cannot +infer a convolution kernel shape through Q/DQ. + +**Keep System 1 in BF16.** Quantization effort belongs on System 2, which is 98 % of both the +weights and the latency. The scripts stay in the recipe so the measurement is reproducible +and so the finding can be re-checked on a different System-1 configuration. + #### Both systems, one planning step System 2 latency here is the full multi-image VLN step (~9 images, ~1764 image tokens) from @@ -86,8 +127,11 @@ so do not mix the two columns. | Configuration | System 2 quantization | System 2 | System 1 | total | vs PyTorch | |---|---|---|---|---|---| | all PyTorch | BF16 | 1631 ms | 175.4 ms | 1806 ms | 1.00x | -| TensorRT, unquantized | FP16 | 770 ms | 63.1 ms | 833 ms | 2.17x | -| **TensorRT, recommended** | **FP8 E4M3 W8A8** | **646 ms** | **63.1 ms** | **709 ms** | **2.55x** | +| TensorRT, unquantized | FP16 | 770 ms | 63.8 ms | 834 ms | 2.17x | +| **TensorRT, recommended** | **FP8 E4M3 W8A8** (System 1 stays BF16) | **646 ms** | **63.8 ms** | **710 ms** | **2.54x** | +| TensorRT, System 1 also FP8 | FP8 both systems | 646 ms | 51.6 ms | 698 ms | 2.59x ✗ | + +The last row is why System 1 stays BF16: 1.7 % off the step for 6x the waypoint deviation. Total on-device weights for the recommended configuration: **7.62 + 1.36 + 0.18 = 9.16 GB**, against 15.7 GB for the unquantized TensorRT path. @@ -336,6 +380,8 @@ navigation model), but do not expect it to buy accuracy. ├── engine_runner.py # direct-TensorRT LLM harness (hand-built 3D mRoPE) ├── export_traj_dit.py # System 1 diffusion head -> ONNX -> BF16 engine ├── export_memory_block.py # System 1 memory block -> ONNX -> BF16 engine + ├── dump_system1_calib.py # real System 1 calibration tensors (needs InternNav) + ├── quantize_system1.py # System 1 FP8 PTQ -> ONNX -> engine (measured, not shipped) ├── internvla_compat.py # the three patches needed to load System 1 ├── traj_dit_loader.py memblock.py ├── trt_torch.py # NVIDIA Apache-2.0 — header kept, not restamped diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/dump_system1_calib.py b/recipes/internvla-n1-dualvln/trt-edgellm/dump_system1_calib.py new file mode 100644 index 0000000..931e717 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/dump_system1_calib.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Capture real System-1 calibration tensors, for quantizing traj_dit and the memory block. + +FP8 in TensorRT is explicit-quantization only, so the ONNX has to carry Q/DQ nodes, which +means a PTQ pass in PyTorch first, which means calibration data. For System 1 that data is +awkward to synthesize: ``traj_dit`` consumes ``z_latents`` produced by System 2 through +``cond_projector``, and the memory block consumes normalized navigation frames. Random +tensors have the wrong scale in both cases, and FP8 calibration is amax-based, so wrong +scale means wrong scale factors. + +This runs the real thing. For each VLN sample it takes real frames, runs System 2's own +``generate_latents`` for a real ``z``, then taps every ``traj_dit`` call the diffusion loop +makes. The result is cached to disk because capturing it runs System 2 on every sample: +re-quantizing with a different config then costs a minute rather than the whole pipeline. + +Run this under the transformers 4.51 environment (Python 3.10 here):: + + INTERNNAV_PATH=~/InternNav PYTHONPATH=~/InternNav \\ + python dump_system1_calib.py --calib_data_root work/calib_scenes \\ + --output_path work/system1_calib.pt +""" +import argparse +import os +import sys + +import torch + +_HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, _HERE) +sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "quantize")) + +import internvla_compat # noqa: E402 + +SEED = 12345 + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--internvla_ckpt", + default=os.path.expanduser( + os.environ.get("INTERNVLA_CKPT", + "~/InternNav/checkpoints/InternVLA-N1-DualVLN"))) + p.add_argument("--calib_data_root", required=True, + help="LeRobot scene root, as used by quantize/benchmark_accuracy.py") + p.add_argument("--output_path", required=True) + p.add_argument("--num_samples", type=int, default=4, + help="VLN samples to draw. Each contributes num_inference_steps traj_dit " + "batches, so 4 gives 40 -- ample for amax calibration.") + p.add_argument("--num_sample_trajs", type=int, default=32) + p.add_argument("--num_inference_steps", type=int, default=10) + p.add_argument("--num_frames", type=int, default=2) + p.add_argument("--device", default="cuda") + return p.parse_args() + + +def main() -> int: + args = parse_args() + internvla_compat.apply_all(need_system1=True, allow_missing_depth=False) + + import numpy as np + import prompt_builder as pb + from PIL import Image + from benchmark_accuracy import discover_samples + from internnav.model.basemodel.internvla_n1.internvla_n1 import ( + InternVLAN1ForCausalLM, InternVLAN1ModelConfig) + from transformers import AutoProcessor + + print(f"[1/4] Loading {args.internvla_ckpt}") + config = InternVLAN1ModelConfig.from_pretrained(args.internvla_ckpt) + model = InternVLAN1ForCausalLM.from_pretrained( + args.internvla_ckpt, config=config, torch_dtype=torch.bfloat16, + attn_implementation="sdpa", low_cpu_mem_usage=True).to(args.device).eval() + processor = AutoProcessor.from_pretrained(args.internvla_ckpt, + min_pixels=128 * 28 * 28, + max_pixels=2048 * 32 * 32) + + print(f"[2/4] Drawing {args.num_samples} real VLN samples") + samples = discover_samples(args.calib_data_root, args.num_samples, seed=SEED) + if not samples: + print(f"[ERROR] no samples under {args.calib_data_root}") + return 1 + + inner = model.get_model() + dit_batches, image_batches = [], [] + + _real_fwd = inner.traj_dit.forward + + def _tap(x, timestep, z_latents, *a, **kw): + dit_batches.append((x.detach().half().cpu(), + timestep.detach().cpu(), + z_latents.detach().half().cpu())) + return _real_fwd(x, timestep, z_latents, *a, **kw) + + print("[3/4] Running System 2 -> System 1 on each sample") + torch.manual_seed(SEED) + for i, sample in enumerate(samples): + paths = sample["images"][-args.num_frames:] + frames = [np.asarray(Image.open(p).convert("RGB").resize((224, 224))) / 255.0 + for p in paths] + images_dp = torch.from_numpy(np.stack(frames)).unsqueeze(0) + images_dp = images_dp.to(args.device, torch.bfloat16) # [1, T, 224, 224, 3] + + enc = pb.build_sample_inputs(sample, processor) + enc = {k: v.to(args.device) for k, v in enc.items() if isinstance(v, torch.Tensor)} + with torch.no_grad(): + # The real bridge output, not a draw: generate_latents runs the LLM with the + # learned latent queries appended and returns the normalized TRAJ hidden states. + z = model.generate_latents(enc["input_ids"], enc.get("pixel_values"), + enc.get("image_grid_thw")) + + chw = images_dp.permute(0, 1, 4, 2, 3) + norm = ((chw - model._resnet_mean) / model._resnet_std).flatten(0, 1) + image_batches.append(norm.to(torch.bfloat16).float().cpu()) + + inner.traj_dit.forward = _tap + model.generate_traj(z, images_dp, + num_sample_trajs=args.num_sample_trajs, + num_inference_steps=args.num_inference_steps) + inner.traj_dit.forward = _real_fwd + print(f" sample {i + 1}/{len(samples)}: {len(dit_batches)} traj_dit batches") + + payload = { + "dit_batches": dit_batches, + "image_batches": image_batches, + "num_sample_trajs": args.num_sample_trajs, + "num_inference_steps": args.num_inference_steps, + } + print("[4/4] Writing") + os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True) + torch.save(payload, args.output_path) + print(f"\nWrote {args.output_path} " + f"({os.path.getsize(args.output_path) / 1e6:.1f} MB): " + f"{len(dit_batches)} traj_dit batches, {len(image_batches)} image batches") + print("Now run quantize_system1.py (same environment) to PTQ and build.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/quantize_system1.py b/recipes/internvla-n1-dualvln/trt-edgellm/quantize_system1.py new file mode 100644 index 0000000..0b840aa --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/quantize_system1.py @@ -0,0 +1,262 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""FP8 PTQ for System 1: quantize traj_dit and the memory block, export ONNX, build engines. + +System 1 ships BF16 by default. This is the experiment that asks whether it should not. + +The route differs from System 2's. System 2 goes through TensorRT-Edge-LLM, which consumes a +quantized HF checkpoint and inserts the scaling itself. System 1 is a plain +``torch.onnx.export`` into ``trtexec``, and **TensorRT supports FP8 only through explicit +quantization** -- there is no implicit FP8 calibration the way there is for INT8. So the Q/DQ +nodes have to be in the ONNX, which means a ModelOpt PTQ pass in PyTorch first. + +Calibration comes from ``dump_system1_calib.py`` rather than from random tensors, and that +matters more here than usual: FP8 scale factors are amax-based, and both of these modules +consume tensors whose scale is set by something upstream -- ``z_latents`` by System 2 through +``cond_projector``, the frames by the ResNet normalization. A synthetic draw has the wrong +amax and produces wrong scales, quietly. + +The calibration bundle is read from disk rather than regenerated because capturing it runs +System 2 on every sample; caching it means re-quantizing with a different config costs a +minute instead of the whole pipeline. + +Run under the transformers 4.51 environment (Python 3.10 here), which has both ModelOpt and +a working InternNav:: + + INTERNNAV_PATH=~/InternNav PYTHONPATH=~/InternNav \\ + python quantize_system1.py --calib_path work/system1_calib.pt --out_dir work/onnx_fp8 +""" +import argparse +import os +import subprocess +import sys + +import torch + +_HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, _HERE) + +import internvla_compat # noqa: E402 +from memblock import MemBlock # noqa: E402 + +TRTEXEC = os.environ.get("TRTEXEC", "/usr/src/tensorrt/bin/trtexec") + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--internvla_ckpt", + default=os.path.expanduser( + os.environ.get("INTERNVLA_CKPT", + "~/InternNav/checkpoints/InternVLA-N1-DualVLN"))) + p.add_argument("--calib_path", required=True, + help="Bundle from dump_system1_calib.py") + p.add_argument("--out_dir", required=True) + p.add_argument("--components", default="traj_dit,memory", + help="Comma-separated subset of traj_dit,memory") + p.add_argument("--zlen", type=int, default=36, + help="Optimization profile z_latents length (32 memory tokens + 4 TRAJ)") + p.add_argument("--exclude_memory", default="nn.Conv2d", + help="Comma-separated exclusions for the memory block: 'nn.X' matches a " + "module class, anything else is a name glob. The default leaves " + "Conv2d alone because the legacy ONNX exporter cannot infer a " + "convolution kernel shape through Q/DQ and dies with 'convolution " + "for kernel of unknown shape'. In DepthAnythingV2 that is the " + "patch_embed projection -- one layer of a ViT, so almost no compute " + "is left behind.") + p.add_argument("--skip_build", action="store_true", + help="Export ONNX only, do not call trtexec") + p.add_argument("--device", default="cuda") + return p.parse_args() + + +def fp8_config(exclude: str = ""): + """FP8_DEFAULT_CFG with name patterns disabled.""" + import copy + + import modelopt.torch.quantization as mtq + + # In ModelOpt 0.44 quant_cfg is an ordered *list* of rules, not a dict, and later rules + # win -- so exclusions go on the end. Entries starting with "nn." match by module class + # (the form the built-in BatchNorm exclusions use); anything else matches by name glob. + cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + for pat in (e.strip() for e in exclude.split(",")): + if not pat: + continue + rule = {"quantizer_name": "*", "enable": False} + if pat.startswith("nn."): + rule["parent_class"] = pat + else: + rule["quantizer_name"] = pat + cfg["quant_cfg"].append(rule) + return cfg + + +def count_qdq(onnx_path: str) -> tuple[int, int]: + """Q/DQ node counts. Zero means the PTQ pass did not survive the export, and the + engine below would silently be BF16 with extra steps. + + ModelOpt emits TensorRT's own ``trt::TRT_FP8QuantizeLinear`` for FP8 rather than the + standard ONNX ``QuantizeLinear``, so counting only the latter reports a correctly + quantized graph as unquantized.""" + import onnx + + model = onnx.load(onnx_path, load_external_data=False) + ops = [n.op_type for n in model.graph.node] + q = ops.count("QuantizeLinear") + ops.count("TRT_FP8QuantizeLinear") + dq = ops.count("DequantizeLinear") + ops.count("TRT_FP8DequantizeLinear") + return q, dq + + +def build(onnx_path: str, engine_path: str, shapes: list[str]) -> bool: + cmd = [TRTEXEC, f"--onnx={onnx_path}", f"--saveEngine={engine_path}", "--fp8", "--bf16"] + cmd += shapes + print(" " + " ".join(cmd)) + r = subprocess.run(cmd, capture_output=True, text=True) + if not os.path.exists(engine_path): + tail = "\n ".join((r.stdout + r.stderr).strip().splitlines()[-12:]) + print(f" BUILD FAILED\n {tail}") + return False + print(f" engine {os.path.getsize(engine_path) / 1e6:.1f} MB") + return True + + +def quantize_traj_dit(inner, calib, args) -> str: + import modelopt.torch.quantization as mtq + + dit = inner.traj_dit.eval() + batches = calib["dit_batches"] + dev = args.device + + def forward_loop(m): + for x, ts, z in batches: + with torch.no_grad(): + m(x=x.to(dev, torch.bfloat16), timestep=ts.to(dev), + z_latents=z.to(dev, torch.bfloat16)) + + print(f"[traj_dit] PTQ FP8 over {len(batches)} real batches") + dit = mtq.quantize(dit, fp8_config(), forward_loop=forward_loop) + + # Export in FP32: the RoPE freqs_cis are float32 and a mixed-precision trace fails. + # Precision is re-established by the Q/DQ nodes plus trtexec --bf16. + dit = dit.to(torch.float32).eval() + x, ts, z = batches[0] + x, ts, z = x.to(dev).float(), ts.to(dev), z.to(dev).float() + + class Wrap(torch.nn.Module): + def __init__(self, d): + super().__init__() + self.d = d + + def forward(self, x, timestep, z_latents): + return self.d(x=x, timestep=timestep, z_latents=z_latents) + + w = Wrap(dit).eval() + onnx_path = os.path.join(args.out_dir, "system1_traj_dit_fp8.onnx") + dyn = {"x": {0: "batch"}, "timestep": {0: "batch"}, + "z_latents": {0: "batch", 1: "zlen"}, "output": {0: "batch"}} + print(f"[traj_dit] Export ONNX -> {onnx_path}") + with torch.inference_mode(): + torch.onnx.export(w, (x, ts, z), onnx_path, + input_names=["x", "timestep", "z_latents"], + output_names=["output"], opset_version=19, + do_constant_folding=True, export_params=True, + dynamic_axes=dyn, dynamo=False) + q, dq = count_qdq(onnx_path) + print(f" {os.path.getsize(onnx_path) / 1e6:.1f} MB, " + f"{q} QuantizeLinear / {dq} DequantizeLinear") + if q == 0: + print(" [WARN] no Q/DQ in the graph -- the engine will not be FP8") + + if args.skip_build: + return onnx_path + B, WP, DIM = x.shape[0], x.shape[1], x.shape[2] + zdim = z.shape[-1] + engine = os.path.join(args.out_dir, "system1_traj_dit_fp8.engine") + print("[traj_dit] Build engine") + build(onnx_path, engine, [ + f"--minShapes=x:{B}x{WP}x{DIM},timestep:{B},z_latents:{B}x4x{zdim}", + f"--optShapes=x:{B}x{WP}x{DIM},timestep:{B},z_latents:{B}x{args.zlen}x{zdim}", + f"--maxShapes=x:{B}x{WP}x{DIM},timestep:{B},z_latents:{B}x64x{zdim}"]) + return onnx_path + + +def quantize_memory(model, inner, calib, args) -> str: + import modelopt.torch.quantization as mtq + + dev = args.device + block = MemBlock(inner.rgb_model, inner.memory_encoder, inner.rgb_resampler).eval() + batches = calib["image_batches"] + + def forward_loop(m): + for imgs in batches: + with torch.no_grad(): + m(imgs.to(dev, torch.bfloat16)) + + print(f"[memory] PTQ FP8 over {len(batches)} real frame batches" + + (f", excluding {args.exclude_memory}" if args.exclude_memory else "")) + block = mtq.quantize(block, fp8_config(args.exclude_memory), forward_loop=forward_loop) + block = block.to(torch.float32).eval() + + imgs = batches[0].to(dev).float() + onnx_path = os.path.join(args.out_dir, "system1_memory_fp8.onnx") + print(f"[memory] Export ONNX -> {onnx_path}") + with torch.inference_mode(): + torch.onnx.export(block, (imgs,), onnx_path, input_names=["images"], + output_names=["memory_tokens"], opset_version=19, + do_constant_folding=True, export_params=True, + dynamic_axes={"images": {0: "frames"}}, dynamo=False) + q, dq = count_qdq(onnx_path) + print(f" {os.path.getsize(onnx_path) / 1e6:.1f} MB, " + f"{q} QuantizeLinear / {dq} DequantizeLinear") + if q == 0: + print(" [WARN] no Q/DQ in the graph -- the engine will not be FP8") + + if args.skip_build: + return onnx_path + T, C, H, W = imgs.shape + engine = os.path.join(args.out_dir, "system1_memory_fp8.engine") + print("[memory] Build engine") + build(onnx_path, engine, [ + f"--minShapes=images:1x{C}x{H}x{W}", + f"--optShapes=images:{T}x{C}x{H}x{W}", + f"--maxShapes=images:8x{C}x{H}x{W}"]) + return onnx_path + + +def main() -> int: + args = parse_args() + wanted = {c.strip() for c in args.components.split(",") if c.strip()} + # The memory block's QFormer takes PyTorch's fused MHA fast path, and + # aten::_transformer_encoder_layer_fwd has no ONNX symbolic. Same toggle the BF16 + # export uses. + try: + torch.backends.mha.set_fastpath_enabled(False) + except Exception as exc: # pragma: no cover + print(f" (mha fastpath toggle unavailable: {exc})") + os.makedirs(args.out_dir, exist_ok=True) + internvla_compat.apply_all(need_system1=True, allow_missing_depth=False) + + from internnav.model.basemodel.internvla_n1.internvla_n1 import ( + InternVLAN1ForCausalLM, InternVLAN1ModelConfig) + + calib = torch.load(args.calib_path, map_location="cpu", weights_only=False) + print(f"Loading {args.internvla_ckpt}") + config = InternVLAN1ModelConfig.from_pretrained(args.internvla_ckpt) + model = InternVLAN1ForCausalLM.from_pretrained( + args.internvla_ckpt, config=config, torch_dtype=torch.bfloat16, + attn_implementation="sdpa", low_cpu_mem_usage=True).to(args.device).eval() + inner = model.get_model() + + if "traj_dit" in wanted: + quantize_traj_dit(inner, calib, args) + if "memory" in wanted: + quantize_memory(model, inner, calib, args) + print("\nDone. Verify with verify/compare_system1_engines.py --engine_dir " + f"{args.out_dir}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/compare_system1_engines.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/compare_system1_engines.py index f64eea6..77ef552 100644 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/compare_system1_engines.py +++ b/recipes/internvla-n1-dualvln/trt-edgellm/verify/compare_system1_engines.py @@ -65,8 +65,9 @@ def parse_args() -> argparse.Namespace: formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument("--reference_path", required=True) p.add_argument("--engine_dir", required=True, - help="Directory holding system1_traj_dit_bf16.engine and " - "system1_memory_bf16.engine") + help="Directory holding the two System-1 engines") + p.add_argument("--engine_suffix", default="bf16", + help="Precision tag in the engine filenames: system1_traj_dit_.engine") p.add_argument("--guidance_scale", type=float, default=1.0) p.add_argument("--device", default="cuda") p.add_argument("--loop_dtype", default="bfloat16", @@ -86,8 +87,9 @@ def main() -> int: ref = torch.load(args.reference_path, map_location="cpu", weights_only=False) dev = args.device - mem_path = os.path.join(args.engine_dir, "system1_memory_bf16.engine") - dit_path = os.path.join(args.engine_dir, "system1_traj_dit_bf16.engine") + tag = args.engine_suffix + mem_path = os.path.join(args.engine_dir, f"system1_memory_{tag}.engine") + dit_path = os.path.join(args.engine_dir, f"system1_traj_dit_{tag}.engine") for path in (mem_path, dit_path): if not os.path.isfile(path): print(f"[ERROR] engine not found: {path}") @@ -97,6 +99,12 @@ def main() -> int: mem = Engine(mem_path) # The memory engine was built from MemBlock, which takes [T, C, H, W]. images = ref["images_chw"].to(dev).float().contiguous() + # The FP8 memory engine is built with a dynamic frame count; the BF16 one is static. + # Setting the shape is a no-op for the static engine and required for the dynamic one. + try: + mem.set_runtime_tensor_shape("images", tuple(images.shape)) + except Exception: + pass tokens = out_of(mem(images=images), "memory_tokens").float().cpu() timing = {} if args.bench_iters: @@ -204,11 +212,19 @@ def run_loop(latents): traj_cos = cos(ref_traj, traj) l2 = float((traj - ref_traj).norm() / ref_traj.norm()) + # Cosine says how aligned the trajectories are; it does not say whether a robot would + # end up somewhere else. Per-waypoint Euclidean deviation is in the trajectory's own + # units (metres) and is the number to judge a scheme on. + dev = (traj - ref_traj).norm(dim=-1) + reach = ref_traj.norm(dim=-1).mean() print("=" * 58) print(f" memory_tokens cosine : {mem_cos:.6f}") print(f" traj_dit single step : {step_cos:.6f}") print(f" trajectory cosine : {traj_cos:.6f}") print(f" trajectory rel-L2 : {l2:.4f}") + print(f" waypoint deviation : mean {dev.mean():.4f} / median " + f"{dev.median():.4f} / p95 {dev.flatten().quantile(0.95):.4f} " + f"(reference waypoint reach {reach:.4f})") ok = traj_cos >= args.gate if timing: print("\n --- latency, mean over " From 279057df1e52bc17e3c4877a8b0109a39f47c3c2 Mon Sep 17 00:00:00 2001 From: hungho77 Date: Fri, 14 Aug 2026 15:51:15 +0700 Subject: [PATCH 27/30] feat(internvla-n1-dualvln): add the six pipeline scripts the Makefile calls Six of the seven scripts the Makefile invokes did not exist, so every entry point failed at the first command -- including the README's own Quick Start. The Python underneath was fine; only the shell wrappers were missing. 00_fetch_calib_scenes.sh pull a diverse InternData-N1 subset (gated dataset) 01_repackage.sh InternVLA -> stock Qwen2.5-VL System 2 02_quantize.sh one scheme x strategy, through the validity gate 04_export_system1.sh traj_dit + memory block -> BF16 engines 05_verify.sh both acceptance gates benchmark retargeted at 06_measure.sh, which exists and takes the engine parent directory via environment, not flags 05_verify.sh --system1 runs the two-stage split, so it takes PYTHON_PT and PYTHON_TRT: InternNav wants transformers 4.x and the TensorRT bindings ship for Python 3.12 where transformers is 5.x, and nothing reconciles the two in one interpreter. Exercised rather than assumed. The System 2 gate reproduces the README's 0.991861 on the FP8 engine; the System 1 gate reports 0.999661 and exits 0; a missing checkpoint and nvfp4 x s3 both fail with their reason. Also fixes two README flags that never existed (--calib_data, --max_seq_len; the real one is --calib_data_root) and drops the unassigned-owner line. --- recipes/internvla-n1-dualvln/Makefile | 7 +- recipes/internvla-n1-dualvln/README.md | 11 ++- .../quantize/scripts/00_fetch_calib_scenes.sh | 84 +++++++++++++++++++ .../quantize/scripts/01_repackage.sh | 38 +++++++++ .../quantize/scripts/02_quantize.sh | 64 ++++++++++++++ .../trt-edgellm/scripts/04_export_system1.sh | 64 ++++++++++++++ .../trt-edgellm/scripts/05_verify.sh | 83 ++++++++++++++++++ 7 files changed, 345 insertions(+), 6 deletions(-) create mode 100755 recipes/internvla-n1-dualvln/quantize/scripts/00_fetch_calib_scenes.sh create mode 100755 recipes/internvla-n1-dualvln/quantize/scripts/01_repackage.sh create mode 100755 recipes/internvla-n1-dualvln/quantize/scripts/02_quantize.sh create mode 100755 recipes/internvla-n1-dualvln/trt-edgellm/scripts/04_export_system1.sh create mode 100755 recipes/internvla-n1-dualvln/trt-edgellm/scripts/05_verify.sh diff --git a/recipes/internvla-n1-dualvln/Makefile b/recipes/internvla-n1-dualvln/Makefile index 456ccbe..0354048 100644 --- a/recipes/internvla-n1-dualvln/Makefile +++ b/recipes/internvla-n1-dualvln/Makefile @@ -73,10 +73,11 @@ verify-latents: --repkg_ckpt $(REPKG_CKPT) \ --calib_data_root $(CALIB_DATA_ROOT) +# 06_measure.sh sweeps every engine under ENGINE_DIR, so it takes the parent directory +# rather than one engine, and is configured by environment rather than flags. benchmark: - bash trt-edgellm/scripts/06_benchmark.sh \ - --engine_dir $(ENGINE_DIR)/$(STRATEGY)_$(SCHEME) \ - --work_dir $(WORK_DIR) + WORK_DIR=$(WORK_DIR) ENGINE_DIR=$(ENGINE_DIR) TRT_EDGELLM_DIR=$(TRT_EDGELLM_DIR) \ + bash trt-edgellm/scripts/06_measure.sh # ── System 1: BF16 engines (requires INTERNNAV_PATH) ────────────────────────── .PHONY: export-system1 verify-system1 quantize-system1-fp8 diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index 6138888..8d504a8 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -5,7 +5,6 @@ vision-language navigation model, targeting NVIDIA Jetson Thor. ## Status -- Owner: unassigned - Model family: `InternVLA-N1-DualVLN` (System 2 = Qwen2.5-VL-7B, System 1 = NextDiT trajectory head) - Quantization presets: `fp8_default`, `fp8_per_channel` (validated) · `nvfp4_*` (experimental, see Notes) - Runtime target: TensorRT-Edge-LLM engines on Jetson Thor (sm_110) @@ -478,6 +477,13 @@ diagnosable, not so you set them by hand: make export-build # ONNX export + FP8 LLM engine + visual engine make verify-latents # the acceptance gate: z_latents cosine > 0.99 +System 1 needs InternNav and its own two interpreters, since no single environment has both +InternNav (transformers 4.x) and the TensorRT bindings (Python 3.12): + + make export-system1 + PYTHON_PT=/path/to/py310/bin/python PYTHON_TRT=/path/to/py312/bin/python \ + make verify-system1 + `make help` lists every target. Each script is also runnable directly; see the per-path READMEs in `quantize/` and `trt-edgellm/`. @@ -490,9 +496,8 @@ READMEs in `quantize/` and `trt-edgellm/`. --strategy {s1,s2,s3,s4} s1 LLM · s2 +KV cache · s3 +ViT · s4 +ViT+KV (default: s1) --scheme NAME Preset from configs/schemes.yaml (default: fp8_default) --calib {auto,text,multimodal,vln} Calibration source (default: auto) - --calib_data PATH Root for VLN calibration scenes + --calib_data_root PATH Root for VLN calibration scenes --num_calib_samples N Calibration samples (default: 512; image paths cap at 128) - --max_seq_len N Calibration truncation length (default: 512) --dtype {fp16,bf16} Load dtype (default: bf16) --device DEVICE Torch device (default: cuda) --resume DIR Layerwise checkpoint dir, for AWQ/Hessian crash recovery diff --git a/recipes/internvla-n1-dualvln/quantize/scripts/00_fetch_calib_scenes.sh b/recipes/internvla-n1-dualvln/quantize/scripts/00_fetch_calib_scenes.sh new file mode 100755 index 0000000..566f898 --- /dev/null +++ b/recipes/internvla-n1-dualvln/quantize/scripts/00_fetch_calib_scenes.sh @@ -0,0 +1,84 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +# +# Fetch a small, diverse subset of InternData-N1 VLN-CE scenes for calibration. +# +# InternData-N1 is gated on HuggingFace and the full VLN-CE traj_data is ~2.5 TB across 914 +# scenes, so this pulls only a handful of per-scene archives. Accept the dataset terms and +# run `huggingface-cli login` first, or every download 401s. +# +# The default set spans r2r + rxr + scalevln -- the same training mix the base model saw -- +# and is deliberately made of small scenes to keep disk in check. Override SCENES to pick +# others. +# +# Note one of these scenes, YmJkqBEsHnH, also appears in the held-out probe set used by +# quantize/benchmark_accuracy.py. Exclude it from either side before reading a number as +# out-of-sample; qat.py already does. + +set -euo pipefail + +OUTPUT_PATH="${CALIB_DATA_ROOT:-$HOME/vln-opt-work/calib_scenes}" +SCENES="${SCENES:-\ +vln_ce/traj_data/r2r/gZ6f7yhEvPG.tar.gz \ +vln_ce/traj_data/r2r/YmJkqBEsHnH.tar.gz \ +vln_ce/traj_data/r2r/XcA2TqTSSAj.tar.gz \ +vln_ce/traj_data/rxr/Pm6F8kyY3z2.tar.gz \ +vln_ce/traj_data/rxr/PuKPg4mmafe.tar.gz \ +vln_ce/traj_data/scalevln/00493-pUneSGJDrvY.tar.gz \ +vln_ce/traj_data/scalevln/00446-tL6i2PtktSh.tar.gz \ +vln_ce/traj_data/scalevln/00351-QxfX5te1gFu.tar.gz \ +vln_ce/traj_data/scalevln/00335-janiYDpzM9j.tar.gz}" + +while [[ $# -gt 0 ]]; do + case "$1" in + --output_path) OUTPUT_PATH="$2"; shift 2 ;; + --scenes) SCENES="$2"; shift 2 ;; + -h|--help) sed -n '2,20p' "$0"; exit 0 ;; + *) echo "[ERROR] unknown argument: $1" >&2; exit 2 ;; + esac +done + +mkdir -p "$OUTPUT_PATH" +export INTERNDATA_DEST="$OUTPUT_PATH" INTERNDATA_SCENES="$SCENES" + +python - <<'PY' +import os +import shutil +import tarfile +import tempfile + +from huggingface_hub import hf_hub_download + +REPO = "InternRobotics/InternData-N1" +dest = os.environ["INTERNDATA_DEST"] +scenes = os.environ["INTERNDATA_SCENES"].split() + +# Download into a scratch dir and delete each archive right after extraction, so the HF blob +# cache never accumulates a second copy. local_dir avoids the shared cache entirely. +scratch = tempfile.mkdtemp(prefix="interndata_") +try: + for rel in scenes: + scene = os.path.basename(rel)[:-len(".tar.gz")] + subset = rel.split("/")[2] # r2r | rxr | scalevln + out_dir = os.path.join(dest, subset, scene) + if os.path.isdir(os.path.join(out_dir, "meta")): + print(f"[skip] already extracted: {out_dir}") + continue + print(f"[get ] {rel}") + tar_path = hf_hub_download(REPO, rel, repo_type="dataset", + local_dir=scratch, local_dir_use_symlinks=False) + with tarfile.open(tar_path) as tf: + tf.extractall(os.path.join(dest, subset)) + os.remove(tar_path) + print(f"[ok ] {scene}") +finally: + shutil.rmtree(scratch, ignore_errors=True) + +found = sum(1 for root, _, files in os.walk(dest) + if root.endswith(os.sep + "meta") and "episodes.jsonl" in files) +print(f"[done] {found} scene(s) with episodes.jsonl under {dest}") +PY + +echo "Calibration data ready: $OUTPUT_PATH" +echo "Pass it as --calib_data_root, or set CALIB_DATA_ROOT." diff --git a/recipes/internvla-n1-dualvln/quantize/scripts/01_repackage.sh b/recipes/internvla-n1-dualvln/quantize/scripts/01_repackage.sh new file mode 100755 index 0000000..38499bd --- /dev/null +++ b/recipes/internvla-n1-dualvln/quantize/scripts/01_repackage.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +# +# Strip System 1 out of the InternVLA checkpoint, leaving a stock Qwen2.5-VL System 2. +# +# This is the step that makes the rest of the recipe ordinary: after it, quantize, export, +# build and the bridge verification all operate on a plain Qwen2.5-VL checkpoint and never +# import InternNav. It is pure safetensors manipulation -- no model is constructed -- so it +# runs anywhere, including without a GPU. +# +# Costs one ~15 GB intermediate copy. Pass --free_source to delete each source shard as it +# is consumed if disk is tight; the peak then is one shard rather than two checkpoints. + +set -euo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +MODEL_PATH="${INTERNVLA_CKPT:-$HOME/InternNav/checkpoints/InternVLA-N1-DualVLN}" +OUTPUT_PATH="${REPKG_CKPT:-$HOME/vln-opt-work/qwen25vl_system2}" +EXTRA=() + +while [[ $# -gt 0 ]]; do + case "$1" in + --model_path) MODEL_PATH="$2"; shift 2 ;; + --output_path) OUTPUT_PATH="$2"; shift 2 ;; + --free_source) EXTRA+=(--free_source); shift ;; + --skip_disk_check) EXTRA+=(--skip_disk_check); shift ;; + -h|--help) sed -n '2,14p' "$0"; exit 0 ;; + *) echo "[ERROR] unknown argument: $1" >&2; exit 2 ;; + esac +done + +[[ -d "$MODEL_PATH" ]] || { echo "[ERROR] checkpoint not found: $MODEL_PATH" >&2; exit 1; } + +exec python -u "$HERE/../repackage_system2.py" \ + --model_path "$MODEL_PATH" \ + --output_path "$OUTPUT_PATH" \ + "${EXTRA[@]}" diff --git a/recipes/internvla-n1-dualvln/quantize/scripts/02_quantize.sh b/recipes/internvla-n1-dualvln/quantize/scripts/02_quantize.sh new file mode 100755 index 0000000..b6eb2bd --- /dev/null +++ b/recipes/internvla-n1-dualvln/quantize/scripts/02_quantize.sh @@ -0,0 +1,64 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +# +# Quantize the repackaged System 2 with one scheme x strategy combination. +# +# Input is the *repackaged* checkpoint, not the InternVLA one -- see 01_repackage.sh. +# The scheme/strategy validity gate lives in quantize/quant_schemes.py and rejects the +# impossible combinations early, with the reason; nvfp4 x {s1,s2} additionally needs +# --allow_experimental because it passes text fluency and still breaks the navigation +# bridge (z_latents 0.931 against a 0.99 gate). +# +# Calibration defaults to 'auto' -- cnn_dailymail text for LLM-only strategies -- matching +# quantize.py. Pass --calib vln to use navigation episodes instead; it needs the scenes from +# 00_fetch_calib_scenes.sh and, measured here, changes held-out z_latents by 0.00003, so +# prefer it for honesty rather than for accuracy. +# +# Pass --dry_run to load the model and validate the configuration without quantizing. + +set -euo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +MODEL_PATH="${REPKG_CKPT:-$HOME/vln-opt-work/qwen25vl_system2}" +OUTPUT_PATH="" +SCHEME="${SCHEME:-fp8_default}" +STRATEGY="${STRATEGY:-s1}" +DEVICE="${DEVICE:-cuda}" +CALIB="${CALIB:-auto}" +CALIB_DATA_ROOT="${CALIB_DATA_ROOT:-$HOME/vln-opt-work/calib_scenes}" +EXTRA=() + +while [[ $# -gt 0 ]]; do + case "$1" in + --model_path) MODEL_PATH="$2"; shift 2 ;; + --output_path) OUTPUT_PATH="$2"; shift 2 ;; + --scheme) SCHEME="$2"; shift 2 ;; + --strategy) STRATEGY="$2"; shift 2 ;; + --device) DEVICE="$2"; shift 2 ;; + --calib) CALIB="$2"; shift 2 ;; + --calib_data_root) CALIB_DATA_ROOT="$2"; shift 2 ;; + --num_calib_samples) EXTRA+=(--num_calib_samples "$2"); shift 2 ;; + --allow_experimental) EXTRA+=(--allow_experimental); shift ;; + --dry_run) EXTRA+=(--dry_run); shift ;; + -h|--help) sed -n "2,19p" "$0"; exit 0 ;; + *) echo "[ERROR] unknown argument: $1" >&2; exit 2 ;; + esac +done + +OUTPUT_PATH="${OUTPUT_PATH:-$HOME/vln-opt-work/qwen25vl_${STRATEGY}_${SCHEME}}" +[[ -d "$MODEL_PATH" ]] || { echo "[ERROR] checkpoint not found: $MODEL_PATH" >&2; exit 1; } + +# ModelOpt's Triton kernels are compiled in-tree on Thor; without this the quantize pass +# fails at import time rather than at use. +export TRITON_BACKENDS_IN_TREE="${TRITON_BACKENDS_IN_TREE:-1}" + +exec python -u "$HERE/../quantize.py" \ + --model_path "$MODEL_PATH" \ + --output_path "$OUTPUT_PATH" \ + --scheme "$SCHEME" \ + --strategy "$STRATEGY" \ + --device "$DEVICE" \ + --calib "$CALIB" \ + --calib_data_root "$CALIB_DATA_ROOT" \ + "${EXTRA[@]}" diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/scripts/04_export_system1.sh b/recipes/internvla-n1-dualvln/trt-edgellm/scripts/04_export_system1.sh new file mode 100755 index 0000000..1b9e1bf --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/scripts/04_export_system1.sh @@ -0,0 +1,64 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +# +# Export System 1 -- the NextDiT diffusion head and the memory block -- to BF16 engines. +# +# Upstream InternNav ships no ONNX or TensorRT path at all, so this is entirely the +# recipe's. Both exporters need InternNav importable, unlike everything under System 2. +# +# System 1 stays BF16 on purpose. FP8 was measured (trt-edgellm/quantize_system1.py) and +# costs 6x the waypoint deviation to save 0.7% of deployed weights and 1.7% of a planning +# step -- see the README. +# +# The two exporters write into $WORK_DIR/onnx, which is also where verify_system1.py looks; +# --engine_dir is linked to those files rather than holding copies, so both paths stay valid +# without a second 176 MB on disk. + +set -euo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ROOT="$(dirname "$HERE")" +INTERNVLA_CKPT="${INTERNVLA_CKPT:-$HOME/InternNav/checkpoints/InternVLA-N1-DualVLN}" +INTERNNAV_PATH="${INTERNNAV_PATH:-$HOME/InternNav}" +WORK_DIR="${WORK_DIR:-$HOME/vln-opt-work}" +ENGINE_DIR="" +COMPONENTS="traj_dit,memory" + +while [[ $# -gt 0 ]]; do + case "$1" in + --internvla_ckpt) INTERNVLA_CKPT="$2"; shift 2 ;; + --internnav_path) INTERNNAV_PATH="$2"; shift 2 ;; + --work_dir) WORK_DIR="$2"; shift 2 ;; + --engine_dir) ENGINE_DIR="$2"; shift 2 ;; + --components) COMPONENTS="$2"; shift 2 ;; + -h|--help) sed -n '2,17p' "$0"; exit 0 ;; + *) echo "[ERROR] unknown argument: $1" >&2; exit 2 ;; + esac +done + +[[ -d "$INTERNNAV_PATH" ]] || { echo "[ERROR] InternNav not found: $INTERNNAV_PATH" >&2; exit 1; } +[[ -d "$INTERNVLA_CKPT" ]] || { echo "[ERROR] checkpoint not found: $INTERNVLA_CKPT" >&2; exit 1; } + +export INTERNNAV_PATH INTERNVLA_CKPT WORK_DIR +# The exporters import both InternNav and the recipe's own modules. +export PYTHONPATH="$INTERNNAV_PATH:$ROOT${PYTHONPATH:+:$PYTHONPATH}" + +if [[ ",$COMPONENTS," == *",traj_dit,"* ]]; then + echo "== traj_dit ==" + python -u "$ROOT/export_traj_dit.py" +fi +if [[ ",$COMPONENTS," == *",memory,"* ]]; then + echo "== memory block ==" + python -u "$ROOT/export_memory_block.py" +fi + +if [[ -n "$ENGINE_DIR" ]]; then + mkdir -p "$ENGINE_DIR" + for f in "$WORK_DIR"/onnx/system1_*.engine; do + [[ -e "$f" ]] || continue + ln -sfn "$f" "$ENGINE_DIR/$(basename "$f")" + done + echo "Engines linked into: $ENGINE_DIR" +fi +ls -lh "$WORK_DIR"/onnx/system1_*.engine 2>/dev/null || true diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/scripts/05_verify.sh b/recipes/internvla-n1-dualvln/trt-edgellm/scripts/05_verify.sh new file mode 100755 index 0000000..43e1054 --- /dev/null +++ b/recipes/internvla-n1-dualvln/trt-edgellm/scripts/05_verify.sh @@ -0,0 +1,83 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +# +# Run the acceptance gates. +# +# Default (System 2): the z_latents bridge -- the last-layer hidden states of the 4 TRAJ +# tokens through the host-side norm and cond_projector, engine against a PyTorch reference. +# This is the number that decides whether a scheme ships. Text fluency is not sufficient +# evidence: NVFP4 stays fluent and fails this at 0.931. +# +# --system1: trajectory parity for the diffusion head and memory block. This one runs in TWO +# stages under TWO interpreters, and that is not incidental. InternNav targets transformers +# 4.x while the TensorRT bindings ship for Python 3.12 where transformers is 5.x, so no +# single environment has both. Stage A writes the PyTorch reference's inputs and outputs to +# a .pt; stage B feeds the engines those same tensors. Set PYTHON_PT and PYTHON_TRT to the +# two interpreters -- if they are left at the default the stages run in whatever is active, +# which works only if one environment happens to satisfy both. + +set -euo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ROOT="$(dirname "$HERE")" +SYSTEM1=0 +ENGINE_DIR="" +ENGINE_SUFFIX="${ENGINE_SUFFIX:-bf16}" +REPKG_CKPT="${REPKG_CKPT:-$HOME/vln-opt-work/qwen25vl_system2}" +CALIB_DATA_ROOT="${CALIB_DATA_ROOT:-$HOME/vln-opt-work/calib_scenes}" +INTERNVLA_CKPT="${INTERNVLA_CKPT:-$HOME/InternNav/checkpoints/InternVLA-N1-DualVLN}" +INTERNNAV_PATH="${INTERNNAV_PATH:-$HOME/InternNav}" +WORK_DIR="${WORK_DIR:-$HOME/vln-opt-work}" +VLN=0 +BENCH_ITERS="${BENCH_ITERS:-0}" +# The System 2 gate needs transformers and the TensorRT bindings in one interpreter; the +# System 1 gate cannot have both and splits across PYTHON_PT / PYTHON_TRT below. +PYTHON="${PYTHON:-python}" + +while [[ $# -gt 0 ]]; do + case "$1" in + --system1) SYSTEM1=1; shift ;; + --vln) VLN=1; shift ;; + --engine_dir) ENGINE_DIR="$2"; shift 2 ;; + --engine_suffix) ENGINE_SUFFIX="$2"; shift 2 ;; + --repkg_ckpt) REPKG_CKPT="$2"; shift 2 ;; + --calib_data_root) CALIB_DATA_ROOT="$2"; shift 2 ;; + --internvla_ckpt) INTERNVLA_CKPT="$2"; shift 2 ;; + --internnav_path) INTERNNAV_PATH="$2"; shift 2 ;; + --work_dir) WORK_DIR="$2"; shift 2 ;; + --bench_iters) BENCH_ITERS="$2"; shift 2 ;; + -h|--help) sed -n '2,19p' "$0"; exit 0 ;; + *) echo "[ERROR] unknown argument: $1" >&2; exit 2 ;; + esac +done + +export WORK_DIR REPKG_CKPT INTERNNAV_PATH INTERNVLA_CKPT CALIB_DATA_ROOT +export PYTHONPATH="$ROOT${PYTHONPATH:+:$PYTHONPATH}" +export EDGELLM_PLUGIN_PATH="${EDGELLM_PLUGIN_PATH:-$HOME/modelopt/Hung-TRT-Edge-LLM/build/libNvInfer_edgellm_plugin.so}" + +if [[ "$SYSTEM1" -eq 1 ]]; then + : "${ENGINE_DIR:=$WORK_DIR/onnx}" + PYTHON_PT="${PYTHON_PT:-python}" + PYTHON_TRT="${PYTHON_TRT:-python}" + REF="${REFERENCE_PATH:-$WORK_DIR/out/system1_reference.pt}" + + echo "== stage A: PyTorch reference (needs InternNav; $PYTHON_PT) ==" + PYTHONPATH="$INTERNNAV_PATH:$PYTHONPATH" "$PYTHON_PT" -u "$ROOT/verify/dump_system1_reference.py" \ + --internvla_ckpt "$INTERNVLA_CKPT" --output_path "$REF" + + echo "== stage B: engines (needs TensorRT; $PYTHON_TRT) ==" + ARGS=(--reference_path "$REF" --engine_dir "$ENGINE_DIR" --engine_suffix "$ENGINE_SUFFIX") + [[ "$BENCH_ITERS" -gt 0 ]] && ARGS+=(--bench_iters "$BENCH_ITERS") + exec "$PYTHON_TRT" -u "$ROOT/verify/compare_system1_engines.py" "${ARGS[@]}" +fi + +[[ -n "$ENGINE_DIR" ]] || { echo "[ERROR] --engine_dir is required" >&2; exit 2; } +ENGINE_PATH="${ENGINE_PATH:-$ENGINE_DIR/llm/llm.engine}" +[[ -f "$ENGINE_PATH" ]] || { echo "[ERROR] engine not found: $ENGINE_PATH" >&2; exit 1; } +export ENGINE_PATH + +if [[ "$VLN" -eq 1 ]]; then + exec "$PYTHON" -u "$ROOT/verify/verify_latents_vln.py" +fi +exec "$PYTHON" -u "$ROOT/verify/verify_latents.py" From 119384cc46f7ef7087110c8c805c176b94d3015a Mon Sep 17 00:00:00 2001 From: hungho77 Date: Tue, 25 Aug 2026 00:22:39 +0700 Subject: [PATCH 28/30] docs(internvla-n1-dualvln): point to the native TensorRT-Edge-LLM support NVIDIA/TensorRT-Edge-LLM#193 exports InternVLA-N1-DualVLN directly, no repackage step, bridge folded into the graph, plus an async C++ runtime and 199-episode closed-loop SR this recipe never measured. Also flags that the recipe's central claim -- z_latents cosine as the acceptance metric -- does not hold. Closed-loop SR showed neither z_latents cosine nor trajectory cosine predicts navigation success for this model, in either direction; they rank this recipe's own FP8 vs NVFP4 call the wrong way. The FP8-on-System-1 measurement and the NVFP4 root-cause analysis are unaffected and still hold. --- recipes/internvla-n1-dualvln/README.md | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index 8d504a8..770a889 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -3,6 +3,28 @@ FP8 quantization and TensorRT-Edge-LLM deployment for InternVLA-N1-DualVLN, a dual-system vision-language navigation model, targeting NVIDIA Jetson Thor. +> **Superseded by native support.** The repackage-and-host-side-bridge approach this recipe +> uses has been replaced by direct support in TensorRT-Edge-LLM itself: +> [NVIDIA/TensorRT-Edge-LLM#193](https://github.com/NVIDIA/TensorRT-Edge-LLM/pull/193) +> (pending review, tracked by +> [NVIDIA/TensorRT-Edge-LLM#190](https://github.com/NVIDIA/TensorRT-Edge-LLM/issues/190)). +> That work exports `InternVLA-N1-DualVLN` directly from the released checkpoint — no +> `repackage_system2.py` step, no `model_type` rewrite — and folds the `z_latents` bridge +> (`final_norm` + `cond_projector`) into the exported graph, so the engine emits `z_latents` +> on its own instead of a host-side Python step computing it. It also ships an async C++ +> runtime (`internvla_n1_dual_system_inference` / `internvla_n1_dual_system_server`) running +> both systems in one process, and 199-episode closed-loop navigation SR on Jetson Thor — +> validation this recipe never had. +> +> It also corrects this recipe's central claim below. **"The bridge is the acceptance +> metric, not text quality" is not true for this model.** Measured closed-loop, `z_latents` +> cosine and trajectory cosine each rank a different quantization scheme wrong, in opposite +> directions — including ranking this recipe's own FP8 recommendation *below* NVFP4, which +> this recipe's Notes section rules out on cosine grounds. The only metric that agreed with +> itself across schemes was closed-loop success rate. Kept below for the FP8-on-System-1 +> measurement and the `investigate_nvfp4.py` root-cause analysis, both still accurate; treat +> everything under "Results" and "Notes" that reasons from `z_latents` cosine as superseded. + ## Status - Model family: `InternVLA-N1-DualVLN` (System 2 = Qwen2.5-VL-7B, System 1 = NextDiT trajectory head) From 19d9d5f9e84fc87ce9041f4b89075c0647ad88fc Mon Sep 17 00:00:00 2001 From: hungho77 Date: Tue, 25 Aug 2026 01:00:37 +0700 Subject: [PATCH 29/30] feat(internvla-n1-dualvln): add TensorRT-Edge-LLM quantize and run instructions (PR #193) NVIDIA/TensorRT-Edge-LLM#193 (open, not yet merged) adds native InternVLA-N1-DualVLN export -- direct checkpoint export, the z_latents bridge (final_norm + cond_projector) folded into the graph -- so once building from that branch, the repackage pass, host-side bridge computation, and custom export/quantize scripts this recipe used to carry are no longer needed. What is left as this recipe's job is the one piece that stays outside TensorRT-Edge-LLM regardless: building a navigation-domain calibration set, since calibrating on the deployment prompt's own domain rather than generic news text measurably improves quantization quality (FP8 trajectory cosine 0.909 -> 0.978 in earlier testing). Replaces the old repackage/quantize/export/verify script tree (45 files) with a single build_calib_jsonl.py plus a two-target Makefile, and rewrites the README with quantize + export + build instructions for PR #193's branch, plus the 199-episode closed-loop SR results, which supersede the recipe's earlier z_latents-cosine-based FP8-over-NVFP4 recommendation. --- pyproject.toml | 23 +- recipes/internvla-n1-dualvln/Makefile | 170 +---- recipes/internvla-n1-dualvln/README.md | 649 ++---------------- .../internvla-n1-dualvln/configs/schemes.yaml | 109 --- .../quantize/benchmark_accuracy.py | 259 ------- .../quantize/build_calib_jsonl.py | 79 +++ .../quantize/calibration.py | 421 ------------ .../quantize/compare_fake_quant.py | 120 ---- .../quantize/load_quantized.py | 218 ------ .../quantize/model_loader.py | 290 -------- .../quantize/prompt_builder.py | 237 ------- recipes/internvla-n1-dualvln/quantize/qat.py | 342 --------- .../quantize/quant_schemes.py | 224 ------ .../internvla-n1-dualvln/quantize/quantize.py | 323 --------- .../quantize/repackage_system2.py | 213 ------ .../quantize/scripts/00_fetch_calib_scenes.sh | 84 --- .../quantize/scripts/01_repackage.sh | 38 - .../quantize/scripts/02_quantize.sh | 64 -- .../requirements-torch.txt | 31 - recipes/internvla-n1-dualvln/run_matrix.py | 124 ---- .../trt-edgellm/benchmark/bench_memory.py | 71 -- .../trt-edgellm/benchmark/bench_system1.py | 77 --- .../benchmark/benchmark_system2.py | 159 ----- .../trt-edgellm/deploy/run_eval_engine.py | 74 -- .../trt-edgellm/diagnose_engine_gap.py | 130 ---- .../trt-edgellm/dump_system1_calib.py | 142 ---- .../trt-edgellm/engine_policy.py | 119 ---- .../trt-edgellm/engine_runner.py | 235 ------- .../trt-edgellm/export_memory_block.py | 123 ---- .../trt-edgellm/export_traj_dit.py | 70 -- .../trt-edgellm/internvla_compat.py | 216 ------ .../trt-edgellm/investigate_nvfp4.py | 310 --------- .../trt-edgellm/memblock.py | 25 - .../trt-edgellm/quantize_system1.py | 262 ------- .../scripts/03_export_build_system2.sh | 145 ---- .../trt-edgellm/scripts/04_export_system1.sh | 64 -- .../trt-edgellm/scripts/05_verify.sh | 83 --- .../trt-edgellm/scripts/06_measure.sh | 55 -- .../trt-edgellm/traj_dit_loader.py | 109 --- .../trt-edgellm/trt_torch.py | 230 ------- .../verify/compare_system1_engines.py | 241 ------- .../verify/dump_system1_reference.py | 183 ----- .../trt-edgellm/verify/verify_accuracy.py | 161 ----- .../trt-edgellm/verify/verify_e2e_agent.py | 220 ------ .../verify/verify_engine_policy.py | 142 ---- .../trt-edgellm/verify/verify_latents.py | 202 ------ .../trt-edgellm/verify/verify_latents_vln.py | 219 ------ .../trt-edgellm/verify/verify_pixelgoal_gt.py | 217 ------ .../trt-edgellm/verify/verify_system1.py | 142 ---- 49 files changed, 192 insertions(+), 8252 deletions(-) delete mode 100644 recipes/internvla-n1-dualvln/configs/schemes.yaml delete mode 100644 recipes/internvla-n1-dualvln/quantize/benchmark_accuracy.py create mode 100644 recipes/internvla-n1-dualvln/quantize/build_calib_jsonl.py delete mode 100644 recipes/internvla-n1-dualvln/quantize/calibration.py delete mode 100644 recipes/internvla-n1-dualvln/quantize/compare_fake_quant.py delete mode 100644 recipes/internvla-n1-dualvln/quantize/load_quantized.py delete mode 100644 recipes/internvla-n1-dualvln/quantize/model_loader.py delete mode 100644 recipes/internvla-n1-dualvln/quantize/prompt_builder.py delete mode 100644 recipes/internvla-n1-dualvln/quantize/qat.py delete mode 100644 recipes/internvla-n1-dualvln/quantize/quant_schemes.py delete mode 100755 recipes/internvla-n1-dualvln/quantize/quantize.py delete mode 100644 recipes/internvla-n1-dualvln/quantize/repackage_system2.py delete mode 100755 recipes/internvla-n1-dualvln/quantize/scripts/00_fetch_calib_scenes.sh delete mode 100755 recipes/internvla-n1-dualvln/quantize/scripts/01_repackage.sh delete mode 100755 recipes/internvla-n1-dualvln/quantize/scripts/02_quantize.sh delete mode 100644 recipes/internvla-n1-dualvln/requirements-torch.txt delete mode 100644 recipes/internvla-n1-dualvln/run_matrix.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_memory.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_system1.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/benchmark/benchmark_system2.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/deploy/run_eval_engine.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/diagnose_engine_gap.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/dump_system1_calib.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/engine_policy.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/engine_runner.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/export_traj_dit.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/internvla_compat.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/investigate_nvfp4.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/memblock.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/quantize_system1.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/scripts/03_export_build_system2.sh delete mode 100755 recipes/internvla-n1-dualvln/trt-edgellm/scripts/04_export_system1.sh delete mode 100755 recipes/internvla-n1-dualvln/trt-edgellm/scripts/05_verify.sh delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/scripts/06_measure.sh delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/traj_dit_loader.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/trt_torch.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/verify/compare_system1_engines.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/verify/dump_system1_reference.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_accuracy.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_e2e_agent.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_engine_policy.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents_vln.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_pixelgoal_gt.py delete mode 100644 recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_system1.py diff --git a/pyproject.toml b/pyproject.toml index 8cd8e28..ae713e0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -83,27 +83,12 @@ cosmos-reason2 = [ "tqdm==4.67.3", ] -# numpy and scipy are deliberately omitted here. This recipe requires numpy 1.x -# (OpenCV and diffusers break under numpy 2 on Jetson) while three other extras pin -# numpy==2.2.6, so rather than assert a resolution outcome we install them as a -# documented post-sync step; see recipes/internvla-n1-dualvln/requirements-torch.txt. +# `build_calib_jsonl.py` is stdlib-only (argparse, gzip, json, random). The one real +# dependency is the `huggingface-cli` entrypoint used to fetch the calibration source data. +# Quantize/export/build run against TensorRT-Edge-LLM's own environment, not this one -- +# see recipes/internvla-n1-dualvln/README.md. internvla-n1-dualvln = [ - "torch>=2.9.0", - "torchvision>=0.24.0", - "transformers==4.51.3", - "diffusers==0.33.1", - "accelerate==1.13.0", - "safetensors>=0.8.0", "huggingface-hub>=0.36.0", - "nvidia-modelopt>=0.44.0", - "onnx==1.22.0", - "onnxscript==0.7.1", - "onnx-graphsurgeon==0.6.1", - "pandas>=2.0.0", - "pillow>=10.0.0", - "numpy-quaternion>=2023.0.0", - "pyyaml>=6.0.0", - "tqdm>=4.66.0", ] [tool.ruff] diff --git a/recipes/internvla-n1-dualvln/Makefile b/recipes/internvla-n1-dualvln/Makefile index 0354048..4c1b340 100644 --- a/recipes/internvla-n1-dualvln/Makefile +++ b/recipes/internvla-n1-dualvln/Makefile @@ -1,149 +1,53 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics # SPDX-License-Identifier: BSD-3-Clause -# Makefile — Top-level command panel for internvla-n1-dualvln +# Makefile — internvla-n1-dualvln # -# Override paths via environment variables, e.g.: -# make quantize-fp8 INTERNVLA_CKPT=/data/InternVLA-N1-DualVLN WORK_DIR=/mnt/scratch -# -# Targets that run Python/bash assume the virtualenv is already activated. +# Builds the navigation-domain calibration set this recipe is responsible for. Quantize, +# export, and build are the standard TensorRT-Edge-LLM CLI, run in that repo's own +# environment -- see README.md for why they are not wrapped here. # -# Targets are grouped by dependency boundary. Everything under "System 2" runs -# without InternNav; System 1 and the agent-level checks require INTERNNAV_PATH. - -# ── Configurable paths ───────────────────────────────────────────────────────── -INTERNVLA_CKPT ?= $(HOME)/InternNav/checkpoints/InternVLA-N1-DualVLN -INTERNNAV_PATH ?= $(HOME)/InternNav -TRT_EDGELLM_DIR ?= $(HOME)/modelopt/TensorRT-Edge-LLM -WORK_DIR ?= $(HOME)/vln-opt-work -ENGINE_DIR ?= $(WORK_DIR)/engines -CALIB_DATA_ROOT ?= $(WORK_DIR)/calib_scenes -SCHEME ?= fp8_default -STRATEGY ?= s1 -DEVICE ?= cuda -REPO_ROOT := $(abspath ../..) - -REPKG_CKPT := $(WORK_DIR)/qwen25vl_system2 -QUANT_CKPT := $(WORK_DIR)/qwen25vl_$(STRATEGY)_$(SCHEME) - -# ── System 2: quantize (no InternNav required) ──────────────────────────────── -.PHONY: fetch-calib repackage quantize quantize-fp8 quantize-nvfp4 - -fetch-calib: - bash quantize/scripts/00_fetch_calib_scenes.sh \ - --output_path $(CALIB_DATA_ROOT) - -repackage: - bash quantize/scripts/01_repackage.sh \ - --model_path $(INTERNVLA_CKPT) \ - --output_path $(REPKG_CKPT) - -quantize: - bash quantize/scripts/02_quantize.sh \ - --model_path $(REPKG_CKPT) \ - --output_path $(QUANT_CKPT) \ - --scheme $(SCHEME) \ - --strategy $(STRATEGY) \ - --device $(DEVICE) - -quantize-fp8: - $(MAKE) quantize SCHEME=fp8_default STRATEGY=s1 - -quantize-nvfp4: - $(MAKE) quantize SCHEME=nvfp4_default STRATEGY=s1 - -# ── System 2: export, build, verify (no InternNav required) ──────────────────── -.PHONY: export-build export-build-base-fp16 verify-latents benchmark - -export-build: - bash trt-edgellm/scripts/03_export_build_system2.sh \ - --model_path $(QUANT_CKPT) \ - --engine_dir $(ENGINE_DIR)/$(STRATEGY)_$(SCHEME) \ - --trt_edgellm_dir $(TRT_EDGELLM_DIR) - -export-build-base-fp16: - bash trt-edgellm/scripts/03_export_build_system2.sh \ - --model_path $(REPKG_CKPT) \ - --engine_dir $(ENGINE_DIR)/base_fp16 \ - --trt_edgellm_dir $(TRT_EDGELLM_DIR) \ - --no_quantization - -verify-latents: - bash trt-edgellm/scripts/05_verify.sh \ - --engine_dir $(ENGINE_DIR)/$(STRATEGY)_$(SCHEME) \ - --repkg_ckpt $(REPKG_CKPT) \ - --calib_data_root $(CALIB_DATA_ROOT) - -# 06_measure.sh sweeps every engine under ENGINE_DIR, so it takes the parent directory -# rather than one engine, and is configured by environment rather than flags. -benchmark: - WORK_DIR=$(WORK_DIR) ENGINE_DIR=$(ENGINE_DIR) TRT_EDGELLM_DIR=$(TRT_EDGELLM_DIR) \ - bash trt-edgellm/scripts/06_measure.sh - -# ── System 1: BF16 engines (requires INTERNNAV_PATH) ────────────────────────── -.PHONY: export-system1 verify-system1 quantize-system1-fp8 - -export-system1: - bash trt-edgellm/scripts/04_export_system1.sh \ - --internvla_ckpt $(INTERNVLA_CKPT) \ - --internnav_path $(INTERNNAV_PATH) \ - --engine_dir $(ENGINE_DIR)/system1 +# Override paths via environment variables, e.g.: +# make build-calib CALIB_DATA_ROOT=/mnt/scratch/calib -# Measured and deliberately not shipped: FP8 costs 6x the waypoint deviation for 1.7% of -# the planning step. Kept so the finding stays reproducible. See README. -quantize-system1-fp8: - python trt-edgellm/dump_system1_calib.py \ - --calib_data_root $(CALIB_DATA_ROOT) \ - --output_path $(WORK_DIR)/system1_calib.pt - python trt-edgellm/quantize_system1.py \ - --calib_path $(WORK_DIR)/system1_calib.pt \ - --out_dir $(ENGINE_DIR)/system1_fp8 +INTERNVLA_CKPT ?= $(HOME)/InternNav/checkpoints/InternVLA-N1-DualVLN +CALIB_DATA_ROOT ?= $(HOME)/vln-opt-work/calib +TRAIN_JSON := $(CALIB_DATA_ROOT)/vln_ce/raw_data/r2r/train/train.json.gz +NAV_CALIB_JSONL := $(CALIB_DATA_ROOT)/nav_calib.jsonl +NUM_SAMPLES ?= 512 -verify-system1: - bash trt-edgellm/scripts/05_verify.sh \ - --system1 \ - --internvla_ckpt $(INTERNVLA_CKPT) \ - --internnav_path $(INTERNNAV_PATH) \ - --engine_dir $(ENGINE_DIR)/system1 +.PHONY: fetch-calib-source build-calib all clean help -# ── NVFP4 investigation ─────────────────────────────────────────────────────── -.PHONY: investigate-nvfp4 +fetch-calib-source: + huggingface-cli download InternRobotics/InternData-N1 \ + vln_ce/raw_data/r2r/train/train.json.gz --repo-type dataset \ + --local-dir $(CALIB_DATA_ROOT) -investigate-nvfp4: - python trt-edgellm/investigate_nvfp4.py \ - --repkg_ckpt $(REPKG_CKPT) \ - --calib_data_root $(CALIB_DATA_ROOT) \ - --work_dir $(WORK_DIR) +build-calib: $(TRAIN_JSON) + python quantize/build_calib_jsonl.py \ + --train_json $(TRAIN_JSON) \ + --output $(NAV_CALIB_JSONL) \ + --num_samples $(NUM_SAMPLES) -# ── Housekeeping ────────────────────────────────────────────────────────────── -.PHONY: clean-onnx help +all: fetch-calib-source build-calib + @echo "" + @echo "Calibration set ready: $(NAV_CALIB_JSONL)" + @echo "Quantize/export/build with the TensorRT-Edge-LLM CLI -- see README.md:" + @echo " export EDGELLM_QUANT_DATASET_CNN_DAILYMAIL=$(NAV_CALIB_JSONL)" + @echo " tensorrt-edgellm-quantize llm --model_dir $(INTERNVLA_CKPT) \\" + @echo " --output_dir --quantization {fp8,nvfp4}" -# ONNX is a pure intermediate and the largest reclaimable artifact (9-16 GB). -clean-onnx: - rm -rf $(WORK_DIR)/onnx - @echo "Removed $(WORK_DIR)/onnx" +clean: + rm -rf $(CALIB_DATA_ROOT) help: @echo "internvla-n1-dualvln — targets" @echo "" - @echo " System 2 (no InternNav needed):" - @echo " fetch-calib download VLN calibration scenes" - @echo " repackage InternVLA checkpoint -> stock Qwen2.5-VL System 2" - @echo " quantize quantize with SCHEME/STRATEGY (default fp8_default/s1)" - @echo " quantize-fp8 shorthand for SCHEME=fp8_default STRATEGY=s1" - @echo " quantize-nvfp4 shorthand for SCHEME=nvfp4_default STRATEGY=s1 (experimental)" - @echo " export-build ONNX export + LLM and visual engines" - @echo " export-build-base-fp16 unquantized FP16 reference engine" - @echo " verify-latents acceptance gate: z_latents cosine vs FP32" - @echo " benchmark latency and memory" - @echo "" - @echo " System 1 (requires INTERNNAV_PATH):" - @echo " export-system1 traj_dit + memory block -> BF16 engines" - @echo " verify-system1 trajectory parity vs PyTorch" - @echo " quantize-system1-fp8 System 1 FP8 PTQ (measured, not recommended)" + @echo " fetch-calib-source download InternData-N1's R2R train split (2.4 MB, gated)" + @echo " build-calib build the navigation-prompt calibration JSONL" + @echo " all both of the above, then print the CLI commands to run next" + @echo " clean remove CALIB_DATA_ROOT" @echo "" - @echo " Other:" - @echo " investigate-nvfp4 why NVFP4 breaks the System2 -> System1 bridge" - @echo " clean-onnx remove the ONNX intermediate (9-16 GB)" + @echo " Variables: INTERNVLA_CKPT CALIB_DATA_ROOT NUM_SAMPLES" @echo "" - @echo " Variables: INTERNVLA_CKPT INTERNNAV_PATH TRT_EDGELLM_DIR WORK_DIR" - @echo " ENGINE_DIR CALIB_DATA_ROOT SCHEME STRATEGY DEVICE" + @echo " Quantize / export / build / run are the standard TensorRT-Edge-LLM CLI," + @echo " not wrapped here -- see README.md." diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index 770a889..388f0c5 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -1,603 +1,98 @@ # internvla-n1-dualvln -FP8 quantization and TensorRT-Edge-LLM deployment for InternVLA-N1-DualVLN, a dual-system -vision-language navigation model, targeting NVIDIA Jetson Thor. - -> **Superseded by native support.** The repackage-and-host-side-bridge approach this recipe -> uses has been replaced by direct support in TensorRT-Edge-LLM itself: -> [NVIDIA/TensorRT-Edge-LLM#193](https://github.com/NVIDIA/TensorRT-Edge-LLM/pull/193) -> (pending review, tracked by -> [NVIDIA/TensorRT-Edge-LLM#190](https://github.com/NVIDIA/TensorRT-Edge-LLM/issues/190)). -> That work exports `InternVLA-N1-DualVLN` directly from the released checkpoint — no -> `repackage_system2.py` step, no `model_type` rewrite — and folds the `z_latents` bridge -> (`final_norm` + `cond_projector`) into the exported graph, so the engine emits `z_latents` -> on its own instead of a host-side Python step computing it. It also ships an async C++ -> runtime (`internvla_n1_dual_system_inference` / `internvla_n1_dual_system_server`) running -> both systems in one process, and 199-episode closed-loop navigation SR on Jetson Thor — -> validation this recipe never had. -> -> It also corrects this recipe's central claim below. **"The bridge is the acceptance -> metric, not text quality" is not true for this model.** Measured closed-loop, `z_latents` -> cosine and trajectory cosine each rank a different quantization scheme wrong, in opposite -> directions — including ranking this recipe's own FP8 recommendation *below* NVFP4, which -> this recipe's Notes section rules out on cosine grounds. The only metric that agreed with -> itself across schemes was closed-loop success rate. Kept below for the FP8-on-System-1 -> measurement and the `investigate_nvfp4.py` root-cause analysis, both still accurate; treat -> everything under "Results" and "Notes" that reasons from `z_latents` cosine as superseded. +FP8 / NVFP4 quantization and TensorRT-Edge-LLM deployment for InternVLA-N1-DualVLN, a +dual-system vision-language navigation model, on NVIDIA Jetson Thor. ## Status -- Model family: `InternVLA-N1-DualVLN` (System 2 = Qwen2.5-VL-7B, System 1 = NextDiT trajectory head) -- Quantization presets: `fp8_default`, `fp8_per_channel` (validated) · `nvfp4_*` (experimental, see Notes) -- Runtime target: TensorRT-Edge-LLM engines on Jetson Thor (sm_110) -- End to end on TensorRT: System 2 FP8 **and** System 1 BF16, both verified against PyTorch - (bridge 0.9919, System-1 trajectory 0.9997) — **2.55x** faster per planning step at - 9.16 GB of weights against 15.7 GB unquantized - -## What this model is, and what the metric has to be - -InternVLA-N1 is a **dual-system** navigation policy: - - camera image ─► ViT (BF16/FP8 TRT) ─► LLM (FP16/FP8 TRT) ─► branch - ├─ discrete action (↑ ← → STOP) - └─ pixel goal (coordinate) - System 2 (planner, ~2 Hz) ▼ - ──────────────────────────── latent_queries ─► LLM hidden ─► norm ─► cond_projector - │ z_latents - System 1 (controller, ~15 Hz) ▼ - ──────────────────────────── RGB ─► memory block (BF16 TRT) ─┐ - ├─► traj_dit (BF16 TRT, x10) - z_latents ───┘ - -Only **System 2** is quantized. System 1 stays BF16. - -The two are joined by `z_latents`: the last-layer hidden states of 4 trajectory tokens, run -through a host-side `final_norm` and `cond_projector`. **That bridge is the acceptance -metric, not text quality.** A quantized checkpoint can emit perfectly fluent captions and -still be useless for navigation — NVFP4 does exactly that here (see Notes). Every claim in -this recipe is gated on `z_latents` cosine against the FP32 reference, not on generated text. - -## Results - -### Benchmark matrix — both systems +Native support landed in TensorRT-Edge-LLM: +[NVIDIA/TensorRT-Edge-LLM#193](https://github.com/NVIDIA/TensorRT-Edge-LLM/pull/193) +(pending review, tracked by +[NVIDIA/TensorRT-Edge-LLM#190](https://github.com/NVIDIA/TensorRT-Edge-LLM/issues/190)). This +recipe used to repackage the checkpoint and compute the `z_latents` bridge on the host in +Python; both steps are gone now that TensorRT-Edge-LLM exports the checkpoint directly and +folds the bridge (`final_norm` + `cond_projector`) into the graph. What is left here is the one +thing that stays outside TensorRT-Edge-LLM: building a navigation-domain calibration set. -Everything below was measured on one Jetson Thor, batch 1, on an idle GPU. **System 2 is the -only part that gets quantized**; System 1 stays BF16 by design, so its row is a conversion -result rather than a quantization one. +**Measured on Jetson Thor, 199 R2R val_unseen episodes, closed-loop:** -#### System 2 — Qwen2.5-VL-7B planner (quantized) - -| Variant | weights / activations | KV cache | vision tower | checkpoint | LLM engine | visual engine | prefill (1024) | decode (pastKV 1024) | z_latents (engine) | pixel L2 mean / median | -|---|---|---|---|---|---|---|---|---|---|---| -| BF16 baseline | BF16 W16A16 -> FP16 engine | FP16 | BF16 | 16.6 GB | 14.15 GB | 1.36 GB | 135.8 ms | 56.4 ms | **0.999471** | 47.24 / 27.05 px | -| **FP8 s1** | **FP8 E4M3, W8A8, per-channel** | FP16 | BF16 | 10.1 GB | **7.62 GB** | 1.36 GB | **82.1 ms** | **31.5 ms** | **0.991861** | 46.26 / **22.51 px** | -| NVFP4 s1 (experimental) | NVFP4 E2M1, W4A4, block 16 w/ FP8 block scales | FP16 | BF16 | 7.2 GB | 4.77 GB | 1.36 GB | 73.2 ms | 20.2 ms | 0.931005 ✗ | 40.69 / 23.54 px | - -The KV cache is FP16 in all three: NVFP4 KV needs `sm100f` (datacenter Blackwell) and Thor is -sm110, and FP8 KV is a separate strategy (`s2`). The vision tower stays BF16 under `s1`; -quantizing it is `s3`/`s4` and is FP8-only, because the ViT MLP `intermediate_size` is 3420 -and 3420 / 16 = 213.75 does not divide by the NVFP4 block size. - -Weight quantization error, mean relative over 21 projections in layers 0/13/27: -FP8 **2.67 %**, NVFP4 **9.45 %**. - -#### System 1 — NextDiT diffusion head + memory block - -System 1 ships **BF16**. FP8 was measured rather than assumed, and the measurement is why it -does not ship: see below. Engine weights are BF16; the engine *interface* is fp32/int64, and -passing bf16 tensors trips an assertion in the wrapper rather than converting silently. - -| Component | weights | ONNX | engine | latency | cosine vs PyTorch | +| | prefill | decode | control rate | engine | SR vs PyTorch (69.8%) | |---|---|---|---|---|---| -| memory block (DAv2 + MemoryEncoder + QFormer) | BF16 | 200 MB | **104 MB** | **2.04 ms** | **0.999981** | -| `traj_dit`, one diffusion step | BF16 | 134 MB | **72 MB** | **5.85 ms** | **0.999508** | -| full trajectory (10 steps x 32 samples x 32 waypoints) | BF16 | — | — | **61.8 ms** | **0.999670** | -| **System 1 total** (memory + trajectory) | BF16 | 334 MB | **176 MB** | **63.8 ms · 15.7 Hz** | — | -| PyTorch `generate_traj` baseline | BF16 | — | — | 175.4 ms · 5.7 Hz | — | - -TensorRT gives System 1 a **2.75x** speedup at identical output. Latency is flat in -`num_sample_trajs` on the PyTorch side (175.4 / 175.6 / 173.1 ms at 32 / 4 / 1), so the head -is launch-bound rather than compute-bound — which is also why moving it to engines pays. - -#### FP8 on System 1 — measured, and not recommended - -Unlike System 2, System 1 goes `torch.onnx.export` -> `trtexec`, and TensorRT does FP8 only -through **explicit** quantization: the Q/DQ nodes must already be in the ONNX. So this needed -a ModelOpt PTQ pass (`quantize_system1.py`) calibrated on real tensors captured from a live -System 2 -> System 1 run (`dump_system1_calib.py`, 40 real `traj_dit` batches over 4 VLN -samples). Verified in the graph: 328 / 328 FP8 Q/DQ pairs in `traj_dit`, 160 / 160 in the -memory block. - -Waypoint deviation is the number to read. Cosine says how *aligned* two trajectories are; -deviation says how far apart the robot would actually end up, in the trajectory's own units, -against a mean per-waypoint reach of **0.2052**. - -| Config | traj_dit | memory | engines | trajectory | traj cosine | waypoint dev. mean / median / p95 | -|---|---|---|---|---|---|---| -| **BF16 (shipped)** | BF16 72 MB | BF16 104 MB | **176 MB** | 61.8 ms | **0.999670** | **0.0032 / 0.0012 / 0.0117** | -| mixed | **FP8 39 MB** | BF16 104 MB | 143 MB | 48.4 ms | 0.988032 ✗ | 0.0154 / 0.0050 / 0.0521 | -| all FP8 | **FP8 39 MB** | **FP8 75 MB** | **114 MB** | 49.6 ms | 0.980213 ✗ | 0.0198 / 0.0053 / 0.0793 | - -FP8 works — 1.55x smaller, 20 % faster on the diffusion loop, both configs well clear of -garbage — and it is still the wrong trade here: - -- **The size saving is irrelevant at the system level.** 62 MB off 9.16 GB of deployed - weights is 0.7 %. -- **The latency saving is nearly as small.** 12 ms off a 709 ms planning step is 1.7 %, - because System 2 dominates by an order of magnitude. -- **The fidelity cost is not small.** Mean waypoint deviation goes 0.0032 -> 0.0198, a **6x** - increase, and p95 goes 0.0117 -> 0.0793, which is 39 % of a typical waypoint's reach. - -Both FP8 configs fall below the 0.99 gate, and the split shows why: quantizing `traj_dit` -alone already costs 0.9997 -> 0.9880. Its per-step error is only 0.999508 -> 0.997811, but the -sampler runs 10 steps and each one feeds the next, so a small per-call error compounds. The -memory block adds the rest (0.999981 -> 0.990204) and buys **nothing** in latency (2.09 ms -against 2.04 ms) — its Conv2d stays unquantized anyway, since the legacy ONNX exporter cannot -infer a convolution kernel shape through Q/DQ. - -**Keep System 1 in BF16.** Quantization effort belongs on System 2, which is 98 % of both the -weights and the latency. The scripts stay in the recipe so the measurement is reproducible -and so the finding can be re-checked on a different System-1 configuration. - -#### Both systems, one planning step - -System 2 latency here is the full multi-image VLN step (~9 images, ~1764 image tokens) from -12 held-out samples, **not** the synthetic `llm_bench` figures above — different measurement, -so do not mix the two columns. - -| Configuration | System 2 quantization | System 2 | System 1 | total | vs PyTorch | -|---|---|---|---|---|---| -| all PyTorch | BF16 | 1631 ms | 175.4 ms | 1806 ms | 1.00x | -| TensorRT, unquantized | FP16 | 770 ms | 63.8 ms | 834 ms | 2.17x | -| **TensorRT, recommended** | **FP8 E4M3 W8A8** (System 1 stays BF16) | **646 ms** | **63.8 ms** | **710 ms** | **2.54x** | -| TensorRT, System 1 also FP8 | FP8 both systems | 646 ms | 51.6 ms | 698 ms | 2.59x ✗ | - -The last row is why System 1 stays BF16: 1.7 % off the step for 6x the waypoint deviation. - -Total on-device weights for the recommended configuration: **7.62 + 1.36 + 0.18 = 9.16 GB**, -against 15.7 GB for the unquantized TensorRT path. - -**FP8 is the recommended scheme.** Against BF16 it is 1.86x smaller and 1.65x/1.79x faster, -holds the bridge at 0.9919, and the median waypoint error does not get worse — it improves -slightly (27.05 -> 22.51 px), which is within the spread of a 42-sample set and should be read -as "unchanged", not as a gain from quantization. - -**NVFP4 is faster and smaller still but fails the gate.** Its bridge sits at 0.931, below the -0.99 threshold, so it is not recommended for navigation despite the attractive size and -latency. See the NVFP4 section for where that number comes from. - -### Where NVFP4's loss actually comes from - -Three measurements per scheme, each isolating one layer. The middle one — PyTorch with live -quantizers, so weights *and* activations are simulated exactly as the engine does them — is -what makes this decomposable: - -| | weights only | fake quant (W4A4) | engine | -|---|---|---|---| -| what is quantized | weights | weights + activations | weights + activations, real kernels | -| FP16 | — | — | 0.999471 | -| FP8 | 0.998020 | — | 0.991861 | -| **NVFP4** | **0.987986** | **0.978631** | **0.931005** | - -For NVFP4 that splits the 0.057 total loss as: - -- weight quantization: 0.012 -- activation quantization: **0.009** -- **everything else, inside the engine: 0.048** - -**This corrects an earlier claim in this file.** The gap was previously attributed to -activation quantization. It is not: activations account for 16 % of it, and 84 % appears -only once the model runs as a TensorRT engine. FP8 shows nothing comparable — its entire -PyTorch-to-engine gap is 0.006, while NVFP4 loses 0.048 there, nearly eight times more. - -The FP16 engine at 0.999471 bounds TensorRT's generic cost at about 0.0005, so this is not -export or runtime overhead in general — it is specific to the NVFP4 path. That is the same -neighbourhood as the known CASK epilogue miscompile, which `-cask_fusion:max_num_epilogues=1` -already improves from 0.647 to 0.931 but evidently does not fully resolve. - -#### Where the engine loses it — measured, not inferred - -`trt-edgellm/diagnose_engine_gap.py` compares the engine against the **fake-quant** model -rather than against the unquantized one, in a single process on the same inputs, so the -difference is the engine alone. On a text prompt, hidden states before the final norm: - -| engine | vs its own fake quant | -|---|---| -| FP8 | **0.998256** | -| NVFP4, `maxBatchSize 1` + CASK cap | **0.986790** | -| NVFP4, `maxBatchSize 2`, no CASK flag | **0.986790** | - -Two things follow. - -**The NVFP4 engine carries about eight times FP8's engine-side error** — 0.013 against -0.002 — and the bridge amplifies it: a 0.013 hidden-state deviation becomes the 0.048 seen -in z_latents once it passes through the final norm, GELU and `cond_projector`. - -**It is not the batch-1 miscompile.** The two NVFP4 engines are genuinely different builds -(different checksums, `maxBatchSize` 1 and 2, and the fork correctly withholds -`-cask_fusion:max_num_epilogues=1` from the batch-2 build) and they measure identically to -six decimal places. So `max_num_epilogues=1` fully recovers whatever the batch-1 path loses, -and the residual deficit is inherent to the NVFP4 kernels, independent of batch size. The -earlier guess that a second batch-1 miscompile was hiding here is wrong. - -A harness bug found this section's control worth having: comparing against -`output_hidden_states[-1]` reads 0.49 for a *known-good* FP8 engine, because the engine -emits hidden states before the final norm while that tensor is after it. The FP8 control -caught it; without one, 0.48 for NVFP4 would have looked like a broken kernel. - -**Practical consequence:** NVFP4 on this model is better than its engine number suggests. -At 0.9786 the quantization itself is close to the 0.99 gate. Anyone wanting to make NVFP4 -viable here should look at the TensorRT NVFP4 kernel path, not at better quantization -algorithms — which is also why AWQ, local-Hessian and QAT all failed to move it. - -### Two things to know before reading these numbers - -**Run everything on an idle GPU.** Latency measured while another job shared the device came -out 40-60 % higher (FP8 prefill 117 ms against 82 ms). The measurement script does not -enforce this. - -**The "identical replies" count is a weak indicator.** It moved from 31/42 to 11/42 for FP8 -across code revisions of the checkpoint loader, while `pixel_goal_l2` barely changed. Greedy -decoding over a 152k vocabulary flips on tiny logit differences, so treat the L2 median as the -number that matters and the identity count as colour. - -### Task accuracy — does the quantized model still pick the same waypoint? - -Measured with `quantize/benchmark_accuracy.py` on 42 held-out samples from two scenes, -identical samples and greedy decoding, PyTorch on Jetson Thor: - -| Checkpoint | pixel_goal_l2 mean | median | parse rate | -|---|---|---|---| -| System 2, BF16 (unquantized) | 47.24 px | **27.05 px** | 100 % | -| System 2, FP8 (s1) | 50.12 px | **26.97 px** | 100 % | - -31 of 42 replies are byte-identical and the **median deviation between the two is 0.00 px**. -The 2.88 px gap in the mean is carried by 8 samples, worst case 89 px — it is a small number -of disagreements, not a systematic shift, which is why both statistics are reported. For -reference, the source project's self-validation gate for this metric is a median under 60 px. - -Caveat worth keeping in view: this is a **weight-quantization** measurement. FP8 W8A8 also -quantizes activations, which the TensorRT engine does and the PyTorch path here does not, so -treat it as a lower bound on the engine's deviation rather than a prediction of it. - -The official InternVLA-N1 metrics (SR, SPL, NE, OS, nDTW) are all closed-loop and need -Habitat or InternUtopia plus MP3D scenes. Neither is installed here and neither is practical -on a Jetson, so the published table (DualVLN: NE 4.05 / SR 64.3 / SPL 58.5 on VLN-CE R2R) is -a literature reference point, not something this recipe reproduces. - -### Engine size and latency - -Both engines built from the same repackaged System 2 and measured with `llm_bench` on -Jetson Thor, batch 1: - -| | LLM engine | visual engine | prefill (1024 tok) | decode (pastKV 1024) | -|---|---|---|---|---| -| base FP16 (unquantized) | 15.0 GB | 1.36 GB | 196.22 ms | 86.16 ms | -| FP8 (s1) | **7.62 GB** | 1.36 GB | **95.80 ms** | **37.35 ms** | -| FP8 gain | 1.97x smaller | — | **2.05x faster** | **2.31x faster** | +| TensorRT FP8 | 90.6 ms | 32.8 ms | 61.3 ms (16.3 Hz) | 7.10 GB | 68.3% (p = 0.728) | +| TensorRT NVFP4 | 75.4 ms | 20.3 ms | 55.4 ms (18.0 Hz) | 4.45 GB | 67.8% (p = 0.572) | -Both produce the same text on the same prompt, so this is a straight win: FP8 halves the -engine and roughly doubles throughput while leaving the median waypoint error unchanged. +Neither differs from PyTorch significantly. **This replaces the recipe's earlier +recommendation.** The old version gated acceptance on `z_latents` cosine and, on that basis, +recommended FP8 and ruled NVFP4 out (cosine 0.647 against a 0.99 gate). Closed-loop SR shows +that gate does not predict the outcome it stands in for — full validation and the "no offline +metric predicts SR" finding are in the PR. Pick NVFP4 for speed and size, FP8 for a wider +margin; both are viable. -The FP16 engine is only correct because the build applies -`__LUNOWUD=-peep:fc_h_fusion=off`. The build log confirms it -(`Using __LUNOWUD=-peep:fc_h_fusion=off -peep:match_dual_gemm=off`). Without it, TensorRT -10.13 miscompiles Myelin's horizontal gate/up fusion on sm_110 and the engine emits fluent -gibberish — see the note further down. +## What this recipe does -### Bridge fidelity — z_latents - -`trt-edgellm/verify/verify_latents.py` reconstructs the System 2 -> System 1 bridge against -a PyTorch BF16 reference: embed, scatter image embeddings, append the 4 trajectory tokens, -run the engine with hand-built 3D mRoPE, take the last-layer hidden states, then apply the -host-side norm and `cond_projector`. - -| Engine | hidden pre-norm | hidden post-norm | **z_latents** | rel-L2 | -|---|---|---|---|---| -| base FP16 | 0.999843 | 0.999123 | **0.999471** | 0.0293 | -| FP8 (s1) | 0.997793 | 0.987343 | **0.991861** | 0.1023 | - -Both pass the > 0.99 gate. This is the check that NVFP4 fails (0.647), and it is the reason -it ships as experimental: text stays fluent while the waypoint bridge collapses. - -### System 1 engines (BF16, not quantized) - -Upstream InternNav ships **no** TensorRT or ONNX export at all — nothing in `internnav/` -references `trtexec`, `tensorrt` or `torch.onnx`. The System-1 conversion here is entirely -this recipe's, ported from the source project. - -Sizes, latency and fidelity are in the matrix above. The engine I/O, verified by execution: - -| Engine | inputs | output | -|---|---|---| -| `traj_dit` | `x[64,32,384]` f32, `timestep[64]` i64, `z_latents[64,*,768]` f32 | `output[64,32,384]` f32 | -| memory block | `images[T,3,224,224]` f32 | `memory_tokens[1,32,768]` f32 | - -The `64` is `2 x num_sample_trajs`: `generate_traj` runs classifier-free guidance, so the -conditioning is `[null, real]` and the latents are duplicated. - -**What stays in PyTorch on the host**, by design rather than omission: `action_encoder` and -`action_decoder` (two 3x384 linears), `pos_encoding`, `cond_projector` (the System 2 bridge), -and the `FlowMatchEulerDiscreteScheduler` loop itself. These are tiny or control-flow heavy; -the two engines cover the compute. - -**End-to-end parity: both engines reproduce PyTorch.** - -| | cosine vs PyTorch | -|---|---| -| `memory_tokens` (DAv2 + MemoryEncoder + QFormer engine) | **0.999981** | -| `traj_dit`, one diffusion step on the reference's own tensors | **0.999508** | -| full trajectory, 32 samples x 32 waypoints x 3 (rel-L2 0.0258) | **0.999670** | - -The check runs in **two stages across two environments**, because no single interpreter has -both halves: InternNav is written against transformers 4.x (under 5.x it fails first on -`config.hidden_size`, then on `apply_chunking_to_forward`, with no natural end), while the -TensorRT bindings ship for Python 3.12 only. - -Splitting it removes the conflict entirely. `verify/dump_system1_reference.py` runs the -PyTorch reference where InternNav works and writes its inputs *and* outputs to a `.pt`; -`verify/compare_system1_engines.py` reads that file where TensorRT works. The comparison -stays exact because the **same tensors** cross the boundary — the engines are fed the -reference's own inputs rather than regenerated ones. +Build a calibration set of realistic navigation prompts. Calibrating the quantized backbone on +its own prompt domain, rather than generic news text, measurably changes activation scales — +the same FP8 recipe moved from trajectory cosine 0.909 (`cnn_dailymail`, the CLI's default) to +0.978 (navigation prompts) in earlier testing on this model. ```bash -# stage A, transformers 4.51 environment (Python 3.10) -INTERNNAV_PATH=~/InternNav PYTHONPATH=~/InternNav \ -python verify/dump_system1_reference.py --output_path work/system1_reference.pt - -# stage B, TensorRT environment (Python 3.12) -python verify/compare_system1_engines.py \ - --reference_path work/system1_reference.pt --engine_dir work/onnx \ - --bench_iters 20 # optional: also time each engine and the whole trajectory +huggingface-cli download InternRobotics/InternData-N1 \ + vln_ce/raw_data/r2r/train/train.json.gz --repo-type dataset \ + --local-dir $CALIB_DATA_ROOT +# gated on Hugging Face -- accept the dataset terms and `huggingface-cli login` first + +python quantize/build_calib_jsonl.py \ + --train_json $CALIB_DATA_ROOT/vln_ce/raw_data/r2r/train/train.json.gz \ + --output $CALIB_DATA_ROOT/nav_calib.jsonl ``` -Two details are worth keeping, because both produce a confident wrong answer: - -- **Normalize before the memory block.** `generate_traj` divides by `_resnet_mean/_resnet_std` - before `rgb_model`; `MemBlock`, the module that was exported to ONNX, does not, so it - expects the already-normalized tensor. Feeding raw pixels made `memory_tokens` disagree - with what `generate_traj` used (0.315) while each half stayed internally consistent, which - reads exactly like a broken engine. -- **Capture the reference's starting noise, and probe one step.** `generate_traj` draws its - latents mid-function, so reseeding in stage B does not reproduce them, and a different draw - gives a different-but-valid trajectory (cosine ~0.31). Stage A therefore dumps the actual - noise, plus one real `traj_dit` call with its inputs and output. The single-step number is - what separates a bad engine from a bad reimplementation of the sampler loop — here it read - 0.9995 while the trajectory still read 0.31, which localized the fault to the harness. - -### A note on the older numbers - -An earlier revision reported z_latents of 0.99974 / 0.99985 / 0.99559 / 0.647 for -PyTorch / FP16 / FP8 / NVFP4 against an FP32 reference on 12 samples. The matrix above -supersedes those: it uses a BF16 reference, 42 samples, and the corrected checkpoint loader. -The ordering is the same and the conclusion is unchanged — FP8 passes, NVFP4 does not. - -**Calibration data made no measurable difference.** Held-out z_latents came out at 0.99143 -with generic `cnn_dailymail` text versus 0.99146 with a domain-specific VLN set — equal -within noise. An earlier apparent gain turned out to be overlap between the calibration and -probe sets. The VLN calibration loader ships anyway (it is the honest default to offer for a -navigation model), but do not expect it to buy accuracy. - -## Files - - . - ├── README.md - ├── Makefile # command panel for both paths - ├── requirements.txt # PyPI dependencies (excluding PyTorch) - ├── requirements-torch.txt # PyTorch installation guide - ├── configs/ - │ └── schemes.yaml # scheme x strategy validity matrix - ├── quantize/ # HF checkpoint -> quantized HF checkpoint - │ ├── repackage_system2.py # strip System 1 -> stock Qwen2.5-VL checkpoint - │ ├── quantize.py # ModelOpt driver - │ ├── quant_schemes.py # scheme registry + validity gate - │ ├── calibration.py # text / multimodal / VLN calibration loaders - │ ├── model_loader.py # load, calibrate, export - │ ├── load_quantized.py # read a ModelOpt checkpoint back WITH its scales - │ ├── benchmark_accuracy.py # pixel-goal L2 on held-out VLN episodes - │ ├── prompt_builder.py # VLN prompt — single source of truth - │ └── scripts/ # 00_fetch_calib_scenes, 01_repackage, 02_quantize - └── trt-edgellm/ # quantized checkpoint -> engines -> verification - ├── engine_runner.py # direct-TensorRT LLM harness (hand-built 3D mRoPE) - ├── export_traj_dit.py # System 1 diffusion head -> ONNX -> BF16 engine - ├── export_memory_block.py # System 1 memory block -> ONNX -> BF16 engine - ├── dump_system1_calib.py # real System 1 calibration tensors (needs InternNav) - ├── quantize_system1.py # System 1 FP8 PTQ -> ONNX -> engine (measured, not shipped) - ├── internvla_compat.py # the three patches needed to load System 1 - ├── traj_dit_loader.py memblock.py - ├── trt_torch.py # NVIDIA Apache-2.0 — header kept, not restamped - ├── engine_policy.py # simulator adapter (untested: needs Habitat) - ├── investigate_nvfp4.py # why NVFP4 breaks the System 1 bridge - ├── verify/ # 7 fidelity checks - ├── benchmark/ # 3 latency and memory benchmarks - ├── deploy/run_eval_engine.py - └── scripts/03_export_build_system2.sh - -## The repackage step, and why it matters - -`InternVLA-N1-DualVLN` declares `model_type: internvla_n1`, ships **no** modeling code in the -checkpoint, and therefore cannot be loaded with `trust_remote_code` — the class has to come -from the InternNav repository. - -`quantize/repackage_system2.py` sidesteps that for the entire quantization flow. It is pure -safetensors surgery: it streams the checkpoint, drops the eight System-1 tensor prefixes, -and rewrites `config.json` to `model_type: qwen2_5_vl` / -`architectures: ["Qwen2_5_VLForConditionalGeneration"]`. **It imports nothing from -InternNav.** - -After that step everything downstream is a stock Qwen2.5-VL flow: - -| Step | Needs `INTERNNAV_PATH`? | -|---|---| -| `00_fetch_calib_scenes.sh` | no | -| `01_repackage.sh` | **no** — pure file surgery | -| `02_quantize.sh` | no — operates on stock Qwen2.5-VL | -| `03_export_build_system2.sh` | no | -| `04_export_system1.sh` | **yes** | -| `05_verify.sh` (latents) | no | -| `05_verify.sh` (agent-level) | **yes** | - -Cost of this approach is one ~15 GB intermediate checkpoint. Pass `--free_source` to delete -each input shard as its converted copy is written if disk is tight. - -## Setup - -**Step 1 — PyTorch.** On Jetson, use the JetPack wheel; do not install from PyPI. - -| Platform | Command | -|---|---| -| Jetson Thor (JetPack 7.1, CUDA 13.0) | use the JetPack-provided `torch==2.10.0`; see requirements-torch.txt | -| x86 CUDA 12.8 | `pip install torch==2.10.0+cu128 --extra-index-url https://download.pytorch.org/whl/cu128` | +Everything after that is the standard TensorRT-Edge-LLM flow, in its own environment: -**Step 2 — remaining dependencies:** - - pip install -r requirements.txt - -**Step 2b — pin numpy afterwards.** This recipe needs numpy 1.x (OpenCV and diffusers break -under numpy 2 on Jetson), while the other recipes in this repository pin `numpy==2.2.6`. -Since `uv lock` resolves every extra into one universal lock, declaring both would make the -lock unsatisfiable, so numpy is deliberately absent from this recipe's extra. Run once after -`uv sync`: - - pip install "numpy==1.26.4" "scipy==1.13.1" - -**Step 2c — System 1 needs a different transformers.** The InternNav modeling code reads -`config.hidden_size` off the top-level config, which transformers 5.x no longer flattens, so -System-1 export fails there with `'InternVLAN1ModelConfig' object has no attribute -'hidden_size'`. Run the System-1 steps under **transformers 4.51.3** with `diffusers==0.33.1` -and `onnx==1.22.0` present. The System-2 path (repackage, quantize, export, engine build, -latent verification) is unaffected and runs on either. - -**Step 3 — dependencies that are not on PyPI.** These must be present before running anything: - -| Dependency | How | -|---|---| -| TensorRT 10.13 | ships with JetPack at `/usr/lib/python3.12/dist-packages` | -| `tensorrt-edgellm` 0.8.0 | build from source, then `pip install --no-deps -e $TRT_EDGE_LLM` | -| InternNav | `git clone` it; export `INTERNNAV_PATH`. Only needed for System 1 and agent-level checks | -| OpenCV | system package | - -## Environment - - export INTERNVLA_CKPT=/path/to/InternVLA-N1-DualVLN # source checkpoint - export INTERNNAV_PATH=/path/to/InternNav # System 1 only - export TRT_EDGE_LLM=/path/to/TensorRT-Edge-LLM # build root - export VLN_OPT_WORK=$HOME/vln-opt-work # intermediate artifacts - export VLN_OPT_ENGINES=$VLN_OPT_WORK/engines # engine output - -Two flags are mandatory and set by the scripts themselves — listed here so their absence is -diagnosable, not so you set them by hand: - -- `TRITON_BACKENDS_IN_TREE=1` for every quantization run. -- `__LUNOWUD="-peep:fc_h_fusion=off"` for every engine build on TensorRT 10.13/10.14. Without - it the FP16 engine emits gibberish — a Myelin miscompile on sm_110, not a precision problem. - -## Quick Start - - make repackage # InternVLA checkpoint -> stock Qwen2.5-VL System 2 - make quantize-fp8 # s1 FP8, text calibration - make export-build # ONNX export + FP8 LLM engine + visual engine - make verify-latents # the acceptance gate: z_latents cosine > 0.99 - -System 1 needs InternNav and its own two interpreters, since no single environment has both -InternNav (transformers 4.x) and the TensorRT bindings (Python 3.12): - - make export-system1 - PYTHON_PT=/path/to/py310/bin/python PYTHON_TRT=/path/to/py312/bin/python \ - make verify-system1 - -`make help` lists every target. Each script is also runnable directly; see the per-path -READMEs in `quantize/` and `trt-edgellm/`. - -## CLI Reference - - quantize/quantize.py [options] - - --model_path PATH Repackaged System 2 checkpoint (required) - --output_path PATH Destination for the quantized checkpoint (required) - --strategy {s1,s2,s3,s4} s1 LLM · s2 +KV cache · s3 +ViT · s4 +ViT+KV (default: s1) - --scheme NAME Preset from configs/schemes.yaml (default: fp8_default) - --calib {auto,text,multimodal,vln} Calibration source (default: auto) - --calib_data_root PATH Root for VLN calibration scenes - --num_calib_samples N Calibration samples (default: 512; image paths cap at 128) - --dtype {fp16,bf16} Load dtype (default: bf16) - --device DEVICE Torch device (default: cuda) - --resume DIR Layerwise checkpoint dir, for AWQ/Hessian crash recovery - --allow_experimental Permit schemes marked experimental (NVFP4) - --dry_run Load and validate, skip quantization - - quantize/repackage_system2.py [options] - - --model_path PATH Source InternVLA-N1-DualVLN checkpoint (required) - --output_path PATH Destination stock Qwen2.5-VL checkpoint (required) - --free_source Delete each source shard once converted (destructive) - -## Notes - -### What is and is not quantized - -Quantized: the System 2 LLM backbone (`model.layers.*`), and the vision tower under `s3`/`s4`. -Never quantized: System 1 (`traj_dit`, memory block), `lm_head`, and the host-side bridge ops -(`final_norm`, `cond_projector`). - -### Scheme validity - -| | s1 (LLM) | s2 (+KV) | s3 (+ViT) | s4 (+ViT+KV) | -|---|---|---|---|---| -| `fp8_default`, `fp8_per_channel` | yes | yes | yes | yes | -| `nvfp4_*` | experimental | experimental | **blocked** | **blocked** | - -- **NVFP4 with the vision tower is blocked, permanently.** The Qwen2.5-VL ViT MLP has - `intermediate_size = 3420`, and 3420 / 16 = 213.75 — not divisible by the NVFP4 block size. - The recipe rejects these combinations up front rather than letting them crash mid-run. -- **KV cache is always FP8**, even with NVFP4 weights. NVFP4 KV requires `sm100f` (datacenter - Blackwell); Thor is sm110. This is a hardware limit, not a configuration mistake — do not - "fix" it by selecting an NVFP4 KV preset. -- **NVFP4 weights are experimental and currently not usable for navigation.** They quantize, - export and generate fluent text, but `z_latents` cosine falls to 0.647, which breaks the - System 2 → System 1 bridge. Requires `--allow_experimental`. See - `trt-edgellm/investigate_nvfp4.py`. +```bash +# Point the CLI's default text_dataset at the local file instead of the Hub. +export EDGELLM_QUANT_DATASET_CNN_DAILYMAIL=$CALIB_DATA_ROOT/nav_calib.jsonl -### FP16 on Thor is fine — an earlier claim to the contrary was wrong +tensorrt-edgellm-quantize llm --model_dir $INTERNVLA_CKPT \ + --output_dir $QUANT_CKPT --quantization {fp8,nvfp4} +tensorrt-edgellm-export $QUANT_CKPT $ONNX_DIR -Earlier notes in the source project stated that a plain FP16 LLM engine "produces garbage on -Thor, so FP8 is mandatory". **That is incorrect and has been retracted.** The garbage came -from a Myelin `fc_h_fusion` miscompile on sm_110 at TensorRT 10.13: the horizontal fusion of -the gate/up projections is wrong at batch 1. TensorRT-Edge-LLM already disables it, but only -for TensorRT >= 10.15, so 10.13 slips through the gap. FP8 quietly dodged the same bug -because its Q/DQ nodes break the fusion pattern — which is why FP8 *looked* mandatory. +export EDGELLM_PLUGIN_PATH=.../libNvInfer_edgellm_plugin.so +export __LUNOWUD="-cask_fusion:max_num_epilogues=1" # NVFP4 at maxBatchSize 1 only +llm_build --onnxDir $ONNX_DIR/llm --engineDir $ENGINE_DIR/llm \ + --maxBatchSize 1 --maxInputLen 3072 --maxKVCacheCapacity 4096 +``` -Exporting `__LUNOWUD="-peep:fc_h_fusion=off"` fixes FP16 completely; the base FP16 engine is -the highest-fidelity variant in the results table. FP8 remains the recommended deployment -choice on size and latency, not on correctness. +System 1 (the trajectory expert) is not part of this checkpoint's quantization — it stays +BF16, built separately with `trtexec` — and the async runtime +(`internvla_n1_dual_system_inference` / `internvla_n1_dual_system_server`) that drives both +systems is not part of this repo either; it ships with TensorRT-Edge-LLM. See +[`experimental_models/internvla_n1/README.md`](https://github.com/NVIDIA/TensorRT-Edge-LLM/blob/main/experimental_models/internvla_n1/README.md) +in that repo (or the PR branch, until it merges) for the full export → build → run flow and the +resident-server protocol for driving it from a Python agent. -### One visual engine, sized for multi-image prompts +## Why calibration text needs no special tokens here -The visual encoder is built once at `--minImageTokens 4 --maxImageTokens 4096 ---maxImageTokensPerImage 1024`, and every consumer reads that one engine. +An earlier version of this pipeline appended four trajectory-query placeholder tokens to every +calibration prompt, because its own driver ran the System-2 → System-1 bridge forward pass +during calibration. `tensorrt-edgellm-quantize` does not: it loads System 2 as a stock +Qwen2.5-VL (see `internvla_n1_loader.py` in TensorRT-Edge-LLM) and calibrates with ordinary +text forward passes, never touching the bridge. Appending tokens the tokenizer does not even +have registered yet — they are added later, at export time — would only add noise. -This is worth stating because the source project got it wrong in a way that is easy to -repeat: its default build used `128 / 512 / 512`, taken from a single-image demo. A VLN -prompt carries 9–10 images and roughly 1,764 image tokens, so it does not fit in 512, and the -multi-image verification scripts were quietly pointed at a hand-built engine that no script -in the repository produced. If a verification here reports a shape or capacity error, check -the visual engine's sizing before suspecting the weights. +## What is and is not quantized -### Scope +Quantized: the System 2 LLM backbone. Never quantized: System 1 (`traj_dit`, memory block), +and the bridge (`cond_projector`, `latent_queries`) — four rows through a +Linear/GELU/Linear, kept at source precision because quantizing them saves nothing measurable +and puts error directly on the tensor System 1 steers by. -These are **conversion-fidelity** numbers plus offline planner metrics. They are not a -closed-loop navigation success rate — that needs the Habitat/InternUtopia simulator, which is -not part of this recipe. `trt-edgellm/verify/verify_engine_policy.py` checks that the -simulator adapter is wired correctly, but cannot exercise it. +**NVFP4 with the vision tower is blocked, permanently.** The Qwen2.5-VL ViT MLP has +`intermediate_size = 3420`, and 3420 / 16 = 213.75 — not divisible by the NVFP4 block size. +Only the LLM backbone is quantized here, so this does not apply, but it is worth knowing if you +extend this to a strategy that includes the vision tower. -## Tested Environments +## Tested environment -- **OS:** Ubuntu 24.04 (JetPack 7.1) -- **Hardware:** NVIDIA Jetson Thor (Blackwell, sm_110, 128 GB unified memory) -- **Python:** 3.12.3 -- **PyTorch:** 2.10.0 -- **CUDA:** 13.0 -- **TensorRT:** 10.13.3.9 -- **tensorrt-edgellm:** 0.8.0 -- **nvidia-modelopt:** 0.44.0 -- **transformers:** 4.51.3 · **diffusers:** 0.33.1 · **onnx:** 1.22.0 · **numpy:** 1.26.4 +Jetson Thor, JetPack 7.1 (TensorRT 10.13.3.9, CUDA 13). `pip install -e ".[tools]"` in the +TensorRT-Edge-LLM checkout pulls `nvidia-modelopt` and `datasets`; without it +`tensorrt-edgellm-quantize` fails at import. diff --git a/recipes/internvla-n1-dualvln/configs/schemes.yaml b/recipes/internvla-n1-dualvln/configs/schemes.yaml deleted file mode 100644 index 42fb6c6..0000000 --- a/recipes/internvla-n1-dualvln/configs/schemes.yaml +++ /dev/null @@ -1,109 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -# -# Quantization schemes and the scheme x strategy validity matrix for -# InternVLA-N1-DualVLN on Jetson Thor (sm_110). -# -# Two of the constraints below are hardware facts, not tuning choices. Encoding them -# here means an impossible combination is rejected before the model is loaded, instead -# of crashing forty minutes into a run: -# -# * NVFP4 cannot touch the Qwen2.5-VL vision tower. Its MLP intermediate_size is -# 3420, and 3420 / 16 = 213.75 -- not divisible by the NVFP4 block size. -# * NVFP4 KV cache requires sm100f (datacenter Blackwell). Thor is sm110, so the KV -# cache is FP8 even when the weights are NVFP4. This is deliberate; do not -# "fix" it by pointing at an NVFP4 KV preset. - -# --------------------------------------------------------------------------- # -# Strategies -- which parts of System 2 get quantized -# --------------------------------------------------------------------------- # -strategies: - s1: - description: LLM backbone only -- fastest, safest, the validated default - quantize_kv_cache: false - quantize_visual: false - s2: - description: LLM + KV cache -- lower memory for long multi-image prompts - quantize_kv_cache: true - quantize_visual: false - s3: - description: LLM + vision tower -- full VLM quantization - quantize_kv_cache: false - quantize_visual: true - s4: - description: LLM + vision tower + KV cache -- maximum compression - quantize_kv_cache: true - quantize_visual: true - -# --------------------------------------------------------------------------- # -# Schemes -- weight/activation formats, mapped to ModelOpt presets -# -# calib_batch_size applies to text calibration only; image calibration is always 1. -# --------------------------------------------------------------------------- # -schemes: - fp8_default: - modelopt_preset: FP8_DEFAULT_CFG - description: FP8 W8A8 per-tensor -- balanced accuracy and speed - calib_batch_size: 8 - status: validated - fp8_per_channel: - modelopt_preset: FP8_PER_CHANNEL_PER_TOKEN_CFG - description: FP8 W8A8 per-channel weight, per-token activation -- highest FP8 accuracy - calib_batch_size: 8 - status: validated - nvfp4_default: - modelopt_preset: NVFP4_DEFAULT_CFG - description: NVFP4 W4A4 -- smallest footprint - calib_batch_size: 4 - status: experimental - nvfp4_awq_full: - modelopt_preset: NVFP4_AWQ_FULL_CFG - description: NVFP4 W4A4 + AWQ (lite + clip) -- best NVFP4 accuracy - calib_batch_size: 1 - status: experimental - nvfp4_local_hessian: - modelopt_preset: NVFP4_W4A4_WEIGHT_LOCAL_HESSIAN_CFG - description: NVFP4 W4A4 + local Hessian -- highest-accuracy NVFP4, slow - calib_batch_size: 1 - status: experimental - -# --------------------------------------------------------------------------- # -# KV cache pairing -# --------------------------------------------------------------------------- # -kv_cache: - preset: FP8_KV_CFG - applies_to: [fp8_default, fp8_per_channel, nvfp4_default, nvfp4_awq_full, nvfp4_local_hessian] - note: >- - FP8 KV for every weight format. NVFP4 KV needs sm100f at inference time, which Thor - (sm110) does not provide; FP8 KV pairs portably with NVFP4 weights. - -# --------------------------------------------------------------------------- # -# Modules excluded when a strategy leaves the vision tower alone. -# -# ModelOpt presets exclude *lm_head* by default but NOT the vision tower, so without -# these globs s1/s2 would silently quantize the ViT. -# --------------------------------------------------------------------------- # -visual_exclude_patterns: - - "*visual*" - - "*vision_tower*" - - "*multi_modal_projector*" - -# --------------------------------------------------------------------------- # -# Validity matrix -- every combination not listed as blocked or experimental is allowed -# --------------------------------------------------------------------------- # -blocked: - - schemes: [nvfp4_default, nvfp4_awq_full, nvfp4_local_hessian] - strategies: [s3, s4] - reason: >- - NVFP4 cannot quantize the Qwen2.5-VL vision tower: its MLP intermediate_size is - 3420 and 3420 / 16 = 213.75, which is not an integer, so the tensor cannot be - split into NVFP4 blocks. Use an fp8_* scheme for s3/s4. - -experimental: - - schemes: [nvfp4_default, nvfp4_awq_full, nvfp4_local_hessian] - strategies: [s1, s2] - reason: >- - NVFP4 weights quantize, export and generate fluent text, but the System 2 -> System 1 - bridge breaks: z_latents cosine against the FP32 reference falls to 0.647 versus - 0.9956 for FP8. The model will caption an image correctly and still navigate wrongly. - Pass --allow_experimental to proceed. See trt-edgellm/investigate_nvfp4.py. diff --git a/recipes/internvla-n1-dualvln/quantize/benchmark_accuracy.py b/recipes/internvla-n1-dualvln/quantize/benchmark_accuracy.py deleted file mode 100644 index b3bb700..0000000 --- a/recipes/internvla-n1-dualvln/quantize/benchmark_accuracy.py +++ /dev/null @@ -1,259 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Score a System-2 checkpoint on held-out VLN episodes, without a simulator. - -The official InternVLA-N1 metrics (SR, SPL, NE, OS, nDTW) are all closed-loop and need -Habitat or InternUtopia plus MP3D scenes. Neither is practical on a Jetson, so this -measures the thing that quantization actually threatens and that *can* be measured -offline: whether System 2 still emits the same navigation decision. - -Two metrics, both against the LeRobot ground truth: - -* ``pixel_goal_l2`` -- Euclidean distance between the predicted and annotated waypoint, - in pixels of the 384x384 prompt image. This is the primary number: the waypoint is what - System 1 consumes, so an error here is an error in where the robot goes. -* ``action_accuracy`` -- agreement on the discrete action token (STOP / up / left / right). - -Both are computed per checkpoint on identical samples, so two runs are directly -comparable. The interesting quantity is the *delta* between an unquantized and a -quantized checkpoint, not either absolute value. - -Prompts are built through ``prompt_builder``, the same path the deployed agent uses, so -what is scored is the deployed prompt rather than a hand-rebuilt approximation. -""" -import argparse -import glob -import json -import os -import re -import sys -import time - -import numpy as np -import torch - -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -import prompt_builder as pb # noqa: E402 -from load_quantized import load_for_eval # noqa: E402 - -# From InternNav's dataset: idx2actions = {0: STOP, 1: up, 2: left, 3: right, 5: down} -IDX2ACTION = {0: "STOP", 1: "↑", 2: "←", 3: "→", 5: "↓"} -# The agent runs two turns. Turn 1 shows the level camera plus history and the model -# replies with a discrete action -- or with the look-down token when the next thing it -# needs is a waypoint. The robot then tilts, and turn 2 appends the single pitched frame, -# where the coordinate is produced. That is why every goal.125cm_0deg entry in the data is -# the [-1,-1] sentinel: a level camera cannot see a point on the floor. -# -# Scoring has to follow the same two turns. Feeding a pitched frame through a turn-1 -# prompt puts the model off-distribution and it answers with an action instead of a -# coordinate, which reads as a 0% parse rate and looks like a model failure when it is a -# harness failure. -LEVEL_CAMERA = "125cm_0deg" -DEFAULT_CAMERA = "125cm_30deg" -LOOKDOWN_TOKEN = "\u2193" - - -def parse_xy(text: str): - """Pull the first two integers out of a model reply, matching InternNav's parser.""" - nums = re.findall(r"-?\d+", text) - if len(nums) < 2: - return None - return float(nums[0]), float(nums[1]) - - -def parse_action(text: str): - for token in ("STOP", "↑", "←", "→", "↓"): - if token in text: - return token - return None - - -def discover_samples(data_root: str, max_samples: int, seed: int = 0, - camera: str = DEFAULT_CAMERA) -> list[dict]: - """Collect (images, instruction, GT waypoint, GT action) tuples from LeRobot episodes.""" - import pyarrow.parquet as pq - - rng = np.random.default_rng(seed) - rgb_key = f"observation.images.rgb.{camera}" - goal_col = f"goal.{camera}" - samples: list[dict] = [] - for meta in sorted(glob.glob(os.path.join(data_root, "**", "meta", "episodes.jsonl"), - recursive=True)): - scene = os.path.dirname(os.path.dirname(meta)) - episodes = [json.loads(line) for line in open(meta)] - for ep in episodes: - idx = ep["episode_index"] - length = ep["length"] - if length < 2: - continue - table = None - for parquet in glob.glob(os.path.join(scene, "data", "**", - f"episode_{idx:06d}.parquet"), - recursive=True): - table = pq.read_table(parquet) - break - if table is None or goal_col not in table.schema.names: - continue - - goals = table.column(goal_col).to_pylist() - actions = table.column("action").to_pylist() - # goal is [-1, -1] on frames that carry no waypoint (a pure turn or stop). - # Scoring L2 against that sentinel would be meaningless, so sample only from - # frames that actually annotate one. - usable = [i for i in range(1, min(length, len(goals))) - if goals[i] is not None and int(goals[i][0]) >= 0] - if not usable: - continue - t = int(usable[rng.integers(0, len(usable))]) - - history = np.unique(np.linspace(0, t - 1, pb.NUM_HISTORY, dtype=np.int32)).tolist() - level_key = f"observation.images.rgb.{LEVEL_CAMERA}" - frames = [os.path.join(scene, "videos", "chunk-000", level_key, - f"episode_{idx:06d}_{i}.jpg") for i in history + [t]] - lookdown = os.path.join(scene, "videos", "chunk-000", rgb_key, - f"episode_{idx:06d}_{t}.jpg") - if not all(os.path.isfile(p) for p in frames + [lookdown]): - continue - - action = actions[t] - samples.append({ - "episode": os.path.basename(scene), - # prompt_builder's "episode_idx" is the step index within the episode -- - # it derives the history frames from linspace(0, episode_idx-1) -- not the - # episode number. Passing the episode number yields one placeholder - # and a count mismatch against the history frames. - "episode_idx": t, - "episode_number": idx, - "instruction": (ep.get("tasks") or [""])[0], - "images": frames + [lookdown], - "turn": 2, - "assistant_turn1": LOOKDOWN_TOKEN, - "gt_goal": goals[t], - "gt_action": IDX2ACTION.get(int(action) if action is not None else -1), - }) - if len(samples) >= max_samples: - return samples - return samples - - -def evaluate(model_path: str, samples: list[dict], max_new_tokens: int, - device: str) -> dict: - model, processor, algo = load_for_eval(model_path, device=device) - # The prompt images are resized to 384x384, so predicted and GT pixel coordinates - # live in the same frame and the L2 is directly interpretable. - l2, n_parsed, n_action_ok, n_action_total = [], 0, 0, 0 - replies = [] - t0 = time.time() - - for i, sample in enumerate(samples, 1): - inputs = pb.build_sample_inputs(sample, processor).to(device) - with torch.inference_mode(): - out = model.generate(**inputs, max_new_tokens=max_new_tokens, - do_sample=False, temperature=None, top_p=None, top_k=None) - reply = processor.batch_decode(out[:, inputs["input_ids"].shape[1]:], - skip_special_tokens=True)[0].strip() - replies.append(reply) - - xy = parse_xy(reply) - if xy is not None and sample["gt_goal"] is not None: - n_parsed += 1 - gt = sample["gt_goal"] - l2.append(float(np.hypot(xy[0] - float(gt[0]), xy[1] - float(gt[1])))) - - # Only meaningful on turn 1. At turn 2 the model emits a coordinate by design, - # so scoring it against a movement action would always read 0% and look like a - # regression rather than the protocol working. - if sample.get("turn", 1) == 1 and sample["gt_action"] is not None: - n_action_total += 1 - if parse_action(reply) == sample["gt_action"]: - n_action_ok += 1 - - if i % 5 == 0 or i == len(samples): - print(f" [{i}/{len(samples)}] {time.time() - t0:.0f}s elapsed", flush=True) - - del model - torch.cuda.empty_cache() - - return { - "model_path": model_path, - "quant_algo": algo, - "n_samples": len(samples), - "pixel_goal_l2_mean": float(np.mean(l2)) if l2 else None, - "pixel_goal_l2_median": float(np.median(l2)) if l2 else None, - "pixel_goal_parse_rate": n_parsed / max(len(samples), 1), - "action_accuracy": n_action_ok / n_action_total if n_action_total else None, - "n_action_scored": n_action_total, - "seconds": round(time.time() - t0, 1), - "replies": replies, - } - - -def parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--model_path", required=True, action="append", - help="Checkpoint to score. Repeat to compare several on identical samples.") - p.add_argument("--data_root", required=True, help="Held-out LeRobot episodes") - p.add_argument("--num_samples", type=int, default=32) - p.add_argument("--max_new_tokens", type=int, default=32) - p.add_argument("--device", default="cuda") - p.add_argument("--seed", type=int, default=0) - p.add_argument("--camera", default=DEFAULT_CAMERA, - help="LeRobot camera setting to score against. Must be a pitched\n" - "one; the level 125cm_0deg carries no pixel goal.") - p.add_argument("--output_path", default=None, help="Write results JSON here") - return p.parse_args() - - -def main() -> int: - args = parse_args() - - samples = discover_samples(args.data_root, args.num_samples, seed=args.seed, - camera=args.camera) - if not samples: - print(f"[ERROR] no usable samples under {args.data_root}") - return 1 - print(f"Scoring {len(samples)} held-out samples from " - f"{len({s['episode'] for s in samples})} scene(s)\n") - - results = [] - for path in args.model_path: - print(f"=== {path}") - res = evaluate(path, samples, args.max_new_tokens, args.device) - results.append(res) - print(f" quant : {res['quant_algo'] or 'none (bf16)'}") - print(f" pixel_goal_l2 : mean {res['pixel_goal_l2_mean']:.2f} px, " - f"median {res['pixel_goal_l2_median']:.2f} px" - if res["pixel_goal_l2_mean"] is not None else " pixel_goal_l2 : n/a") - print(f" parse_rate : {100 * res['pixel_goal_parse_rate']:.1f}%") - if res["action_accuracy"] is not None and res["n_action_scored"]: - print(f" action_accuracy : {100 * res['action_accuracy']:.1f}% " - f"({res['n_action_scored']} scored)") - print(f" took : {res['seconds']}s\n") - - if len(results) > 1: - base = results[0] - print("=== delta vs " + os.path.basename(base["model_path"])) - for res in results[1:]: - name = os.path.basename(res["model_path"]) - if base["pixel_goal_l2_mean"] and res["pixel_goal_l2_mean"]: - d = res["pixel_goal_l2_mean"] - base["pixel_goal_l2_mean"] - print(f" {name}: pixel_goal_l2 {d:+.2f} px") - if (base["action_accuracy"] is not None and res["action_accuracy"] is not None - and base["n_action_scored"]): - d = 100 * (res["action_accuracy"] - base["action_accuracy"]) - print(f" {name}: action_accuracy {d:+.1f} pp") - agree = sum(a == b for a, b in zip(base["replies"], res["replies"])) - print(f" {name}: identical replies {agree}/{len(res['replies'])}") - - if args.output_path: - os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True) - with open(args.output_path, "w") as f: - json.dump(results, f, indent=2) - print(f"\nWrote {args.output_path}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/quantize/build_calib_jsonl.py b/recipes/internvla-n1-dualvln/quantize/build_calib_jsonl.py new file mode 100644 index 0000000..23401bc --- /dev/null +++ b/recipes/internvla-n1-dualvln/quantize/build_calib_jsonl.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics +# SPDX-License-Identifier: BSD-3-Clause +"""Build a navigation-prompt calibration JSONL for `tensorrt-edgellm-quantize`. + +The native `tensorrt-edgellm-quantize llm --text_dataset NAME` flag only accepts a +*registered* dataset name (the default is `cnn_dailymail`), but any built-in can be pointed +at a local file with `EDGELLM_QUANT_DATASET_=/path/to/file.jsonl` -- see +`tensorrt_edgellm/quantization/datasets/__init__.py::local_override_path`. This script writes +that file in the schema the override loader expects: one `{"article": "..."}` object per line +(the field name `cnn_dailymail()` reads, since that is the dataset being overridden). + +Why bother, instead of the CLI's cnn_dailymail default: calibrating on the deployment prompt's +own domain measurably changes activation scales. The same FP8 recipe went from trajectory +cosine 0.909 (cnn_dailymail) to 0.978 (this) in earlier testing -- calibration text that never +resembles a navigation instruction leaves the backbone's activation ranges fitted to news +articles instead of the prompts it actually deploys on. + +One thing this does *not* need, unlike an earlier version of this recipe: the four trailing +trajectory-query tokens (`<|latent_q0..3|>`). Those only matter for calibrating the +System-2 -> System-1 bridge forward pass, and `tensorrt-edgellm-quantize` never runs that -- +it loads System 2 as a stock Qwen2.5-VL (see `internvla_n1_loader.py` in TensorRT-Edge-LLM) +and calibrates with ordinary text forward passes. Appending tokens the tokenizer does not yet +know about (they are registered later, at export time) would only add noise. + +Usage: + huggingface-cli download InternRobotics/InternData-N1 \\ + vln_ce/raw_data/r2r/train/train.json.gz --repo-type dataset \\ + --local-dir $CALIB_DATA_ROOT + python build_calib_jsonl.py \\ + --train_json $CALIB_DATA_ROOT/vln_ce/raw_data/r2r/train/train.json.gz \\ + --output $CALIB_DATA_ROOT/nav_calib.jsonl +""" +import argparse +import gzip +import json +import random + +PROMPT_TEMPLATE = ( + "You are an autonomous navigation assistant. Your task is to {instruction} " + "Where should you go next to stay on track? Please output the next waypoint's " + "coordinates in the image. Please output STOP when you have successfully completed " + "the task.") + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--train_json", required=True, + help="InternData-N1 vln_ce/raw_data/r2r/train/train.json.gz") + ap.add_argument("--output", required=True, help="Destination JSONL") + ap.add_argument("--num_samples", type=int, default=512, + help="Matches the native CLI's calibration sample count (default: 512)") + ap.add_argument("--seed", type=int, default=0) + args = ap.parse_args() + + with gzip.open(args.train_json) as f: + episodes = json.load(f)["episodes"] + + random.Random(args.seed).shuffle(episodes) + + written = 0 + with open(args.output, "w") as out: + for ep in episodes: + if written >= args.num_samples: + break + text = ep["instruction"]["instruction_text"].strip() + if not text: + continue + instruction = text.rstrip(". ") + prompt = PROMPT_TEMPLATE.format(instruction=instruction + ".") + out.write(json.dumps({"article": prompt}) + "\n") + written += 1 + + print(f"wrote {written} navigation calibration prompts to {args.output}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/quantize/calibration.py b/recipes/internvla-n1-dualvln/quantize/calibration.py deleted file mode 100644 index a992c15..0000000 --- a/recipes/internvla-n1-dualvln/quantize/calibration.py +++ /dev/null @@ -1,421 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Calibration dataloaders for Qwen2.5-VL quantization. - -Three dataloader paths: - * ``text_calib_dataloader`` — ``cnn_dailymail`` ``input_ids`` only, used when - quantizing LLM backbone alone (no visual quantization). - * ``multimodal_calib_dataloader`` — ``lmms-lab/MMMU`` image+text pairs streamed - through the model's own ``AutoProcessor`` chat template, used when visual - tower quantization is enabled. - * ``vln_calib_dataloader`` — the in-distribution set: real InternData-N1 VLN-CE - episodes (the data System 2 was fine-tuned on) assembled through the deployed - agent's own prompt path, so ModelOpt sees the true activation ranges of the - farthest-pixel-goal navigation task instead of out-of-distribution web text/QA. - -The first two mirror NVIDIA's reference implementation. -""" - -import os -from typing import Any, Optional - -import torch -from torch.utils.data import DataLoader - - -# --------------------------------------------------------------------------- # -# Offline fallbacks -# --------------------------------------------------------------------------- # -def _hub_reachable(timeout: float = 5.0) -> bool: - """Cheap reachability probe for huggingface.co. - - Returns False on a TLS interception proxy as well as on a real outage, which is - what we want: in both cases the cached parquet is the usable path. - """ - import urllib.error - import urllib.request - - try: - urllib.request.urlopen("https://huggingface.co/api/whoami-v2", timeout=timeout) - return True - except urllib.error.HTTPError: - return True # reachable, just unauthenticated - except Exception: - return False - - -def _local_parquet_split(dataset_name: str, split: str) -> list[str]: - """Find cached parquet shards for ``dataset_name``/``split`` in the HF hub cache.""" - cache = os.environ.get("HF_HUB_CACHE") or os.path.expanduser( - "~/.cache/huggingface/hub") - repo_dir = os.path.join(cache, "datasets--" + dataset_name.replace("/", "--")) - snapshots = os.path.join(repo_dir, "snapshots") - if not os.path.isdir(snapshots): - return [] - found: list[str] = [] - for root, _dirs, files in os.walk(snapshots): - for name in sorted(files): - if name.startswith(split) and name.endswith(".parquet"): - found.append(os.path.join(root, name)) - return sorted(found) - - -# --------------------------------------------------------------------------- # -# Text calibration (LLM-only strategies S1, S2) -# --------------------------------------------------------------------------- # -def text_calib_dataloader( - tokenizer, - dataset_name: str = "abisee/cnn_dailymail", - batch_size: int = 1, - num_samples: int = 512, - max_length: int = 512, -) -> DataLoader: - """Return a DataLoader of tokenised ``input_ids`` for calibration. - - Mirrors NVIDIA's ``_text_calib_dataloader``. ``max_length`` only governs - tokenizer truncation during this call and does NOT modify - ``tokenizer.model_max_length`` — verified by an assert in the caller. - """ - from datasets import load_dataset - - # A Jetson is often behind a TLS-intercepting proxy or fully air-gapped, where - # load_dataset() fails even though the parquet shards are already in the hub cache - # (HF's own offline mode does not help: it still wants the dataset script). Fall - # back to reading those shards directly so a calibration run does not depend on - # network reachability. Set HF_DATASETS_LOCAL_ONLY=1 to skip the hub attempt. - local = _local_parquet_split(dataset_name, split="train") - if local and (os.environ.get("HF_DATASETS_LOCAL_ONLY") == "1" or not _hub_reachable()): - print(f" [calib] reading cached parquet ({len(local)} shard(s)) instead of the Hub") - ds = load_dataset("parquet", data_files=local, split="train") - col = "article" if "article" in ds.column_names else ds.column_names[0] - texts = ds[col][:num_samples] - elif "abisee/cnn_dailymail" in dataset_name: - ds = load_dataset(dataset_name, name="3.0.0", split="train") - texts = ds["article"][:num_samples] - else: - ds = load_dataset(dataset_name, split="train") - if "text" in ds.column_names: - col = "text" - elif "article" in ds.column_names: - col = "article" - else: - raise ValueError( - f"Dataset {dataset_name!r} has no 'text' or 'article' column: " - f"{ds.column_names}" - ) - texts = ds[col][:num_samples] - - enc = tokenizer( - texts, - return_tensors="pt", - padding=True, - truncation=True, - max_length=max_length, - ) - return DataLoader(enc["input_ids"], batch_size=batch_size, shuffle=False) - - -# --------------------------------------------------------------------------- # -# Multimodal calibration (visual-quantization strategies S3, S4) -# --------------------------------------------------------------------------- # -def _iter_image_question_pairs(dataset_name: str): - """Yield ``(image, question)`` pairs from a HuggingFace calibration dataset. - - Mirrors NVIDIA's ``_iter_image_question_pairs``. Tolerant of two common - schemas: - * ScienceQA-style: single ``image`` column. - * MMMU-style: numbered ``image_1`` / ``image_2`` / ... columns. - - Splits are tried in the order ``dev`` → ``validation`` → ``train``. - """ - from datasets import load_dataset - - last_err: Optional[Exception] = None - ds = None - for split in ("dev", "validation", "train"): - try: - ds = load_dataset(dataset_name, split=split, streaming=True) - break - except Exception as e: # noqa: BLE001 - last_err = e - if ds is None: - raise RuntimeError( - f"Could not load {dataset_name!r} via any of " - f"split=dev/validation/train" - ) from last_err - - for example in ds: - image = example.get("image") - if image is None: - for i in range(1, 8): - image = example.get(f"image_{i}") - if image is not None: - break - question = example.get("question") or "" - if image is not None and question: - yield image, question - - -def multimodal_calib_dataloader( - processor, - dataset_name: str = "lmms-lab/MMMU", - num_samples: int = 128, - max_length: int = 512, -) -> list[dict[str, Any]]: - """Materialise a list of ``BatchFeature`` dicts with ``input_ids`` + ``pixel_values``. - - Mirrors NVIDIA's ``_multimodal_calib_dataloader``. Streams image-question - pairs through the model's own ``AutoProcessor`` chat template so the visual - tower receives real activations. - - NVIDIA caps multimodal calibration at 128 samples because VLM calibration - is GPU-memory bound. The caller is expected to enforce this cap. - - Returns a *list* (not generator) so ModelOpt can re-iterate forward_loop - during algorithm selection (mirrors NVIDIA's design). - """ - batches: list[dict[str, Any]] = [] - for image, question in _iter_image_question_pairs(dataset_name): - messages = [{ - "role": "user", - "content": [ - {"type": "image", "image": image}, - {"type": "text", "text": question}, - ], - }] - - inputs = processor.apply_chat_template( - messages, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - return_tensors="pt", - ) - - batches.append({ - k: v - for k, v in inputs.items() - if v is not None - and not (isinstance(v, torch.Tensor) and v.numel() == 0) - }) - if len(batches) >= num_samples: - break - - if not batches: - raise RuntimeError( - f"No usable multimodal samples from {dataset_name!r}. " - "Check dataset access / processor chat template." - ) - return batches - - -# --------------------------------------------------------------------------- # -# VLN calibration (in-distribution — InternData-N1 VLN-CE) -# --------------------------------------------------------------------------- # -def _import_prompt_builder(): - """Import ``lib/prompt_builder`` (the single source of truth for the System 2 - prompt) regardless of how this module was launched. Kept lazy so the text/MMMU - paths never pay for it.""" - import os - import sys - - lib = os.path.join( - os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), - "lib", - ) - if lib not in sys.path: - sys.path.insert(0, lib) - import prompt_builder # noqa: E402 - - return prompt_builder - - -def _discover_vln_episodes(data_root: str, rgb_key: str): - """Walk a LeRobot dataset tree and yield one record per episode. - - ``data_root`` may be a single scene dir or any parent of several — every dir - holding ``meta/episodes.jsonl`` is picked up. Frames are stored as per-frame - JPGs under ``videos/chunk-XXX//episode_{idx:06d}_{frame}.jpg`` (no - video decoding needed). Falls back to the first available ``*.rgb.*`` stream - when ``rgb_key`` is absent in a scene. - """ - import glob - import json - import os - - records = [] - meta_files = glob.glob( - os.path.join(data_root, "**", "meta", "episodes.jsonl"), recursive=True - ) - if not os.path.isdir(os.path.join(data_root, "meta")) and not meta_files: - raise FileNotFoundError( - f"No LeRobot episodes.jsonl found under {data_root!r} " - "(expected /meta/episodes.jsonl)." - ) - # Include data_root itself if it is a scene dir. - if os.path.isfile(os.path.join(data_root, "meta", "episodes.jsonl")): - meta_files.append(os.path.join(data_root, "meta", "episodes.jsonl")) - - for meta_file in sorted(set(meta_files)): - scene_dir = os.path.dirname(os.path.dirname(meta_file)) - info_path = os.path.join(scene_dir, "meta", "info.json") - chunks_size = 1000 - if os.path.isfile(info_path): - chunks_size = json.load(open(info_path)).get("chunks_size", 1000) - - # Resolve the rgb stream dir for this scene (prefer the requested key). - video_root = os.path.join(scene_dir, "videos") - chunk_dirs = sorted(glob.glob(os.path.join(video_root, "chunk-*"))) - if not chunk_dirs: - continue - - def _rgb_dir_for_chunk(chunk_dir): - want = os.path.join(chunk_dir, rgb_key) - if os.path.isdir(want): - return want - alts = sorted(glob.glob(os.path.join(chunk_dir, "*.rgb.*"))) - return alts[0] if alts else None - - with open(meta_file) as f: - for line in f: - line = line.strip() - if not line: - continue - ep = json.loads(line) - ep_idx = ep["episode_index"] - length = ep.get("length", 0) - tasks = ep.get("tasks") or [] - if length <= 0 or not tasks: - continue - chunk_dir = os.path.join( - video_root, f"chunk-{ep_idx // chunks_size:03d}" - ) - rgb_dir = _rgb_dir_for_chunk(chunk_dir) or _rgb_dir_for_chunk( - chunk_dirs[0] - ) - if rgb_dir is None: - continue - records.append( - { - "scene_dir": scene_dir, - "ep_idx": ep_idx, - "length": length, - "instruction": tasks[0], - "rgb_dir": rgb_dir, - } - ) - return records - - -def vln_calib_dataloader( - processor, - data_root: str, - num_samples: int = 128, - rgb_key: str = "observation.images.rgb.125cm_0deg", - seed: int = 0, -) -> list[dict[str, Any]]: - """Materialise calibration ``BatchFeature`` dicts from InternData-N1 VLN-CE. - - Each sample is one navigation step: a real instruction + a sequence of - egocentric RGB frames (sub-sampled history + current), assembled through the - SAME path the deployed agent uses (``prompt_builder.build_sample_inputs``), so - the number of history frames, the prompt template and the ```` layout - all match ``InternVLAN1Net.s2_step`` exactly. - - Returns a *list* (not a generator) so ModelOpt can re-iterate the forward loop - during algorithm selection, mirroring ``multimodal_calib_dataloader``. - """ - import os - import random - import shutil - import tempfile - - import numpy as np - from PIL import Image - - pb = _import_prompt_builder() - num_history = pb.NUM_HISTORY # single source of truth — must match the prompt - - def frame_path(rgb_dir, ep_idx, frame): - return os.path.join(rgb_dir, f"episode_{ep_idx:06d}_{frame}.jpg") - - episodes = _discover_vln_episodes(data_root, rgb_key) - if not episodes: - raise RuntimeError( - f"No usable VLN episodes under {data_root!r} (rgb_key={rgb_key!r})." - ) - - # The deployed agent resizes every RGB frame to (resize_w, resize_h) before the - # processor (internvla_n1_policy.py). Match that here so the visual-token count - # and ViT activations track deployment, not the raw camera resolution. Resized - # frames are cached under a temp dir and cleaned up once every batch is built - # (the returned batches hold materialised tensors, not paths). - tmp_dir = tempfile.mkdtemp(prefix="vln_calib_") - resized_cache: dict[str, str] = {} - - def resized_frame(src_path): - cached = resized_cache.get(src_path) - if cached is not None: - return cached - img = Image.open(src_path).convert("RGB").resize( - (pb.RESIZE_W, pb.RESIZE_H) - ) - dst = os.path.join(tmp_dir, f"f{len(resized_cache):06d}.jpg") - img.save(dst, quality=95) - resized_cache[src_path] = dst - return dst - - rng = random.Random(seed) - batches: list[dict[str, Any]] = [] - attempts = 0 - max_attempts = num_samples * 20 - - try: - while len(batches) < num_samples and attempts < max_attempts: - attempts += 1 - ep = rng.choice(episodes) - length = ep["length"] - # Pick a "current" step; prefer t>=1 so the sample carries history. - t = rng.randint(1, length - 1) if length > 1 else 0 - - # History frame indices — identical rule to - # prompt_builder.build_conversation (same NUM_HISTORY constant). - if t == 0: - hist = [] - else: - hist = np.unique( - np.linspace(0, t - 1, num_history, dtype=np.int32) - ).tolist() - - src_paths = [frame_path(ep["rgb_dir"], ep["ep_idx"], h) for h in hist] - src_paths.append(frame_path(ep["rgb_dir"], ep["ep_idx"], t)) - if not all(os.path.isfile(p) for p in src_paths): - continue # frame missing on disk — skip this sample - - sample = { - "images": [resized_frame(p) for p in src_paths], - "episode_idx": t, - "instruction": ep["instruction"], - } - try: - inputs = pb.build_sample_inputs(sample, processor) - except Exception: # noqa: BLE001 — drop malformed samples, keep going - continue - - batches.append( - { - k: v - for k, v in inputs.items() - if v is not None - and not (isinstance(v, torch.Tensor) and v.numel() == 0) - } - ) - finally: - shutil.rmtree(tmp_dir, ignore_errors=True) - - if not batches: - raise RuntimeError( - f"Could not build any VLN calibration sample from {data_root!r} " - f"after {attempts} attempts (rgb_key={rgb_key!r})." - ) - return batches diff --git a/recipes/internvla-n1-dualvln/quantize/compare_fake_quant.py b/recipes/internvla-n1-dualvln/quantize/compare_fake_quant.py deleted file mode 100644 index a0dd5f1..0000000 --- a/recipes/internvla-n1-dualvln/quantize/compare_fake_quant.py +++ /dev/null @@ -1,120 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Prove — or refute — that activation quantization is what the engine loses. - -Three numbers exist for each scheme and until now only two were measured directly: - - weights only weights quantized then reconstructed, activations bf16 - fake quant weights AND activations quantized in PyTorch <- this script - engine the TensorRT engine - -The claim this recipe makes is that the gap between "weights only" and "engine" is -activation quantization rather than anything about the export or the runtime. That has so -far been an inference from three numbers, including the FP16 engine measuring 0.999471, -which bounds TensorRT's own cost at about 0.0005. - -Measuring fake quant closes it directly. If fake quant lands near the engine, the claim -holds. If it lands near weights-only instead, the claim is wrong and the loss is somewhere -in the export or the runtime, which is a different investigation. - -The comparison runs on the same z_latents bridge the rest of the recipe is gated on. -""" -import argparse -import json -import os -import sys - -import torch - -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -import prompt_builder as pb # noqa: E402 -from benchmark_accuracy import discover_samples # noqa: E402 -from load_quantized import load_fake_quant, load_for_eval # noqa: E402 - -PROMPT = ("You are an autonomous navigation assistant. Your task is to go to the kitchen. " - "Where should you go next to stay on track?") - - -def cos(a: torch.Tensor, b: torch.Tensor) -> float: - a = a.double().flatten() - b = b.double().flatten() - return float(a @ b / (a.norm() * b.norm())) - - -def bridge_tensors(ckpt: str, device: str): - from safetensors import safe_open - - path = os.path.join(ckpt, "bridge.safetensors") - with safe_open(path, framework="pt") as f: - raw = {k: f.get_tensor(k) for k in f.keys()} - return {k.replace("model.cond_projector.", ""): v.float().to(device) - for k, v in raw.items() if "cond_projector" in k} - - -def z_latents(model, processor, cond, device: str) -> torch.Tensor: - enc = processor.tokenizer(PROMPT, return_tensors="pt").to(device) - with torch.inference_mode(): - out = model(**enc, output_hidden_states=True) - h = out.hidden_states[-1][0, -4:].float() - x = torch.nn.functional.linear(h, cond["0.weight"], cond.get("0.bias")) - x = torch.nn.functional.gelu(x) - return torch.nn.functional.linear(x, cond["2.weight"], cond.get("2.bias")).cpu() - - -def main() -> int: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--repkg_ckpt", required=True) - p.add_argument("--quant_ckpt", required=True, help="The PTQ checkpoint, for weights-only") - p.add_argument("--scheme", default="nvfp4_default") - p.add_argument("--strategy", default="s1") - p.add_argument("--calib_data_root", required=True) - p.add_argument("--num_calib_samples", type=int, default=16) - p.add_argument("--device", default="cuda") - p.add_argument("--output_path", default=None) - args = p.parse_args() - - cond = bridge_tensors(args.repkg_ckpt, args.device) - - print("[1/3] reference (unquantized)") - ref_model, ref_proc, _ = load_for_eval(args.repkg_ckpt, device=args.device) - z_ref = z_latents(ref_model, ref_proc, cond, args.device) - del ref_model - torch.cuda.empty_cache() - - print("[2/3] weights only (activations left in bf16)") - wo_model, wo_proc, _ = load_for_eval(args.quant_ckpt, device=args.device) - z_wo = z_latents(wo_model, wo_proc, cond, args.device) - del wo_model - torch.cuda.empty_cache() - - print(f"[3/3] fake quant ({args.scheme}, weights AND activations)") - calib = discover_samples(args.calib_data_root, args.num_calib_samples, seed=0) - fq_model, _, fq_proc = load_fake_quant( - args.repkg_ckpt, args.scheme, args.strategy, calib, - prompt_builder=pb.build_sample_inputs, device=args.device) - z_fq = z_latents(fq_model, fq_proc, cond, args.device) - del fq_model - torch.cuda.empty_cache() - - result = { - "weights_only": cos(z_ref, z_wo), - "fake_quant": cos(z_ref, z_fq), - } - print("\n=== z_latents vs the unquantized reference ===") - print(f" weights only : {result['weights_only']:.6f}") - print(f" fake quant : {result['fake_quant']:.6f} <- comparable to the engine") - print("\n Compare 'fake quant' against the engine number for this scheme. Close means") - print(" activation quantization explains the engine's loss; far means it does not and") - print(" the export or runtime is worth investigating instead.") - - if args.output_path: - os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True) - with open(args.output_path, "w") as f: - json.dump(result, f, indent=2) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/quantize/load_quantized.py b/recipes/internvla-n1-dualvln/quantize/load_quantized.py deleted file mode 100644 index 52edf83..0000000 --- a/recipes/internvla-n1-dualvln/quantize/load_quantized.py +++ /dev/null @@ -1,218 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Load a ModelOpt-exported checkpoint back into PyTorch with its scales applied. - -This exists because the obvious thing does not work. ``export_hf_checkpoint`` writes the -LLM weights as real ``torch.float8_e4m3fn`` with the dequantization scales in separate -``*.weight_scale`` / ``*.input_scale`` tensors. Calling -``AutoModelForImageTextToText.from_pretrained`` on that directory *appears* to succeed -- -it warns that the scale tensors "were not used when initializing" and moves on -- but the -resulting model has every quantized weight off by its scale factor. Any accuracy measured -that way is measuring a broken load, not quantization error, and it will look far worse -than the engine actually is. - -So reconstruct the weights explicitly:: - - w_bf16 = w_fp8.to(bfloat16) * weight_scale - -Two caveats worth stating plainly: - -* This reproduces **weight** quantization error only. Real FP8 W8A8 also quantizes - activations, which the TensorRT engine does and this does not. Weight error is the - dominant term and this is the standard PyTorch-side proxy, but a number from here is a - lower bound on the engine's deviation, not a prediction of it. -* Modules excluded from quantization (the vision tower, ``lm_head``) are stored in bf16 - already and pass through untouched, which is the intended behaviour. -""" -import glob -import json -import os -from typing import Optional - -import torch -from safetensors import safe_open - - -# NVFP4 E2M1: 1 sign, 2 exponent, 1 mantissa bit. Sixteen representable values, two packed -# per stored byte, with a per-16-element FP8 block scale and one float32 global scale. -_E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, - -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], dtype=torch.float32) -NVFP4_BLOCK = 16 - - -def unpack_nvfp4(packed: torch.Tensor, block_scale: torch.Tensor, - global_scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: - """Reconstruct a bf16 weight from ModelOpt's packed NVFP4 representation. - - ``packed`` is uint8 of shape [out, in/2]: low nibble first, then high nibble. - ``block_scale`` is FP8 of shape [out, in/16], one scale per 16 input elements. - ``global_scale`` is a single float32 that rescales the whole tensor. - """ - lut = _E2M1.to(packed.device) - low = lut[(packed & 0x0F).long()] - high = lut[(packed >> 4).long()] - # Interleave back to the original width: low nibble is element 2i, high is 2i+1. - out = torch.stack((low, high), dim=-1).reshape(packed.shape[0], -1) - - scale = block_scale.to(torch.float32) * global_scale.to(torch.float32) - scale = scale.repeat_interleave(NVFP4_BLOCK, dim=-1) - if scale.shape[-1] != out.shape[-1]: - raise ValueError(f"NVFP4 block scale expands to {scale.shape[-1]} columns but the " - f"unpacked weight has {out.shape[-1]}") - return (out * scale).to(dtype) - - -def quant_algo(model_path: str) -> Optional[str]: - """Return the quantization algorithm recorded by ModelOpt, or None if unquantized.""" - cfg = os.path.join(model_path, "hf_quant_config.json") - if not os.path.isfile(cfg): - return None - with open(cfg) as f: - return json.load(f).get("quantization", {}).get("quant_algo") - - -def _iter_shards(model_path: str): - for shard in sorted(glob.glob(os.path.join(model_path, "*.safetensors"))): - if os.path.basename(shard) == "bridge.safetensors": - continue - yield shard - - -def dequantize_state_dict(model_path: str, - dtype: torch.dtype = torch.bfloat16) -> dict[str, torch.Tensor]: - """Read a ModelOpt checkpoint and return a plain state dict with scales folded in. - - Weights that were not quantized are returned as stored. - """ - scales: dict[str, torch.Tensor] = {} - raw: dict[str, torch.Tensor] = {} - - for shard in _iter_shards(model_path): - with safe_open(shard, framework="pt") as f: - for key in f.keys(): - tensor = f.get_tensor(key) - if key.endswith(("weight_scale", "input_scale", "weight_scale_2", - "pre_quant_scale")): - scales[key] = tensor - else: - raw[key] = tensor - - out: dict[str, torch.Tensor] = {} - n_dequant = 0 - n_awq = 0 - for key, tensor in raw.items(): - # NVFP4: packed uint8 plus a block scale and a global scale. - block = scales.get(key + "_scale") - glob = scales.get(key + "_scale_2") - if tensor.dtype == torch.uint8 and block is not None and glob is not None: - w = unpack_nvfp4(tensor, block, glob, torch.float32) - # AWQ-lite scales the activations by a per-input-channel s and stores the - # weight pre-divided by it, so that y = (x * s) @ (W / s) reproduces x @ W. - # A plain matmul needs s multiplied back in. Verified empirically on this - # checkpoint against the unquantized weights, since the direction is easy to - # get backwards: multiplying gives cosine 0.9898, leaving it alone 0.9606, - # dividing 0.8089. Skipping this entirely reads as ~190% relative error and - # looks like AWQ being catastrophically bad rather than loaded wrong. - pqs = scales.get(key.replace(".weight", ".pre_quant_scale")) - if pqs is not None: - w = w * pqs.to(torch.float32) - n_awq += 1 - out[key] = w.to(dtype) - n_dequant += 1 - continue - scale = scales.get(key + "_scale") - if scale is not None and tensor.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): - out[key] = tensor.to(torch.float32).mul_(scale.to(torch.float32)).to(dtype) - n_dequant += 1 - elif tensor.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): - # FP8 storage with no scale would silently corrupt the weight; refuse. - raise ValueError(f"{key} is FP8 but has no matching {key}_scale in the checkpoint") - else: - out[key] = tensor.to(dtype) if tensor.is_floating_point() else tensor - - awq_note = f", {n_awq} with an AWQ pre-quant scale folded out" if n_awq else "" - print(f" [load] dequantized {n_dequant} tensors{awq_note}, " - f"{len(out) - n_dequant} passed through unchanged") - return out - - -def load_for_eval(model_path: str, dtype: torch.dtype = torch.bfloat16, - device: str = "cuda"): - """Load a checkpoint for evaluation, applying ModelOpt scales when present. - - Returns ``(model, processor, algo)`` where ``algo`` is the quantization algorithm - string or None. - """ - from transformers import AutoConfig, AutoProcessor, Qwen2_5_VLForConditionalGeneration - - algo = quant_algo(model_path) - processor = AutoProcessor.from_pretrained( - model_path, min_pixels=128 * 28 * 28, max_pixels=2048 * 32 * 32) - - if algo is None: - model = Qwen2_5_VLForConditionalGeneration.from_pretrained( - model_path, torch_dtype=dtype, low_cpu_mem_usage=True) - else: - # Hand transformers the already-dequantized weights instead of letting it read the - # checkpoint. Two reasons it cannot read this itself: FP8 tensors load without - # their scales (silently wrong by 600-1800x), and NVFP4 tensors are packed two - # values per byte, so from_pretrained fails outright on the halved width. Passing - # state_dict= also avoids loading the whole checkpoint twice. - state = dequantize_state_dict(model_path, dtype=dtype) - config = AutoConfig.from_pretrained(model_path) - if hasattr(config, "quantization_config"): - del config.quantization_config - model = Qwen2_5_VLForConditionalGeneration.from_pretrained( - None, config=config, state_dict=state, torch_dtype=dtype, - low_cpu_mem_usage=True) - - model = model.to(device=device, dtype=dtype).eval() - return model, processor, algo - - -def load_fake_quant(base_model_path: str, scheme: str, strategy: str, - calib_samples: list, prompt_builder=None, - dtype: torch.dtype = torch.bfloat16, device: str = "cuda"): - """Load the *unquantized* checkpoint and insert live quantizers. - - Use this, not :func:`load_for_eval`, whenever a PyTorch number is going to be compared - against an engine number. - - The difference matters and is easy to miss. ``load_for_eval`` reconstructs the weights - and loads them into a plain model with no quantizers, so it reproduces **weight** - quantization error only -- activations stay in bf16 and the matmuls are bf16. But NVFP4 - is W4**A4** and FP8 is W8**A8**: the engine also quantizes every activation. Comparing - the two measures different things, and on this model the difference is not small -- - weight-only NVFP4 reads 0.988 while the engine reads 0.931. - - Inserting real quantizers via ``mtq.quantize`` simulates both halves, so the PyTorch - number becomes directly comparable to the engine's. It costs a calibration pass, which - is why the cheaper path still exists for weight-only questions. - - Returns ``(model, tokenizer, processor)``. - """ - import modelopt.torch.quantization as mtq - - from model_loader import load_model - from quant_schemes import build_quant_config - - model, tokenizer, processor = load_model(base_model_path, dtype="bf16", device=device) - quant_cfg = build_quant_config(scheme, strategy) - - def forward_loop(m): - for sample in calib_samples: - with torch.no_grad(): - if prompt_builder is not None: - inputs = prompt_builder(sample, processor) - m(**{k: (v.to(device) if torch.is_tensor(v) else v) - for k, v in inputs.items()}) - else: - enc = tokenizer(sample, return_tensors="pt", truncation=True, - max_length=512) - m(enc["input_ids"].to(device)) - - model = mtq.quantize(model, quant_cfg, forward_loop=forward_loop) - n_quant = sum(1 for n, _ in model.named_modules() if "quantizer" in n.lower()) - print(f" [load] fake-quant active: {n_quant} quantizers " - f"(weights AND activations simulated, as the engine does)") - return model.eval(), tokenizer, processor diff --git a/recipes/internvla-n1-dualvln/quantize/model_loader.py b/recipes/internvla-n1-dualvln/quantize/model_loader.py deleted file mode 100644 index 20bc76e..0000000 --- a/recipes/internvla-n1-dualvln/quantize/model_loader.py +++ /dev/null @@ -1,290 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Model loading, calibration forward loops, and export for Qwen2.5-VL. - -Logic mirrors NVIDIA's reference quantize_and_export implementation, with the -transformers >= 5.x workaround patches kept intact. - -There is deliberately no InternVLA-N1 branch here. ``repackage_system2.py`` runs first -and turns the checkpoint into a stock Qwen2.5-VL, so this module never sees an -``internvla_n1`` config and never needs to import InternNav. Loading the original -checkpoint directly (with ``config.system1 = "none"`` to suppress System 1) also works -and saves the intermediate copy, but it puts a gated third-party repository -- and the -three monkeypatches needed to load it -- on the critical path of every quantization run. -""" - -import os -import shutil -from typing import Optional - -import torch -from tqdm import tqdm -from transformers import ( - AutoModel, - AutoModelForCausalLM, - AutoModelForImageTextToText, - AutoProcessor, - AutoTokenizer, -) - - -# --------------------------------------------------------------------------- # -# Model loading -# --------------------------------------------------------------------------- # -def load_model( - model_dir: str, - dtype: str = "bf16", - device: str = "cuda", -): - """Load model + tokenizer + optional processor via Auto* classes. - - Mirrors NVIDIA's ``_load_model`` for the standard VLM path: tries - ``AutoModelForImageTextToText`` first (so multimodal architectures are - not silently downgraded to their text-only registration), then falls - back to ``AutoModelForCausalLM`` and finally ``AutoModel``. - """ - if dtype == "fp16": - torch_dtype = torch.float16 - elif dtype == "bf16": - torch_dtype = torch.bfloat16 - else: - raise ValueError(f"Unsupported dtype: {dtype!r}") - - tokenizer = AutoTokenizer.from_pretrained( - model_dir, trust_remote_code=True - ) - - try: - processor = AutoProcessor.from_pretrained( - model_dir, - trust_remote_code=True, - min_pixels=128 * 28 * 28, - max_pixels=2048 * 32 * 32, - ) - except Exception: - processor = None - - last_err: Optional[Exception] = None - model = None - for factory in ( - AutoModelForImageTextToText, - AutoModelForCausalLM, - AutoModel, - ): - try: - model = factory.from_pretrained( - model_dir, - torch_dtype=torch_dtype, - trust_remote_code=True, - ).to(device) - break - except (ValueError, KeyError) as e: - last_err = e - if model is None: - raise RuntimeError( - f"Could not load {model_dir} via any AutoModel factory" - ) from last_err - - model.to(torch_dtype) - - # ModelOpt export_hf_checkpoint crashes when architectures is None. - if getattr(model.config, "architectures", None) is None: - model.config.architectures = [type(model).__name__] - - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - - return model, tokenizer, processor - - -# --------------------------------------------------------------------------- # -# Calibration forward loops -# --------------------------------------------------------------------------- # -def calibrate_text(model, dataloader): - """Forward-loop calibration pass for text-only DataLoader. - - Mirrors NVIDIA's ``_calibrate``. - """ - for data in tqdm(dataloader, desc="Calibrating (text)"): - data = data.to(model.device) - model(data) - - -def calibrate_multimodal(model, batches): - """Forward-loop calibration pass for multimodal ``BatchFeature`` dicts. - - Mirrors NVIDIA's ``_calibrate_multimodal``. Float inputs (pixel_values) - inherit the model's dtype; int inputs (input_ids, attention_mask) are - untouched. ``use_cache=False`` matches the reference. - - NaN amax assertions from ModelOpt's internal calibrator are caught and - the offending batch is skipped so calibration can complete on the - remaining valid samples. - """ - device = model.device - model_dtype = next(model.parameters()).dtype - valid_batches = 0 - skipped_nan_batches = 0 - - for batch in tqdm(batches, desc="Calibrating (multimodal)"): - kwargs = {} - for k, v in batch.items(): - if isinstance(v, torch.Tensor): - v = v.to(device) - if v.dtype.is_floating_point: - v = v.to(model_dtype) - kwargs[k] = v - kwargs.setdefault("use_cache", False) - - with torch.no_grad(): - try: - model(**kwargs) - valid_batches += 1 - except AssertionError as exc: - if "detected nan values in amax" in str(exc): - skipped_nan_batches += 1 - continue - raise - - if valid_batches == 0: - raise RuntimeError( - "All multimodal calibration batches were skipped due to NaN amax." - ) - if skipped_nan_batches > 0: - print( - f"[WAR] Skipped {skipped_nan_batches} multimodal calibration " - f"batch(es) with NaN amax." - ) - - -# --------------------------------------------------------------------------- # -# WAR patches (transformers >= 5.x) -# --------------------------------------------------------------------------- # -def normalize_tied_weights_keys(model) -> None: - """WAR for transformers >= 5.x ``_tied_weights_keys`` format change. - - Mirrors NVIDIA's ``_normalize_tied_weights_keys``. Newer transformers - expects each submodule's ``_tied_weights_keys`` attribute to be a - dict-like (so ``modeling_utils._get_tied_weight_keys`` can call ``.keys()``). - Older custom modeling code still declares it as a list, which crashes - ``model.save_pretrained``. - - Convert list-shaped attributes to ``{key: key}`` dicts in place. The dict's - keys exactly match the original list, preserving behavior for downstream - tied-weight tracking. No-op for modules already in the dict format. - """ - for module in model.modules(): - attr = getattr(module, "_tied_weights_keys", None) - if isinstance(attr, list): - module._tied_weights_keys = {k: k for k in attr} - - -def fix_generation_config_for_strict_validate(model) -> None: - """WAR for transformers >= 5.x ``GenerationConfig.validate(strict=True)``. - - Mirrors NVIDIA's ``_fix_generation_config_for_strict_validate``. ModelOpt's - ``export_hf_checkpoint`` calls ``model.save_pretrained`` which runs - ``validate(strict=True)`` on the generation config; that validator rejects - HF checkpoints whose ``generation_config.json`` sets sampling-only kwargs - (``top_p`` / ``top_k`` / ``temperature``) without ``do_sample = True``. - - Force ``do_sample = True`` when any sampling kwarg is present. This only - changes the saved ``generation_config.json``; runtime params are read - separately and are not affected. - """ - gc = getattr(model, "generation_config", None) - if gc is None: - return - sampling_set = ( - getattr(gc, "top_p", None) not in (None, 1.0) - or getattr(gc, "top_k", None) not in (None, 0, 50) - or getattr(gc, "temperature", None) not in (None, 1.0) - ) - if sampling_set and not getattr(gc, "do_sample", False): - gc.do_sample = True - - -# --------------------------------------------------------------------------- # -# Export helpers -# --------------------------------------------------------------------------- # - -# Files copied verbatim from source model dir — preserves preprocessing config -# untouched by calibration. Matches NVIDIA's copy list. -PROCESSOR_FILES = ( - "preprocessor_config.json", - "processor_config.json", - "video_preprocessor_config.json", - "chat_template.jinja", - "chat_template.json", -) - -# Tokenizer files copied verbatim from source. Quantization only rewrites model -# weights (+ config.json / generation_config.json); it never alters the -# tokenizer, so the source tokenizer files are byte-for-byte authoritative. -# Copying them verbatim — instead of trusting tokenizer.save_pretrained() to -# round-trip cleanly — makes it impossible for calibration-time truncation state -# (max_length=512, truncation_strategy, ...) to leak into the deployed -# tokenizer_config.json. That leak silently caps inference sequence length -# (e.g. MME benchmark truncated to 512). Non-existent files are skipped, so -# listing both fast (tokenizer.json) and slow (vocab.json/merges.txt) artifacts -# is safe across tokenizer variants. -TOKENIZER_FILES = ( - "tokenizer_config.json", - "tokenizer.json", - "vocab.json", - "merges.txt", - "special_tokens_map.json", - "added_tokens.json", -) - - -# --------------------------------------------------------------------------- # -# Export -# --------------------------------------------------------------------------- # -def export_quantized_model( - model, - tokenizer, - processor, - model_dir: str, - output_dir: str, -) -> None: - """Save a quantized model with HF-compatible layout. - - Mirrors NVIDIA's export tail: - 1. Apply transformers >= 5.x WARs. - 2. ``export_hf_checkpoint`` writes weights + ``config.json`` + - ``generation_config.json`` + ``hf_quant_config.json``. - 3. Save tokenizer + processor (so the full file set always exists). - 4. Overwrite the tokenizer and preprocessor config files with verbatim - source copies. This is the authoritative step: quantization never - changes these, so the source files are correct by definition, and a - verbatim copy guarantees calibration's ``max_length=512`` truncation - state can never leak into ``tokenizer_config.json``. Only allow-listed - filenames are overwritten — the quantized weights, ``config.json``, - ``generation_config.json`` and ``hf_quant_config.json`` from step 2 are - left untouched. - """ - from modelopt.torch.export import export_hf_checkpoint - - fix_generation_config_for_strict_validate(model) - normalize_tied_weights_keys(model) - - os.makedirs(output_dir, exist_ok=True) - - with torch.inference_mode(): - export_hf_checkpoint(model, export_dir=output_dir) - - tokenizer.save_pretrained(output_dir) - - if processor is not None: - processor.save_pretrained(output_dir) - - # Overwrite tokenizer + processor configs with verbatim source copies so - # quantization/calibration cannot alter them. This is the authoritative - # final write: it runs after both save_pretrained() calls and only touches - # the allow-listed filenames, so weights, config.json, generation_config.json - # and hf_quant_config.json produced by export_hf_checkpoint are untouched. - for fname in (*TOKENIZER_FILES, *PROCESSOR_FILES): - src = os.path.join(model_dir, fname) - if os.path.isfile(src): - shutil.copy2(src, os.path.join(output_dir, fname)) diff --git a/recipes/internvla-n1-dualvln/quantize/prompt_builder.py b/recipes/internvla-n1-dualvln/quantize/prompt_builder.py deleted file mode 100644 index 6d7a217..0000000 --- a/recipes/internvla-n1-dualvln/quantize/prompt_builder.py +++ /dev/null @@ -1,237 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Prompt / conversation builders for InternVLA-N1 System 2. - -Reconstructs the exact multi-image chat prompt that ``InternVLAN1Net.s2_step`` builds, for both -the normal turn and the look-down (turn 2) follow-up. Kept as a single source of truth so the -verification and calibration scripts never diverge from the agent's real prompt format. - -Read-only: only a processor/tokenizer is loaded; no model weights, no writes to any source tree. -""" -from __future__ import annotations - -import argparse -import json -import os -import sys -from typing import Any - -import numpy as np -from PIL import Image -from transformers import AutoProcessor - -# Eval config values (InternNav h1_internvla_n1_async_cfg). -RESIZE_W = 384 -RESIZE_H = 384 -NUM_HISTORY = 8 - -# Special token ids (internvla_n1.py). -IMAGE_TOKEN_INDEX = 151655 -TRAJ_TOKEN_INDEX = 151667 - -# Base prompt template (internvla_n1_policy.py), verbatim. -PROMPT_TEMPLATE = ( - "You are an autonomous navigation assistant. Your task is to . " - "Where should you go next to stay on track? Please output the next waypoint's " - "coordinates in the image. Please output STOP when you have successfully " - "completed the task." -) -CONJUNCTION = "you can see " -DEFAULT_IMAGE_TOKEN = "" - -SAMPLE_INSTRUCTION = ( - "walk out of the bathroom and turn left, then walk down the hallway and " - "stop in front of the second door on your right" -) - -# Two processor construction variants: the deploy default vs the calibration override. -PROCESSOR_VARIANTS: dict[str, dict[str, Any]] = { - "deploy": {}, - "calib": {"min_pixels": 128 * 28 * 28, "max_pixels": 2048 * 32 * 32}, -} - - -def split_and_clean(text: str) -> list[str]: - """Split around '' and drop empty parts. - - Mirrors ``internnav.model.utils.vln_utils.split_and_clean`` (re-implemented so this module - runs without the full internnav dependency tree). - """ - import re - - parts = re.split(r"()", text) - return [p for p in (s.strip() if s != DEFAULT_IMAGE_TOKEN else s for s in parts) if p] - - -def build_conversation(episode_idx: int, instruction: str) -> tuple[list[dict], int]: - """Build the normal-turn prompt (look_down=False), matching ``internvla_n1_policy.py``. - - Returns ``(conversation, n_images)``. - """ - sources_value = PROMPT_TEMPLATE.replace(".", instruction) - - if episode_idx == 0: - history_id: list[int] = [] - else: - history_id = np.unique( - np.linspace(0, episode_idx - 1, NUM_HISTORY, dtype=np.int32) - ).tolist() - placeholder = (DEFAULT_IMAGE_TOKEN + "\n") * len(history_id) - sources_value += f" These are your historical observations: {placeholder}." - - # The current frame is always appended last. - sources_value += f" {CONJUNCTION}{DEFAULT_IMAGE_TOKEN}." - - n_images = len(history_id) + 1 - - content: list[dict] = [] - for part in split_and_clean(sources_value): - if part == DEFAULT_IMAGE_TOKEN: - content.append({"type": "image", "image": None}) # placeholder - else: - content.append({"type": "text", "text": part}) - - return [{"role": "user", "content": content}], n_images - - -def build_conversation_lookdown( - episode_idx: int, instruction: str, assistant_reply: str -) -> list[dict]: - """Build the turn-2 (look-down) conversation, matching the agent. - - The real flow is two turns: turn 1 returns '↓' (action 5 = look down); the robot tilts its - camera; turn 2 sends the extra downward view **without resetting history**, and the coordinate - is produced here. The turn-2 user message repeats no instruction/history — only - " you can see ." — and only the newly appended look-down image binds to it. - """ - turn1, _ = build_conversation(episode_idx, instruction) - value = f" {CONJUNCTION}{DEFAULT_IMAGE_TOKEN}." - - content: list[dict] = [] - for part in split_and_clean(value): - if part == DEFAULT_IMAGE_TOKEN: - content.append({"type": "image", "image": None}) - else: - content.append({"type": "text", "text": part}) - - return [ - turn1[0], - {"role": "assistant", "content": [{"type": "text", "text": assistant_reply}]}, - {"role": "user", "content": content}, - ] - - -def build_sample_inputs(sample: dict, processor): - """Build processor ``inputs`` for one golden sample, handling both turn 1 and turn 2. - - Single source of truth for gate / calibration / profiling: every hand-rebuilt prompt is a - place that can drift from real behavior. ``turn`` defaults to 1. - """ - from PIL import Image - - turn = sample.get("turn", 1) - images = [Image.open(p).convert("RGB") for p in sample["images"]] - - if turn == 2: - reply = sample.get("assistant_turn1") - if not reply: - raise ValueError( - f"Turn-2 sample missing 'assistant_turn1' " - f"({sample['episode']} ep{sample['episode_idx']})." - ) - conv = build_conversation_lookdown( - sample["episode_idx"], sample["instruction"], reply) - else: - conv, _ = build_conversation(sample["episode_idx"], sample["instruction"]) - - j = 0 - for t in conv: - for it in t["content"]: - if it["type"] == "image" and it.get("image") is None: - it["image"] = images[j] - j += 1 - if j != len(images): - raise ValueError( - f" token count ({j}) != number of images ({len(images)}) — " - f"{sample['episode']} ep{sample['episode_idx']} turn {turn}" - ) - - text = processor.apply_chat_template(conv, tokenize=False, add_generation_prompt=True) - return processor(text=[text], images=images, return_tensors="pt") - - -def census_one(processor, episode_idx: int, instruction: str) -> dict[str, Any]: - """Count visual vs text tokens for one episode step (dummy images sized to the resize).""" - conversation, n_images = build_conversation(episode_idx, instruction) - - images = [ - Image.fromarray(np.zeros((RESIZE_H, RESIZE_W, 3), dtype=np.uint8)).convert("RGB") - for _ in range(n_images) - ] - img_i = 0 - for item in conversation[0]["content"]: - if item["type"] == "image": - item["image"] = images[img_i] - img_i += 1 - - text = processor.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True) - inputs = processor(text=[text], images=images, return_tensors="pt") - - ids = inputs["input_ids"][0] - total = int(ids.numel()) - n_visual = int((ids == IMAGE_TOKEN_INDEX).sum()) - n_traj = int((ids == TRAJ_TOKEN_INDEX).sum()) - grid = inputs.get("image_grid_thw") - return { - "episode_idx": episode_idx, - "n_images": n_images, - "total_tokens": total, - "visual_tokens": n_visual, - "text_tokens": total - n_visual - n_traj, - "traj_tokens": n_traj, - "visual_pct": round(100.0 * n_visual / total, 2), - "tokens_per_image": (n_visual // n_images) if n_images else 0, - "image_grid_thw": grid.tolist() if grid is not None else None, - } - - -def main() -> int: - """Token census: report the visual/text token split the LLM actually sees at runtime.""" - ap = argparse.ArgumentParser(description=main.__doc__) - ap.add_argument("--model-path", default=os.environ.get( - "INTERNVLA_CKPT", os.path.expanduser("~/InternNav/checkpoints/InternVLA-N1-DualVLN"))) - ap.add_argument("--episodes", type=int, nargs="+", default=[0, 1, 2, 4, 8, 16, 32, 64]) - ap.add_argument("--out", default="token_census.json") - args = ap.parse_args() - - if not os.path.isdir(args.model_path): - print(f"ERROR: model-path not found: {args.model_path}", file=sys.stderr) - return 1 - - report: dict[str, Any] = {"model_path": args.model_path, "variants": {}} - for variant, kwargs in PROCESSOR_VARIANTS.items(): - print(f"\n### processor variant: {variant} {kwargs or '(default)'}") - try: - processor = AutoProcessor.from_pretrained( - args.model_path, trust_remote_code=True, **kwargs) - except Exception as exc: # noqa: BLE001 - print(f" !! could not load processor: {exc}") - report["variants"][variant] = {"error": str(exc)} - continue - rows = [census_one(processor, ep, SAMPLE_INSTRUCTION) for ep in args.episodes] - for r in rows: - print(f" ep={r['episode_idx']:>3} imgs={r['n_images']:>2} " - f"total={r['total_tokens']:>6} visual={r['visual_tokens']:>6} " - f"visual%={r['visual_pct']:>6.2f}") - report["variants"][variant] = {"processor_kwargs": kwargs, "rows": rows} - - os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True) - with open(args.out, "w") as f: - json.dump(report, f, indent=2) - print(f"\nWrote: {args.out}") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/quantize/qat.py b/recipes/internvla-n1-dualvln/quantize/qat.py deleted file mode 100644 index 0006702..0000000 --- a/recipes/internvla-n1-dualvln/quantize/qat.py +++ /dev/null @@ -1,342 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Quantization-aware fine-tuning for the InternVLA-N1 System 2 planner. - -Why this exists. Post-training NVFP4 leaves the System 2 -> System 1 bridge at z_latents -0.931 through the engine, under the 0.99 gate that FP8 clears. Weight quantization is not -what costs it -- NVFP4 weights alone measure 0.988 -- so the remaining loss sits in the -4-bit *activations*, which no amount of weight-side scaling can fix after the fact. QAT is -the standard answer: let the model see the quantization noise during training and adapt to -it. - -What this does. ModelOpt has no separate QAT entry point; QAT is ``mtq.quantize`` followed -by ordinary fine-tuning, with the fake-quantizers left in place so gradients flow through -them (straight-through estimation). So this: - - 1. loads the repackaged System 2, - 2. calibrates and inserts quantizers exactly as ``quantize.py`` would, - 3. fine-tunes on real VLN episodes with the deployed prompt, - 4. exports through the same path, producing a checkpoint the existing engine build and - verification scripts accept unchanged. - -Read the result honestly. **Success rate is not measurable here** -- SR, SPL and NE are all -closed-loop and need Habitat or InternUtopia, neither of which runs on a Jetson. What this -optimises and what you can check is the proxy: z_latents cosine and pixel-goal L2 on -held-out episodes. A proxy that improves is a necessary condition for SR to improve, not -evidence that it did. - -Train/eval separation. The calibration and probe sets share one MP3D scene -(``YmJkqBEsHnH`` appears as ``calib_scenes/r2r/`` and ``probe_heldout/rxr/`` -- same -building, different split), and the source project has already been bitten by exactly this -overlap once, reporting a calibration gain that turned out to be leakage. That scene is -excluded from training by default; ``--allow_overlap`` turns the guard off if you want to -measure the effect of the leak deliberately. -""" -import argparse -import json -import math -import os -import sys -import time - -import numpy as np -import torch - -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -import prompt_builder as pb # noqa: E402 -from model_loader import export_quantized_model, load_model # noqa: E402 -from quant_schemes import build_quant_config, calib_batch_size, validate # noqa: E402 - -# Same building as probe_heldout/rxr/YmJkqBEsHnH. Training on it contaminates the -# held-out evaluation even though the episodes and instructions differ. -OVERLAPPING_SCENES = ("YmJkqBEsHnH",) - - -def discover_training_samples(data_root: str, max_samples: int, camera: str, - allow_overlap: bool, seed: int = 0, - samples_per_episode: int = 1) -> list[dict]: - """Collect prompt/target pairs from LeRobot episodes, skipping leaked scenes.""" - import glob - import pyarrow.parquet as pq - - rng = np.random.default_rng(seed) - rgb_key = f"observation.images.rgb.{camera}" - level_key = "observation.images.rgb.125cm_0deg" - goal_col = f"goal.{camera}" - samples: list[dict] = [] - skipped_scenes: set[str] = set() - - for meta in sorted(glob.glob(os.path.join(data_root, "**", "meta", "episodes.jsonl"), - recursive=True)): - scene_dir = os.path.dirname(os.path.dirname(meta)) - scene = os.path.basename(scene_dir) - if not allow_overlap and scene in OVERLAPPING_SCENES: - skipped_scenes.add(scene) - continue - - for ep in (json.loads(line) for line in open(meta)): - idx, length = ep["episode_index"], ep["length"] - table = None - for parquet in glob.glob(os.path.join(scene_dir, "data", "**", - f"episode_{idx:06d}.parquet"), - recursive=True): - table = pq.read_table(parquet) - break - if table is None or goal_col not in table.schema.names: - continue - - goals = table.column(goal_col).to_pylist() - usable = [i for i in range(1, min(length, len(goals))) - if goals[i] is not None and int(goals[i][0]) >= 0] - if not usable: - continue - # One sample per episode caps the set at ~91, too few for a few hundred - # optimizer steps without looping over the same prompts a dozen times. Each - # episode carries many annotated frames, so drawing several spreads the data - # over genuinely different observations rather than repeating one. - k = min(samples_per_episode, len(usable)) - picks = rng.choice(len(usable), size=k, replace=False) - for pick in picks: - t = int(usable[int(pick)]) - - history = np.unique( - np.linspace(0, t - 1, pb.NUM_HISTORY, dtype=np.int32)).tolist() - frames = [os.path.join(scene_dir, "videos", "chunk-000", level_key, - f"episode_{idx:06d}_{i}.jpg") for i in history + [t]] - lookdown = os.path.join(scene_dir, "videos", "chunk-000", rgb_key, - f"episode_{idx:06d}_{t}.jpg") - if not all(os.path.isfile(p) for p in frames + [lookdown]): - continue - - gt = goals[t] - samples.append({ - "episode": scene, - "episode_idx": t, - "instruction": (ep.get("tasks") or [""])[0], - "images": frames + [lookdown], - "turn": 2, - "assistant_turn1": "↓", - # The supervision target is the deployed answer format: "row col". - "target": f"{int(gt[0])} {int(gt[1])}", - }) - if len(samples) >= max_samples: - break - if len(samples) >= max_samples: - break - if len(samples) >= max_samples: - break - - if skipped_scenes: - print(f" [data] excluded {sorted(skipped_scenes)} -- also present in the " - f"held-out probe set") - return samples - - -def build_batch(sample: dict, processor, device: str): - """Prompt plus target, with the prompt masked out of the loss.""" - inputs = pb.build_sample_inputs(sample, processor) - prompt_len = inputs["input_ids"].shape[1] - - target_ids = processor.tokenizer(sample["target"], add_special_tokens=False, - return_tensors="pt")["input_ids"] - input_ids = torch.cat([inputs["input_ids"], target_ids], dim=1) - labels = input_ids.clone() - labels[:, :prompt_len] = -100 # supervise the answer only - - batch = {k: v for k, v in inputs.items() if k != "input_ids"} - batch["input_ids"] = input_ids - batch["labels"] = labels - if "attention_mask" in batch: - pad = torch.ones((1, target_ids.shape[1]), dtype=batch["attention_mask"].dtype) - batch["attention_mask"] = torch.cat([batch["attention_mask"], pad], dim=1) - return {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in batch.items()} - - -def freeze_all_but_last_layers(model, n_last: int) -> int: - """Freeze everything except the final ``n_last`` decoder layers. - - Returns the number of frozen parameters. Full fine-tuning does not fit: at 8.29B - parameters, weights plus gradients plus AdamW moments come to about 100 GB before any - activation memory, against a 122 GB pool shared with the host. Restricting to the last - layers is also where the quantity being repaired lives -- the System 1 bridge reads the - last layer's hidden states. - """ - layers = None - for owner in (getattr(model, "model", None), model): - inner = getattr(owner, "language_model", owner) - if inner is not None and hasattr(inner, "layers"): - layers = inner.layers - break - if layers is None: - raise AttributeError("could not locate the decoder layer list on this model") - - keep = {id(p) for layer in layers[-n_last:] for p in layer.parameters()} - frozen = 0 - for param in model.parameters(): - if id(param) not in keep: - param.requires_grad_(False) - frozen += param.numel() - return frozen - - -def parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--model_path", required=True, help="Repackaged System 2 checkpoint") - p.add_argument("--output_path", required=True) - p.add_argument("--data_root", required=True, help="LeRobot episodes for training") - p.add_argument("--scheme", default="nvfp4_default") - p.add_argument("--strategy", default="s1") - p.add_argument("--num_train_samples", type=int, default=64) - p.add_argument("--num_calib_samples", type=int, default=64) - p.add_argument("--epochs", type=int, default=1) - p.add_argument("--lr", type=float, default=1e-5, - help="Small on purpose: QAT adapts to quantization noise, it does not " - "re-learn the task, and a large step undoes the pretrained planner") - p.add_argument("--grad_accum", type=int, default=4) - p.add_argument("--samples_per_episode", type=int, default=1, - help="Timesteps drawn per episode. The calibration set has ~91\n" - "episodes, so 1 caps training far below what a few hundred\n" - "steps needs.") - p.add_argument("--warmup_frac", type=float, default=0.1, - help="Fraction of steps spent warming the learning rate up, then\n" - "cosine-decayed. A flat rate from step 0 is what made the\n" - "first run diverge.") - p.add_argument("--train_last_n_layers", type=int, default=4, - help="Train only the last N decoder layers; freeze the rest. Full " - "fine-tuning of the 8.29B planner needs ~100 GB for weights, " - "gradients and AdamW state alone, which the 122 GB unified pool " - "cannot hold alongside a 10-image prompt's activations. The bridge " - "reads the last layer's hidden states, so the last layers are also " - "where the signal it depends on is formed. 0 trains everything.") - p.add_argument("--gradient_checkpointing", action="store_true", default=True) - p.add_argument("--camera", default="125cm_30deg") - p.add_argument("--allow_overlap", action="store_true", - help="Permit training on scenes that also appear in the probe set") - p.add_argument("--device", default="cuda") - p.add_argument("--seed", type=int, default=0) - return p.parse_args() - - -def main() -> int: - import modelopt.torch.quantization as mtq - - args = parse_args() - torch.manual_seed(args.seed) - - try: - validate(args.scheme, args.strategy, model_path=args.model_path, - allow_experimental=True) - except ValueError as exc: - print(f"[ERROR] {exc}") - return 1 - - print(f"[1/5] Loading {args.model_path}") - model, tokenizer, processor = load_model(args.model_path, dtype="bf16", - device=args.device) - if processor is None: - print("[ERROR] no processor; VLN prompts cannot be built") - return 1 - - print(f"[2/5] Collecting training samples from {args.data_root}") - train = discover_training_samples(args.data_root, args.num_train_samples, - args.camera, args.allow_overlap, args.seed, - args.samples_per_episode) - if not train: - print("[ERROR] no usable training samples") - return 1 - print(f" {len(train)} samples from " - f"{len({s['episode'] for s in train})} scene(s)") - - print(f"[3/5] Inserting quantizers ({args.scheme} / {args.strategy})") - quant_cfg = build_quant_config(args.scheme, args.strategy) - calib = train[:min(args.num_calib_samples, len(train))] - - def forward_loop(m): - for s in calib: - with torch.no_grad(): - m(**build_batch(s, processor, args.device)) - - model = mtq.quantize(model, quant_cfg, forward_loop=forward_loop) - _ = calib_batch_size(args.scheme, is_image_calib=True) - - print(f"[4/5] Fine-tuning: {args.epochs} epoch(s), lr={args.lr}, " - f"grad_accum={args.grad_accum}") - # Gradients flow through the fake-quantizers by straight-through estimation, which is - # what lets the weights adapt to 4-bit activation noise rather than merely to 4-bit - # weights. - model.train() - model.config.use_cache = False - - if args.train_last_n_layers > 0: - n_frozen = freeze_all_but_last_layers(model, args.train_last_n_layers) - print(f" froze {n_frozen} parameters; training the last " - f"{args.train_last_n_layers} decoder layers") - if args.gradient_checkpointing and hasattr(model, "gradient_checkpointing_enable"): - # With the first layers frozen, the activations entering the first trainable layer - # carry no grad_fn, and reentrant checkpointing then drops the graph entirely -- - # loss.backward() fails with "element 0 of tensors does not require grad". Two - # fixes are needed together: non-reentrant checkpointing, and forcing the input - # embeddings to require grad so a graph exists from the start. - model.gradient_checkpointing_enable( - gradient_checkpointing_kwargs={"use_reentrant": False}) - if hasattr(model, "enable_input_require_grads"): - model.enable_input_require_grads() - - trainable = [p for p in model.parameters() if p.requires_grad] - n_train = sum(p.numel() for p in trainable) - print(f" trainable: {n_train / 1e6:.0f} M parameters " - f"(~{n_train * 12 / 1e9:.1f} GB for weights, grads and AdamW state)") - optim = torch.optim.AdamW(trainable, lr=args.lr, weight_decay=0.0) - - total_steps = max(1, args.epochs * len(train) // args.grad_accum) - warmup_steps = max(1, int(total_steps * args.warmup_frac)) - - def lr_at(step: int) -> float: - """Linear warmup then cosine decay. - - The first attempt ran a flat rate from step 0 and the loss rose monotonically. - Warming up matters more than usual here: the fake-quantizers were only just - calibrated, so the first gradients are the noisiest ones the run will see. - """ - if step < warmup_steps: - return args.lr * (step + 1) / warmup_steps - progress = (step - warmup_steps) / max(1, total_steps - warmup_steps) - return args.lr * 0.5 * (1.0 + math.cos(math.pi * min(1.0, progress))) - - print(f" {total_steps} optimizer steps planned, {warmup_steps} of them warmup") - step, t0, losses = 0, time.time(), [] - for epoch in range(args.epochs): - order = np.random.default_rng(args.seed + epoch).permutation(len(train)) - optim.zero_grad(set_to_none=True) - for n, i in enumerate(order, 1): - out = model(**build_batch(train[i], processor, args.device)) - loss = out.loss / args.grad_accum - loss.backward() - losses.append(float(out.loss)) - if n % args.grad_accum == 0: - for group in optim.param_groups: - group["lr"] = lr_at(step) - torch.nn.utils.clip_grad_norm_(trainable, 1.0) - optim.step() - optim.zero_grad(set_to_none=True) - step += 1 - if step % 10 == 0: - recent = float(np.mean(losses[-10 * args.grad_accum:])) - print(f" epoch {epoch} step {step:4d} loss {recent:.4f} " - f"lr {lr_at(step):.2e} {time.time() - t0:.0f}s", flush=True) - - print(f" trained {step} optimizer steps, final loss " - f"{float(np.mean(losses[-8:])):.4f}") - - print(f"[5/5] Exporting to {args.output_path}") - model.eval() - export_quantized_model(model, tokenizer, processor, - model_dir=args.model_path, output_dir=args.output_path) - print(f"Saved to {args.output_path}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/quantize/quant_schemes.py b/recipes/internvla-n1-dualvln/quantize/quant_schemes.py deleted file mode 100644 index 1e04a7a..0000000 --- a/recipes/internvla-n1-dualvln/quantize/quant_schemes.py +++ /dev/null @@ -1,224 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Quantization scheme registry, validity gate, and ModelOpt config composer. - -Schemes and strategies live in ``configs/schemes.yaml`` rather than in Python so the -matrix is readable without tracing code, following the ``qwen36-27b`` recipe's pattern. - -The validity gate exists because two combinations are impossible on this hardware and -neither fails early on its own: - -* NVFP4 cannot quantize the Qwen2.5-VL vision tower. The ViT MLP ``intermediate_size`` - is 3420 and the NVFP4 block size is 16; 3420 / 16 = 213.75. Without a gate this only - surfaces once quantization is already under way. -* NVFP4 KV cache requires ``sm100f`` (datacenter Blackwell). Jetson Thor is sm110, so the - KV cache is FP8 even when the weights are NVFP4. - -The divisibility check reads ``vision_config.intermediate_size`` from the checkpoint being -quantized rather than hardcoding 3420, so it stays correct for other Qwen2.5-VL sizes. -""" -import copy -import json -import os -from typing import Any, Optional - -import yaml - -_DEFAULT_SCHEMES_YAML = os.path.join( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "configs", "schemes.yaml" -) - -# NVFP4 packs weights in blocks of this many elements along the reduction axis. -NVFP4_BLOCK_SIZE = 16 - - -def load_registry(path: Optional[str] = None) -> dict: - """Load ``configs/schemes.yaml``.""" - path = path or _DEFAULT_SCHEMES_YAML - if not os.path.isfile(path): - raise FileNotFoundError(f"scheme registry not found: {path}") - with open(path) as f: - return yaml.safe_load(f) - - -def scheme_names(registry: Optional[dict] = None) -> list[str]: - return sorted((registry or load_registry())["schemes"]) - - -def strategy_names(registry: Optional[dict] = None) -> list[str]: - return sorted((registry or load_registry())["strategies"]) - - -def is_nvfp4(scheme: str) -> bool: - return scheme.startswith("nvfp4") - - -def _vision_intermediate_size(model_path: str) -> Optional[int]: - """Read ``vision_config.intermediate_size`` from a checkpoint, if it has one.""" - cfg_path = os.path.join(model_path, "config.json") - if not os.path.isfile(cfg_path): - return None - with open(cfg_path) as f: - cfg = json.load(f) - vision = cfg.get("vision_config") - if isinstance(vision, dict): - return vision.get("intermediate_size") - return None - - -def validate(scheme: str, strategy: str, model_path: Optional[str] = None, - allow_experimental: bool = False, - registry: Optional[dict] = None) -> None: - """Raise ``ValueError`` if this combination cannot or should not run. - - Called before the model is loaded, so a rejected combination costs seconds rather - than a full checkpoint load followed by a mid-quantization crash. - """ - registry = registry or load_registry() - - if scheme not in registry["schemes"]: - raise ValueError(f"unknown scheme {scheme!r}; available: {scheme_names(registry)}") - if strategy not in registry["strategies"]: - raise ValueError(f"unknown strategy {strategy!r}; available: {strategy_names(registry)}") - - strat = registry["strategies"][strategy] - - for rule in registry.get("blocked", []): - if scheme in rule["schemes"] and strategy in rule["strategies"]: - detail = "" - # Prefer the checkpoint's own number over the one in the message. - if strat["quantize_visual"] and model_path: - size = _vision_intermediate_size(model_path) - if size is not None: - detail = (f" This checkpoint's vision_config.intermediate_size is {size}; " - f"{size} / {NVFP4_BLOCK_SIZE} = {size / NVFP4_BLOCK_SIZE}.") - raise ValueError(f"{scheme} + {strategy} is not supported. " - f"{rule['reason'].strip()}{detail}") - - for rule in registry.get("experimental", []): - if scheme in rule["schemes"] and strategy in rule["strategies"]: - if not allow_experimental: - raise ValueError(f"{scheme} + {strategy} is experimental. " - f"{rule['reason'].strip()}") - print(f"[WARN] {scheme} + {strategy} is experimental. {rule['reason'].strip()}") - - -def build_quant_config(scheme: str, strategy: str, - layerwise_checkpoint_dir: Optional[str] = None, - registry: Optional[dict] = None) -> dict: - """Compose a ModelOpt quant_cfg from a scheme preset plus a strategy. - - Mirrors NVIDIA's ``build_quant_config`` pattern: start from a base preset, then merge - in the KV-cache entries and append visual exclusions as the strategy dictates. - """ - import modelopt.torch.quantization as mtq - - registry = registry or load_registry() - scheme_cfg = registry["schemes"][scheme] - strat = registry["strategies"][strategy] - - preset_name = scheme_cfg["modelopt_preset"] - if not hasattr(mtq, preset_name): - raise ValueError(f"modelopt has no preset {preset_name!r} " - f"(scheme {scheme!r}); check the installed nvidia-modelopt version") - quant_cfg = copy.deepcopy(getattr(mtq, preset_name)) - - if strat["quantize_kv_cache"]: - kv_name = registry["kv_cache"]["preset"] - if scheme not in registry["kv_cache"]["applies_to"]: - raise ValueError(f"no KV-cache preset registered for scheme {scheme!r}") - # FP8 KV for every weight format; see the note in schemes.yaml. - kv_cfg = getattr(mtq, kv_name) - quant_cfg["quant_cfg"] = quant_cfg["quant_cfg"] + kv_cfg["quant_cfg"] - - if not strat["quantize_visual"]: - # ModelOpt presets exclude *lm_head* by default but not the vision tower, so - # without these the ViT would be quantized on s1/s2 without anyone asking. - for pattern in registry["visual_exclude_patterns"]: - quant_cfg["quant_cfg"].append({"quantizer_name": pattern, "enable": False}) - - if layerwise_checkpoint_dir is not None: - algo: Any = quant_cfg.get("algorithm") - if isinstance(algo, str): - algo = {"method": algo} - elif algo is None: - algo = {} - elif isinstance(algo, dict): - algo = dict(algo) - else: - raise TypeError(f"unexpected algorithm type: {type(algo)}") - algo["layerwise"] = {"enable": True, "checkpoint_dir": layerwise_checkpoint_dir} - quant_cfg["algorithm"] = algo - - return quant_cfg - - -def calib_batch_size(scheme: str, is_image_calib: bool, - registry: Optional[dict] = None) -> int: - """Calibration batch size. Image calibration is always 1 (GPU-memory bound).""" - if is_image_calib: - return 1 - registry = registry or load_registry() - return int(registry["schemes"][scheme].get("calib_batch_size", 1)) - - -def render_matrix(registry: Optional[dict] = None) -> str: - """Render the scheme x strategy matrix for pasting into a README.""" - registry = registry or load_registry() - strategies = strategy_names(registry) - blocked = {(s, st) for r in registry.get("blocked", []) - for s in r["schemes"] for st in r["strategies"]} - experimental = {(s, st) for r in registry.get("experimental", []) - for s in r["schemes"] for st in r["strategies"]} - - lines = ["| scheme | " + " | ".join(strategies) + " |", - "|---" * (len(strategies) + 1) + "|"] - for scheme in scheme_names(registry): - cells = [] - for strategy in strategies: - if (scheme, strategy) in blocked: - cells.append("blocked") - elif (scheme, strategy) in experimental: - cells.append("experimental") - else: - cells.append("yes") - lines.append(f"| `{scheme}` | " + " | ".join(cells) + " |") - return "\n".join(lines) - - -def describe() -> str: - """Human-readable listing of schemes and strategies, for --help epilogs.""" - registry = load_registry() - out = ["schemes:"] - for name in scheme_names(registry): - cfg = registry["schemes"][name] - out.append(f" {name:22s} {cfg['description']} [{cfg['status']}]") - out.append("strategies:") - for name in strategy_names(registry): - out.append(f" {name:22s} {registry['strategies'][name]['description']}") - return "\n".join(out) - - -def main() -> int: - import argparse - - parser = argparse.ArgumentParser(description="Inspect the quantization scheme registry") - parser.add_argument("--list_scheme_names", action="store_true") - parser.add_argument("--list_strategy_names", action="store_true") - parser.add_argument("--print_matrix", action="store_true") - parser.add_argument("--describe", action="store_true") - args = parser.parse_args() - - if args.list_scheme_names: - print("\n".join(scheme_names())) - elif args.list_strategy_names: - print("\n".join(strategy_names())) - elif args.print_matrix: - print(render_matrix()) - else: - print(describe()) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/quantize/quantize.py b/recipes/internvla-n1-dualvln/quantize/quantize.py deleted file mode 100755 index c0eca10..0000000 --- a/recipes/internvla-n1-dualvln/quantize/quantize.py +++ /dev/null @@ -1,323 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Quantize Qwen2.5-VL with NVIDIA ModelOpt. - -Usage: - python quantize.py --strategy s1 --cfg fp8_default - python quantize.py --strategy s3 --cfg nvfp4_awq_full --dry-run - -Supports 4 strategies x 5 schemes; see ``configs/schemes.yaml`` for the matrix -and which combinations are blocked or experimental on this hardware. -""" - -import argparse -import os -import sys -import time - -import modelopt.torch.quantization as mtq -import torch - -from calibration import ( - multimodal_calib_dataloader, - text_calib_dataloader, - vln_calib_dataloader, -) -from quant_schemes import ( - build_quant_config, - calib_batch_size, - describe, - load_registry, - scheme_names, - strategy_names, - validate, -) -from model_loader import ( - calibrate_multimodal, - calibrate_text, - export_quantized_model, - load_model, -) - - -# --------------------------------------------------------------------------- # -# Defaults -# --------------------------------------------------------------------------- # -DEFAULT_MODEL_DIR = os.path.expanduser( - os.environ.get("INTERNVLA_CKPT", "~/InternNav/checkpoints/InternVLA-N1-DualVLN") -) -DEFAULT_OUTPUT_BASE = os.path.expanduser( - os.environ.get("VLN_OPT_WORK", "~/vln-opt-work") -) -TEXT_DATASET = "abisee/cnn_dailymail" -MULTIMODAL_DATASET = "lmms-lab/MMMU" -MULTIMODAL_MAX_SAMPLES = 128 # NVIDIA's cap — VLM calibration is GPU-mem bound -# In-distribution calibration: the InternData-N1 VLN-CE subset System 2 was -# fine-tuned on. Override with --calib-data or $VLN_CALIB_DATA. -VLN_CALIB_DATA = os.path.expanduser( - os.environ.get( - # Default = the output of build/00_fetch_calib_scenes.sh. Point it at any - # InternData-N1 VLN-CE tree (a scene dir or a parent of several) via the env var. - "VLN_CALIB_DATA", - "~/vln-opt-work/calib_scenes", - ) -) -VLN_RGB_KEY = os.environ.get("VLN_RGB_KEY", "observation.images.rgb.125cm_0deg") - - -# --------------------------------------------------------------------------- # -# CLI -# --------------------------------------------------------------------------- # -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Quantize Qwen2.5-VL with NVIDIA ModelOpt", - formatter_class=argparse.RawDescriptionHelpFormatter, - epilog=describe(), - ) - parser.add_argument( - "--strategy", - required=True, - choices=strategy_names(), - help="Quantization strategy (see options below).", - ) - parser.add_argument( - "--scheme", - required=True, - choices=scheme_names(), - help="Quantization scheme (see options below).", - ) - parser.add_argument( - "--model_path", - default=DEFAULT_MODEL_DIR, - help=f"Path to source HF model (default: {DEFAULT_MODEL_DIR}).", - ) - parser.add_argument( - "--output_path", - default=None, - help="Output dir for quantized model. Default: " - f"{DEFAULT_OUTPUT_BASE}/qwen2.5-vl-7b--/", - ) - parser.add_argument( - "--num_calib_samples", - type=int, - default=512, - help="Calibration sample count. Capped at 128 for multimodal " - "strategies (S3/S4) per NVIDIA's reference. Default: 512.", - ) - parser.add_argument( - "--calib", - default="auto", - choices=["auto", "text", "multimodal", "vln"], - help="Calibration data source. 'auto' (default) keeps the legacy " - "behaviour: text (cnn_dailymail) for LLM-only strategies, multimodal " - "(MMMU) when the visual tower is quantized. 'vln' uses the " - "in-distribution InternData-N1 VLN-CE set (recommended for this model).", - ) - parser.add_argument( - "--calib_data_root", - default=VLN_CALIB_DATA, - help=f"Root of the VLN-CE LeRobot data for --calib vln " - f"(default: {VLN_CALIB_DATA}).", - ) - parser.add_argument( - "--dtype", - default="bf16", - choices=["fp16", "bf16"], - help="Model load dtype. Default: bf16 (matches source config).", - ) - parser.add_argument( - "--device", - default="cuda", - help="Device for model + calibration. Default: cuda.", - ) - parser.add_argument( - "--resume", - default=None, - metavar="DIR", - help="Enable layerwise calibration resume from DIR (useful for " - "AWQ/Hessian recovery on crash).", - ) - parser.add_argument( - "--allow_experimental", - action="store_true", - help="Permit schemes marked experimental in configs/schemes.yaml. NVFP4 needs " - "this: it quantizes and generates fluent text, but the System 2 -> System 1 " - "bridge breaks (z_latents cosine 0.647 vs 0.9956 for FP8).", - ) - parser.add_argument( - "--dry_run", - action="store_true", - help="Load model + validate environment, skip quantization. " - "Useful for verifying setup before a long run.", - ) - return parser.parse_args() - - -def derive_output_path(args: argparse.Namespace) -> str: - if args.output_path is not None: - return args.output_path - tag = f"internvla-n1-system2-{args.strategy}-{args.scheme}" - return os.path.join(DEFAULT_OUTPUT_BASE, tag) - - -def print_environment() -> None: - """Quick environment summary for logging at the top of every run.""" - print("=" * 70) - print("Environment:") - print(f" torch: {torch.__version__}") - print(f" CUDA: {torch.cuda.is_available()}") - if torch.cuda.is_available(): - props = torch.cuda.get_device_properties(0) - cc = torch.cuda.get_device_capability(0) - print(f" device: {props.name}") - print(f" arch: sm_{cc[0]}{cc[1]}") - print(f" VRAM: {props.total_memory / 1e9:.1f} GB") - print(f" arch list: {torch.cuda.get_arch_list()}") - print("=" * 70) - - -# --------------------------------------------------------------------------- # -# Main -# --------------------------------------------------------------------------- # -def main() -> int: - args = parse_args() - output_path = derive_output_path(args) - - print_environment() - - registry = load_registry() - - # Reject impossible combinations before the model is loaded: a bad request should - # cost seconds, not a checkpoint load followed by a mid-quantization crash. - try: - validate(args.scheme, args.strategy, model_path=args.model_path, - allow_experimental=args.allow_experimental, registry=registry) - except ValueError as exc: - print(f"[ERROR] {exc}") - return 1 - - strat = registry["strategies"][args.strategy] - cfg_info = registry["schemes"][args.scheme] - - # Resolve calibration source. 'auto' preserves the legacy behaviour; an - # explicit --calib overrides it (e.g. VLN calib for an LLM-only strategy). - calib_mode = args.calib - if calib_mode == "auto": - calib_mode = "multimodal" if strat["quantize_visual"] else "text" - calib_desc = { - "text": "text-only (cnn_dailymail)", - "multimodal": "multimodal (MMMU image+text)", - "vln": "VLN in-distribution (InternData-N1 VLN-CE)", - }[calib_mode] - - print() - print(f"Strategy: {args.strategy} — {strat['description']}") - print(f"CFG preset: {args.scheme} — {cfg_info['description']}") - print(f"Model path: {args.model_path}") - print(f"Output path: {output_path}") - print(f"Calibration: {calib_desc}") - print() - - if not os.path.isdir(args.model_path): - print(f"ERROR: model_path not found: {args.model_path}", file=sys.stderr) - return 1 - - # -- Load model ------------------------------------------------------- # - t_load_start = time.time() - print("[1/4] Loading model...") - model, tokenizer, processor = load_model( - args.model_path, dtype=args.dtype, device=args.device - ) - original_max_length = tokenizer.model_max_length - print(f" done in {time.time() - t_load_start:.1f}s") - print(f" model class: {type(model).__name__}") - print(f" tokenizer.model_max_length (preserved): {original_max_length}") - - if args.dry_run: - print("[dry-run] Skipping quantization. Environment OK.") - return 0 - - # -- Build quant config ---------------------------------------------- # - quant_cfg = build_quant_config( - scheme=args.scheme, - strategy=args.strategy, - layerwise_checkpoint_dir=args.resume, - ) - - # -- Build dataloader & forward loop --------------------------------- # - t_calib_start = time.time() - print("[2/4] Preparing calibration data...") - - if calib_mode in ("multimodal", "vln"): - if processor is None: - raise RuntimeError( - f"{calib_mode} calibration requires an AutoProcessor but none " - "was found in the model dir." - ) - # Both image paths are GPU-memory bound (multi-image forward), so they - # share NVIDIA's multimodal sample cap. - mm_samples = min(args.num_calib_samples, MULTIMODAL_MAX_SAMPLES) - if mm_samples < args.num_calib_samples: - print(f" capping num_samples {args.num_calib_samples} → " - f"{mm_samples} (multimodal GPU-mem cap)") - if calib_mode == "vln": - batches = vln_calib_dataloader( - processor, - data_root=args.calib_data_root, - num_samples=mm_samples, - rgb_key=VLN_RGB_KEY, - ) - print(f" VLN batches prepared: {len(batches)} " - f"(data: {args.calib_data_root})") - else: - batches = multimodal_calib_dataloader( - processor, - dataset_name=MULTIMODAL_DATASET, - num_samples=mm_samples, - ) - print(f" multimodal batches prepared: {len(batches)}") - forward_loop = lambda m: calibrate_multimodal(m, batches) # noqa: E731 - else: - batch_size = calib_batch_size(args.scheme, is_image_calib=False) - loader = text_calib_dataloader( - tokenizer, - dataset_name=TEXT_DATASET, - batch_size=batch_size, - num_samples=args.num_calib_samples, - ) - forward_loop = lambda m: calibrate_text(m, loader) # noqa: E731 - print(f" text loader prepared: batch_size={batch_size}, " - f"samples={args.num_calib_samples}") - - # -- Quantize -------------------------------------------------------- # - print("[3/4] Running quantization...") - mtq.quantize(model, quant_cfg, forward_loop=forward_loop) - mtq.print_quant_summary(model) - print(f" quantization done in {time.time() - t_calib_start:.1f}s") - - # Safety: tokenizer.model_max_length must not have been touched by - # calibration. This is the user's main concern about preprocessor leakage. - assert tokenizer.model_max_length == original_max_length, ( - f"tokenizer.model_max_length changed during calibration: " - f"{original_max_length} → {tokenizer.model_max_length}" - ) - - # -- Export ---------------------------------------------------------- # - t_export_start = time.time() - print("[4/4] Exporting quantized checkpoint...") - export_quantized_model( - model=model, - tokenizer=tokenizer, - processor=processor, - model_dir=args.model_path, - output_dir=output_path, - ) - print(f" export done in {time.time() - t_export_start:.1f}s") - print() - print(f"Saved to {output_path}") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/quantize/repackage_system2.py b/recipes/internvla-n1-dualvln/quantize/repackage_system2.py deleted file mode 100644 index e7e9c5a..0000000 --- a/recipes/internvla-n1-dualvln/quantize/repackage_system2.py +++ /dev/null @@ -1,213 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Repackage the InternVLA-N1 System 2 into a standalone Qwen2.5-VL checkpoint. - -InternVLA-N1-DualVLN declares ``model_type: internvla_n1`` and ships no modeling code, so -it cannot be loaded with ``trust_remote_code`` -- the class has to come from the InternNav -repository. That dependency is avoidable for everything except System 1 itself, because the -System 2 backbone (vision tower + Qwen2.5 LLM) inside the checkpoint already uses standard -Qwen2.5-VL key names. - -This is a lossless subset copy, streamed through ``safe_open`` so peak memory is one shard: - - keep ``visual.*``, ``model.embed_tokens``, ``model.norm``, ``model.layers.*``, ``lm_head`` - drop the System 1 modules (traj_dit, rgb_model, rgb_resampler, memory_encoder, - cond_projector, action_encoder, action_decoder, pos_encoding) - rewrite ``config.json`` to ``model_type=qwen2_5_vl`` / - ``architectures=[Qwen2_5_VLForConditionalGeneration]`` - -Weights are copied bit-for-bit; the source is opened read-only. - -Two System 1 tensors are *not* simply dropped. ``latent_queries`` and ``cond_projector`` form -the System 2 -> System 1 bridge, and the fidelity checks need them to compute z_latents. They -are written to a separate ``bridge.safetensors`` (~25 MB) so that quantization, export, engine -build and verification can all run from this directory alone, without reopening the 16 GB -source checkpoint and without importing InternNav. -""" -import argparse -import json -import os -import shutil - -from safetensors import safe_open -from safetensors.torch import save_file - -SYSTEM1_PREFIXES = ( - "model.traj_dit", - "model.rgb_model", - "model.rgb_resampler", - "model.memory_encoder", - "model.cond_projector", - "model.action_encoder", - "model.action_decoder", - "model.pos_encoding", -) - -# Tensors that belong to System 1 but are needed to evaluate the System 2 -> System 1 -# bridge. Kept aside rather than dropped; see the module docstring. -BRIDGE_KEYS = ("model.latent_queries", "model.cond_projector") - -# Non-weight files carried over verbatim. Matching by prefix rather than an exact list -# keeps this working when a checkpoint ships an extra tokenizer artifact. -COPY_PREFIXES = ( - "tokenizer", - "vocab", - "merges", - "preprocessor", - "chat_template", - "generation_config", - "added_tokens", - "special_tokens", -) - -# A repackaged 7B System 2 is ~15.5 GB. Refuse to start rather than die at 90 %. -MIN_FREE_GB = 18 - - -def is_system2(key: str) -> bool: - """Keep the standard vision + LLM keys; drop every System 1 module.""" - if key.startswith(SYSTEM1_PREFIXES): - return False - return ( - key.startswith("visual.") - or key == "lm_head.weight" - or key.startswith("model.embed_tokens") - or key.startswith("model.norm") - or key.startswith("model.layers.") - ) - - -def is_bridge(key: str) -> bool: - """Tensors kept aside for z_latents evaluation.""" - return key.startswith(BRIDGE_KEYS) - - -def build_config(src: str) -> dict: - """Rewrite the InternVLA config into a plain Qwen2.5-VL one.""" - with open(os.path.join(src, "config.json")) as f: - cfg = json.load(f) - for key in ("system1", "n_query", "model_cfg", "model_type", "architectures", "auto_map"): - cfg.pop(key, None) - cfg["model_type"] = "qwen2_5_vl" - cfg["architectures"] = ["Qwen2_5_VLForConditionalGeneration"] - return cfg - - -def free_gb(path: str) -> float: - stat = os.statvfs(path) - return stat.f_bavail * stat.f_frsize / 1e9 - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument("--model_path", required=True, - help="Source InternVLA-N1-DualVLN checkpoint directory") - parser.add_argument("--output_path", required=True, - help="Destination directory for the stock Qwen2.5-VL System 2 checkpoint") - parser.add_argument("--free_source", action="store_true", - help="Delete each source shard once its converted copy is written. " - "Peak disk becomes one shard rather than two checkpoints. " - "Destructive -- only use this if the source can be re-downloaded.") - parser.add_argument("--skip_disk_check", action="store_true", - help="Skip the free-space preflight") - return parser.parse_args() - - -def main() -> int: - args = parse_args() - src, dst = args.model_path, args.output_path - - if not os.path.isdir(src): - print(f"[ERROR] --model_path does not exist: {src}") - return 1 - index_path = os.path.join(src, "model.safetensors.index.json") - if not os.path.isfile(index_path): - print(f"[ERROR] no model.safetensors.index.json under {src}; " - f"a sharded checkpoint is expected") - return 1 - - os.makedirs(dst, exist_ok=True) - - if not args.skip_disk_check and not args.free_source: - avail = free_gb(dst) - if avail < MIN_FREE_GB: - print(f"[ERROR] only {avail:.1f} GB free under {dst}; need >= {MIN_FREE_GB} GB. " - f"Free space, point --output_path elsewhere, or pass --free_source to " - f"delete each source shard as it is converted.") - return 1 - - with open(index_path) as f: - weight_map = json.load(f)["weight_map"] - - kept = [k for k in weight_map if is_system2(k)] - bridge = [k for k in weight_map if is_bridge(k)] - dropped = [k for k in weight_map if not is_system2(k)] - print(f"keep {len(kept)} System2 keys, drop {len(dropped)} System1 keys " - f"({len(bridge)} of them kept aside as bridge tensors)") - - # Group by source shard so each is opened once and read sequentially. - by_shard: dict[str, list[str]] = {} - for key in kept: - by_shard.setdefault(weight_map[key], []).append(key) - bridge_by_shard: dict[str, list[str]] = {} - for key in bridge: - bridge_by_shard.setdefault(weight_map[key], []).append(key) - - new_weight_map: dict[str, str] = {} - bridge_tensors: dict[str, "object"] = {} - total_bytes = 0 - out_shards = sorted(by_shard) - n_shards = len(out_shards) - - for i, shard in enumerate(out_shards, 1): - out_name = f"model-{i:05d}-of-{n_shards:05d}.safetensors" - tensors = {} - with safe_open(os.path.join(src, shard), framework="pt") as f: - for key in by_shard[shard]: - tensor = f.get_tensor(key) - tensors[key] = tensor - total_bytes += tensor.numel() * tensor.element_size() - new_weight_map[key] = out_name - for key in bridge_by_shard.get(shard, []): - bridge_tensors[key] = f.get_tensor(key) - save_file(tensors, os.path.join(dst, out_name), metadata={"format": "pt"}) - print(f" [{i}/{n_shards}] {out_name}: {len(tensors)} tensors") - del tensors - if args.free_source: - os.remove(os.path.join(src, shard)) - print(f" removed source shard {shard}") - - with open(os.path.join(dst, "model.safetensors.index.json"), "w") as f: - json.dump({"metadata": {"total_size": total_bytes}, "weight_map": new_weight_map}, - f, indent=2) - - with open(os.path.join(dst, "config.json"), "w") as f: - json.dump(build_config(src), f, indent=2) - - if bridge_tensors: - save_file(bridge_tensors, os.path.join(dst, "bridge.safetensors"), - metadata={"format": "pt"}) - print(f" bridge.safetensors: {len(bridge_tensors)} tensors " - f"({', '.join(sorted(bridge_tensors))})") - else: - print("[WARN] no bridge tensors found; z_latents verification will need the " - "original checkpoint") - - for name in os.listdir(src): - if name.endswith(".safetensors"): - continue - if any(name.startswith(p) for p in COPY_PREFIXES): - shutil.copy2(os.path.join(src, name), os.path.join(dst, name)) - - dropped_modules = sorted({k.split(".")[1] for k in dropped if "." in k}) - print(f"\nDone -> {dst}") - print(f" {n_shards} shards, {total_bytes / 1e9:.1f} GB System 2") - print(" config model_type=qwen2_5_vl") - print(f" dropped System 1 modules: {dropped_modules}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/quantize/scripts/00_fetch_calib_scenes.sh b/recipes/internvla-n1-dualvln/quantize/scripts/00_fetch_calib_scenes.sh deleted file mode 100755 index 566f898..0000000 --- a/recipes/internvla-n1-dualvln/quantize/scripts/00_fetch_calib_scenes.sh +++ /dev/null @@ -1,84 +0,0 @@ -#!/usr/bin/env bash -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -# -# Fetch a small, diverse subset of InternData-N1 VLN-CE scenes for calibration. -# -# InternData-N1 is gated on HuggingFace and the full VLN-CE traj_data is ~2.5 TB across 914 -# scenes, so this pulls only a handful of per-scene archives. Accept the dataset terms and -# run `huggingface-cli login` first, or every download 401s. -# -# The default set spans r2r + rxr + scalevln -- the same training mix the base model saw -- -# and is deliberately made of small scenes to keep disk in check. Override SCENES to pick -# others. -# -# Note one of these scenes, YmJkqBEsHnH, also appears in the held-out probe set used by -# quantize/benchmark_accuracy.py. Exclude it from either side before reading a number as -# out-of-sample; qat.py already does. - -set -euo pipefail - -OUTPUT_PATH="${CALIB_DATA_ROOT:-$HOME/vln-opt-work/calib_scenes}" -SCENES="${SCENES:-\ -vln_ce/traj_data/r2r/gZ6f7yhEvPG.tar.gz \ -vln_ce/traj_data/r2r/YmJkqBEsHnH.tar.gz \ -vln_ce/traj_data/r2r/XcA2TqTSSAj.tar.gz \ -vln_ce/traj_data/rxr/Pm6F8kyY3z2.tar.gz \ -vln_ce/traj_data/rxr/PuKPg4mmafe.tar.gz \ -vln_ce/traj_data/scalevln/00493-pUneSGJDrvY.tar.gz \ -vln_ce/traj_data/scalevln/00446-tL6i2PtktSh.tar.gz \ -vln_ce/traj_data/scalevln/00351-QxfX5te1gFu.tar.gz \ -vln_ce/traj_data/scalevln/00335-janiYDpzM9j.tar.gz}" - -while [[ $# -gt 0 ]]; do - case "$1" in - --output_path) OUTPUT_PATH="$2"; shift 2 ;; - --scenes) SCENES="$2"; shift 2 ;; - -h|--help) sed -n '2,20p' "$0"; exit 0 ;; - *) echo "[ERROR] unknown argument: $1" >&2; exit 2 ;; - esac -done - -mkdir -p "$OUTPUT_PATH" -export INTERNDATA_DEST="$OUTPUT_PATH" INTERNDATA_SCENES="$SCENES" - -python - <<'PY' -import os -import shutil -import tarfile -import tempfile - -from huggingface_hub import hf_hub_download - -REPO = "InternRobotics/InternData-N1" -dest = os.environ["INTERNDATA_DEST"] -scenes = os.environ["INTERNDATA_SCENES"].split() - -# Download into a scratch dir and delete each archive right after extraction, so the HF blob -# cache never accumulates a second copy. local_dir avoids the shared cache entirely. -scratch = tempfile.mkdtemp(prefix="interndata_") -try: - for rel in scenes: - scene = os.path.basename(rel)[:-len(".tar.gz")] - subset = rel.split("/")[2] # r2r | rxr | scalevln - out_dir = os.path.join(dest, subset, scene) - if os.path.isdir(os.path.join(out_dir, "meta")): - print(f"[skip] already extracted: {out_dir}") - continue - print(f"[get ] {rel}") - tar_path = hf_hub_download(REPO, rel, repo_type="dataset", - local_dir=scratch, local_dir_use_symlinks=False) - with tarfile.open(tar_path) as tf: - tf.extractall(os.path.join(dest, subset)) - os.remove(tar_path) - print(f"[ok ] {scene}") -finally: - shutil.rmtree(scratch, ignore_errors=True) - -found = sum(1 for root, _, files in os.walk(dest) - if root.endswith(os.sep + "meta") and "episodes.jsonl" in files) -print(f"[done] {found} scene(s) with episodes.jsonl under {dest}") -PY - -echo "Calibration data ready: $OUTPUT_PATH" -echo "Pass it as --calib_data_root, or set CALIB_DATA_ROOT." diff --git a/recipes/internvla-n1-dualvln/quantize/scripts/01_repackage.sh b/recipes/internvla-n1-dualvln/quantize/scripts/01_repackage.sh deleted file mode 100755 index 38499bd..0000000 --- a/recipes/internvla-n1-dualvln/quantize/scripts/01_repackage.sh +++ /dev/null @@ -1,38 +0,0 @@ -#!/usr/bin/env bash -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -# -# Strip System 1 out of the InternVLA checkpoint, leaving a stock Qwen2.5-VL System 2. -# -# This is the step that makes the rest of the recipe ordinary: after it, quantize, export, -# build and the bridge verification all operate on a plain Qwen2.5-VL checkpoint and never -# import InternNav. It is pure safetensors manipulation -- no model is constructed -- so it -# runs anywhere, including without a GPU. -# -# Costs one ~15 GB intermediate copy. Pass --free_source to delete each source shard as it -# is consumed if disk is tight; the peak then is one shard rather than two checkpoints. - -set -euo pipefail - -HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -MODEL_PATH="${INTERNVLA_CKPT:-$HOME/InternNav/checkpoints/InternVLA-N1-DualVLN}" -OUTPUT_PATH="${REPKG_CKPT:-$HOME/vln-opt-work/qwen25vl_system2}" -EXTRA=() - -while [[ $# -gt 0 ]]; do - case "$1" in - --model_path) MODEL_PATH="$2"; shift 2 ;; - --output_path) OUTPUT_PATH="$2"; shift 2 ;; - --free_source) EXTRA+=(--free_source); shift ;; - --skip_disk_check) EXTRA+=(--skip_disk_check); shift ;; - -h|--help) sed -n '2,14p' "$0"; exit 0 ;; - *) echo "[ERROR] unknown argument: $1" >&2; exit 2 ;; - esac -done - -[[ -d "$MODEL_PATH" ]] || { echo "[ERROR] checkpoint not found: $MODEL_PATH" >&2; exit 1; } - -exec python -u "$HERE/../repackage_system2.py" \ - --model_path "$MODEL_PATH" \ - --output_path "$OUTPUT_PATH" \ - "${EXTRA[@]}" diff --git a/recipes/internvla-n1-dualvln/quantize/scripts/02_quantize.sh b/recipes/internvla-n1-dualvln/quantize/scripts/02_quantize.sh deleted file mode 100755 index b6eb2bd..0000000 --- a/recipes/internvla-n1-dualvln/quantize/scripts/02_quantize.sh +++ /dev/null @@ -1,64 +0,0 @@ -#!/usr/bin/env bash -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -# -# Quantize the repackaged System 2 with one scheme x strategy combination. -# -# Input is the *repackaged* checkpoint, not the InternVLA one -- see 01_repackage.sh. -# The scheme/strategy validity gate lives in quantize/quant_schemes.py and rejects the -# impossible combinations early, with the reason; nvfp4 x {s1,s2} additionally needs -# --allow_experimental because it passes text fluency and still breaks the navigation -# bridge (z_latents 0.931 against a 0.99 gate). -# -# Calibration defaults to 'auto' -- cnn_dailymail text for LLM-only strategies -- matching -# quantize.py. Pass --calib vln to use navigation episodes instead; it needs the scenes from -# 00_fetch_calib_scenes.sh and, measured here, changes held-out z_latents by 0.00003, so -# prefer it for honesty rather than for accuracy. -# -# Pass --dry_run to load the model and validate the configuration without quantizing. - -set -euo pipefail - -HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -MODEL_PATH="${REPKG_CKPT:-$HOME/vln-opt-work/qwen25vl_system2}" -OUTPUT_PATH="" -SCHEME="${SCHEME:-fp8_default}" -STRATEGY="${STRATEGY:-s1}" -DEVICE="${DEVICE:-cuda}" -CALIB="${CALIB:-auto}" -CALIB_DATA_ROOT="${CALIB_DATA_ROOT:-$HOME/vln-opt-work/calib_scenes}" -EXTRA=() - -while [[ $# -gt 0 ]]; do - case "$1" in - --model_path) MODEL_PATH="$2"; shift 2 ;; - --output_path) OUTPUT_PATH="$2"; shift 2 ;; - --scheme) SCHEME="$2"; shift 2 ;; - --strategy) STRATEGY="$2"; shift 2 ;; - --device) DEVICE="$2"; shift 2 ;; - --calib) CALIB="$2"; shift 2 ;; - --calib_data_root) CALIB_DATA_ROOT="$2"; shift 2 ;; - --num_calib_samples) EXTRA+=(--num_calib_samples "$2"); shift 2 ;; - --allow_experimental) EXTRA+=(--allow_experimental); shift ;; - --dry_run) EXTRA+=(--dry_run); shift ;; - -h|--help) sed -n "2,19p" "$0"; exit 0 ;; - *) echo "[ERROR] unknown argument: $1" >&2; exit 2 ;; - esac -done - -OUTPUT_PATH="${OUTPUT_PATH:-$HOME/vln-opt-work/qwen25vl_${STRATEGY}_${SCHEME}}" -[[ -d "$MODEL_PATH" ]] || { echo "[ERROR] checkpoint not found: $MODEL_PATH" >&2; exit 1; } - -# ModelOpt's Triton kernels are compiled in-tree on Thor; without this the quantize pass -# fails at import time rather than at use. -export TRITON_BACKENDS_IN_TREE="${TRITON_BACKENDS_IN_TREE:-1}" - -exec python -u "$HERE/../quantize.py" \ - --model_path "$MODEL_PATH" \ - --output_path "$OUTPUT_PATH" \ - --scheme "$SCHEME" \ - --strategy "$STRATEGY" \ - --device "$DEVICE" \ - --calib "$CALIB" \ - --calib_data_root "$CALIB_DATA_ROOT" \ - "${EXTRA[@]}" diff --git a/recipes/internvla-n1-dualvln/requirements-torch.txt b/recipes/internvla-n1-dualvln/requirements-torch.txt deleted file mode 100644 index aff3bb1..0000000 --- a/recipes/internvla-n1-dualvln/requirements-torch.txt +++ /dev/null @@ -1,31 +0,0 @@ -# Install PyTorch BEFORE running: pip install -r requirements.txt -# Pick the line matching your platform: -# -# Jetson Thor (JetPack 7.1, CUDA 13.0) — this is the tested configuration: -# Use the JetPack-provided wheel. Do NOT install torch from PyPI on Jetson; -# the PyPI aarch64 wheels are not built for sm_110 and will fail at runtime. -# Verify with: -# python -c "import torch; print(torch.__version__, torch.cuda.get_arch_list())" -# Expect torch 2.10.0 and 'sm_110' present in the arch list. -# -# x86 CUDA 12.8: -# pip install torch==2.10.0+cu128 torchvision==0.25.0+cu128 --extra-index-url https://download.pytorch.org/whl/cu128 -# -# x86 CUDA 13.x: -# pip install torch==2.10.0+cu130 torchvision==0.25.0+cu130 --extra-index-url https://download.pytorch.org/whl/cu130 -# -# ----------------------------------------------------------------------------- -# numpy must be pinned AFTER `uv sync`, not by it. -# -# This recipe needs numpy 1.x: OpenCV and diffusers on Jetson break under numpy 2, -# while three other extras in this repository pin numpy==2.2.6. Rather than declare a -# pin whose resolution outcome we have not verified, numpy and scipy are left out of -# this recipe's pyproject extra and installed explicitly below. -# -# Run this once after `uv sync --extra internvla-n1-dualvln`: -# -# pip install "numpy==1.26.4" "scipy==1.13.1" -# -# Installing `datasets` later can silently upgrade numpy again; re-run the line above -# if it does. -# ----------------------------------------------------------------------------- diff --git a/recipes/internvla-n1-dualvln/run_matrix.py b/recipes/internvla-n1-dualvln/run_matrix.py deleted file mode 100644 index baa009b..0000000 --- a/recipes/internvla-n1-dualvln/run_matrix.py +++ /dev/null @@ -1,124 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Collect every measurement for every built variant into one comparison table. - -Reads what is already on disk rather than recomputing: checkpoint sizes, engine sizes, -the accuracy JSON from benchmark_accuracy.py, the NVFP4 investigation JSON, and the -z_latents numbers passed in. Anything missing prints as "-" instead of failing, so the -table can be produced at any point in the pipeline. - -Emits Markdown, ready to paste into the recipe README. -""" -import argparse -import json -import os - - -def dir_size_gb(path: str) -> float | None: - if not path or not os.path.isdir(path): - return None - total = 0 - for root, _dirs, files in os.walk(path): - for name in files: - try: - total += os.path.getsize(os.path.join(root, name)) - except OSError: - pass - return total / 1e9 - - -def file_size_gb(path: str) -> float | None: - return os.path.getsize(path) / 1e9 if path and os.path.isfile(path) else None - - -def fmt(value, spec: str = ".2f", suffix: str = "") -> str: - return "-" if value is None else f"{value:{spec}}{suffix}" - - -def load_json(path: str) -> dict | list | None: - if path and os.path.isfile(path): - with open(path) as f: - return json.load(f) - return None - - -def main() -> int: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--work_dir", default=os.path.expanduser("~/vln-opt-work")) - p.add_argument("--output_path", default=None, help="Write the Markdown here") - args = p.parse_args() - w = args.work_dir - - # variant -> (checkpoint dir, engine dir, label) - variants = [ - ("BF16 (unquantized)", "qwen25vl_system2", "engines/base_fp16"), - ("FP8 s1", "qwen25vl_s1_fp8", "engines/s1_fp8"), - ("NVFP4 s1 (experimental)", "qwen25vl_s1_nvfp4", "engines/s1_nvfp4"), - ] - - acc = load_json(os.path.join(w, "out/accuracy_bf16_vs_fp8.json")) or [] - acc_by_path = {os.path.basename(r["model_path"]): r for r in acc} - inv = load_json(os.path.join(w, "nvfp4_investigation.json")) or {} - lat = load_json(os.path.join(w, "out/latency.json")) or {} - zlat = load_json(os.path.join(w, "out/z_latents.json")) or {} - - rows = [] - for label, ckpt_name, eng_name in variants: - ckpt = os.path.join(w, ckpt_name) - eng_llm = os.path.join(w, eng_name, "llm/llm.engine") - eng_vis = os.path.join(w, eng_name, "visual/visual.engine") - a = acc_by_path.get(ckpt_name) - rows.append({ - "label": label, - "ckpt_gb": dir_size_gb(ckpt), - "llm_gb": file_size_gb(eng_llm), - "vis_gb": file_size_gb(eng_vis), - "prefill_ms": lat.get(ckpt_name, {}).get("prefill_ms"), - "decode_ms": lat.get(ckpt_name, {}).get("decode_ms"), - "z_engine": zlat.get(ckpt_name), - "z_weights": inv.get("bridge", {}).get( - {"qwen25vl_s1_fp8": "fp8", "qwen25vl_s1_nvfp4": "nvfp4"}.get(ckpt_name, "")), - "l2_mean": a and a.get("pixel_goal_l2_mean"), - "l2_median": a and a.get("pixel_goal_l2_median"), - "w_err": inv.get("weights", {}).get( - {"qwen25vl_s1_fp8": "fp8", "qwen25vl_s1_nvfp4": "nvfp4"}.get(ckpt_name, ""), - {}).get("rel_err_mean"), - }) - - out = [] - out.append("### Benchmark matrix\n") - out.append("Every column is a real measurement on this machine; `-` means not measured.\n") - out.append("| Variant | checkpoint | LLM engine | visual | prefill | decode " - "| z_latents (engine) | z_latents (weights) | pixel L2 mean / median |") - out.append("|---|---|---|---|---|---|---|---|---|") - for r in rows: - l2 = ("-" if r["l2_mean"] is None - else f"{r['l2_mean']:.2f} / {r['l2_median']:.2f} px") - out.append( - f"| {r['label']} | {fmt(r['ckpt_gb'], '.1f', ' GB')} | " - f"{fmt(r['llm_gb'], '.2f', ' GB')} | {fmt(r['vis_gb'], '.2f', ' GB')} | " - f"{fmt(r['prefill_ms'], '.1f', ' ms')} | {fmt(r['decode_ms'], '.1f', ' ms')} | " - f"{fmt(r['z_engine'], '.6f')} | {fmt(r['z_weights'], '.6f')} | {l2} |") - - out.append("\n**Weight quantization error** (mean relative, 21 projections across " - "layers 0/13/27):\n") - out.append("| Variant | rel-err |") - out.append("|---|---|") - for r in rows: - if r["w_err"] is not None: - out.append(f"| {r['label']} | {100 * r['w_err']:.2f} % |") - - text = "\n".join(out) - print(text) - if args.output_path: - os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True) - with open(args.output_path, "w") as f: - f.write(text + "\n") - print(f"\nWrote {args.output_path}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_memory.py b/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_memory.py deleted file mode 100644 index 7ce5318..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_memory.py +++ /dev/null @@ -1,71 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Measure peak GPU memory (torch.cuda.max_memory_allocated) of the full PyTorch -pipeline (System 2 weights + System 1 forward). -""" -import os -import sys -_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, _R) -sys.path.insert(0, os.path.join(_R, "lib")) -sys.path.append("/usr/lib/python3.12/dist-packages") -import numpy as np # noqa: E402 -import torch # noqa: E402 -from PIL import Image # noqa: E402 - -ACTIVE = os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) -CKPT = os.path.join(ACTIVE, "checkpoints/InternVLA-N1-DualVLN") -IMG = os.path.expanduser("~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg") -GB = 1024**3 - - -def peak(): return torch.cuda.max_memory_allocated() / GB # noqa: E704 - - -def reset(): - torch.cuda.reset_peak_memory_stats() - torch.cuda.empty_cache() - - -def main(): - dev = "cuda" - try: - torch.backends.mha.set_fastpath_enabled(False) - except Exception: - pass - if ACTIVE not in sys.path: - sys.path.insert(0, ACTIVE) - from internvla_compat import apply_all - apply_all(need_system1=True, allow_missing_depth=True) - from internnav.model.basemodel.internvla_n1.internvla_n1 import ( - InternVLAN1ForCausalLM, InternVLAN1ModelConfig) - reset() - print("[1] Load full InternVLA (PyTorch, System1+System2)") - cfg = InternVLAN1ModelConfig.from_pretrained(CKPT) - model = InternVLAN1ForCausalLM.from_pretrained(CKPT, config=cfg, torch_dtype=torch.bfloat16, - attn_implementation="sdpa", low_cpu_mem_usage=True).to(dev).eval() - w_mem = torch.cuda.memory_allocated() / GB - print(f" weights loaded (S2+S1): {w_mem:.2f} GB") - - NQ = model.get_n_query() - z = torch.randn(1, NQ, cfg.hidden_size, device=dev, dtype=torch.bfloat16) - a = np.array(Image.open(IMG).convert("RGB").resize((224, 224))) / 255.0 - tt = torch.from_numpy(a).float() - images_dp = torch.stack([tt, tt]).unsqueeze(0).to(dev) - - print("[2] S1 generate_traj peak (PyTorch)") - reset() - with torch.no_grad(): - model.generate_traj(z, images_dp, num_sample_trajs=32, num_inference_steps=10) - torch.cuda.synchronize() - print(f" S1 peak (incl. weights): {peak():.2f} GB") - - print("\n=== PyTorch peak GPU mem ===") - print(f" weights S2+S1 : {w_mem:.2f} GB") - print(f" peak during S1: {peak():.2f} GB") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_system1.py b/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_system1.py deleted file mode 100644 index 8df6fbf..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/bench_system1.py +++ /dev/null @@ -1,77 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Benchmark System 1 (generate_traj) latency in PyTorch and report Hz. -""" -import os -import sys -import time -_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, _R) -sys.path.insert(0, os.path.join(_R, "lib")) -sys.path.append("/usr/lib/python3.12/dist-packages") # cv2 (system) for depth_anything -import numpy as np # noqa: E402 -import torch # noqa: E402 -from PIL import Image # noqa: E402 - -ACTIVE = os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) -CKPT = os.path.join(ACTIVE, "checkpoints/InternVLA-N1-DualVLN") -REPKG = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "qwen25vl_system2") -IMG = os.path.expanduser("~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg") - - -def timeit(fn, warm=2, n=5): - for _ in range(warm): - fn() - torch.cuda.synchronize() - ts = [] - for _ in range(n): - torch.cuda.synchronize() - t0 = time.perf_counter() - fn() - torch.cuda.synchronize() - ts.append(time.perf_counter() - t0) - return sum(ts) / len(ts), min(ts), max(ts) - - -def main(): - dev = "cuda" - if ACTIVE not in sys.path: - sys.path.insert(0, ACTIVE) - from internvla_compat import apply_all - apply_all(need_system1=True, allow_missing_depth=True) - from internnav.model.basemodel.internvla_n1.internvla_n1 import ( - InternVLAN1ForCausalLM, InternVLAN1ModelConfig) - print("[1/3] Load full InternVLA (System1)") - cfg = InternVLAN1ModelConfig.from_pretrained(CKPT) - model = InternVLAN1ForCausalLM.from_pretrained( - CKPT, config=cfg, torch_dtype=torch.bfloat16, attn_implementation=os.environ.get("ATTN", "sdpa"), - low_cpu_mem_usage=True).to(dev).eval() - - print("[2/3] Prepare z_latents + images_dp (as the agent)") - N_QUERY = model.get_n_query() - # placeholder z_latents (cond_projector is applied inside generate_traj) - traj_latents = torch.randn(1, N_QUERY, cfg.hidden_size, device=dev, dtype=torch.bfloat16) - a = np.array(Image.open(IMG).convert("RGB").resize((224, 224))) / 255.0 - t = torch.from_numpy(a).float() - images_dp = torch.stack([t, t]).unsqueeze(0).to(dev) # [1,2,224,224,3] - - print("[3/3] Time generate_traj (System 1)\n" + "=" * 60) - with torch.no_grad(): - for ns in (32, 4, 1): - def fn(ns=ns): - with torch.no_grad(): - model.generate_traj(traj_latents.to(dev), images_dp, - num_sample_trajs=ns, num_inference_steps=10) - try: - mean, mn, mx = timeit(fn, warm=2, n=5) - print( - f" num_sample_trajs={ns:>2}: {mean*1000:7.1f} ms (min {mn*1000:.0f}, max {mx*1000:.0f}) → {1/mean:5.1f} Hz") # noqa: E501 - except Exception as e: - print(f" num_sample_trajs={ns}: ERR {type(e).__name__}: {e}") - print("\n (paper target S1 = 30 Hz ≈ 33 ms; S2 = 2 Hz ≈ 500 ms)") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/benchmark_system2.py b/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/benchmark_system2.py deleted file mode 100644 index 27afac0..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/benchmark/benchmark_system2.py +++ /dev/null @@ -1,159 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Benchmark System 2 latency: PyTorch vs FP8 TensorRT, same navigation input, -engine load time excluded. Both paths produce identical output. -""" -import os -import sys -import json -import time -import subprocess -_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, _R) -sys.path.insert(0, os.path.join(_R, "lib")) -import torch # noqa: E402 -from engine_runner import ENGINE, REPKG # noqa: E402 - -TRT = os.environ.get("TRT_EDGE_LLM", os.path.expanduser("~/TensorRT-Edge-LLM")) -ENG_LLM = os.path.dirname(ENGINE) -ENG_VIS = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "engines/system2_visual") -# NOTE: nothing in this recipe (or the source project) generates this manifest. It is a -# user-supplied list of golden samples. The preflight below says so rather than letting -# the benchmark die on a missing file. -# prompt_builder is the single source of truth for the VLN prompt and lives in the -# quantize path, so calibration and verification cannot drift apart. Walk up to the -# recipe root instead of counting parent levels -- these scripts sit at two depths. -_d = os.path.dirname(os.path.abspath(__file__)) -while _d != "/" and not os.path.isdir(os.path.join(_d, "quantize")): - _d = os.path.dirname(_d) -sys.path.insert(0, os.path.join(_d, "quantize")) - -MANIFEST = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "golden/manifest_v3.json") -SCRATCH = os.path.expanduser(os.environ.get('VLN_OPT_OUT', '~/vln-opt-work/out')) -IMG_DIR = os.path.join(SCRATCH, "bench_imgs") -os.makedirs(IMG_DIR, exist_ok=True) -MINP, MAXP = 3136, 12845056 -MAXNEW = 16 - - -def fill_paths(conv, image_paths): - out, j = [], 0 - for turn in conv: - content = [] - for it in turn["content"]: - if it["type"] == "image": - content.append({"type": "image", "image": os.path.abspath(image_paths[j])}) - j += 1 - else: - content.append({"type": "text", "text": it["text"]}) - out.append({"role": turn["role"], "content": content}) - return out - - -def run_llm_inference(requests): - inp = os.path.join(SCRATCH, "bench_in.json") - out = os.path.join(SCRATCH, "bench_out.json") - json.dump({"batch_size": 1, "temperature": 0.0, "top_p": 1.0, "top_k": 1, "max_generate_length": MAXNEW, - "requests": requests}, open(inp, "w")) - env = dict(os.environ, EDGELLM_PLUGIN_PATH=f"{TRT}/build/libNvInfer_edgellm_plugin.so") - t0 = time.perf_counter() - r = subprocess.run([f"{TRT}/build/examples/llm/llm_inference", "--engineDir", ENG_LLM, - "--multimodalEngineDir", ENG_VIS, "--inputFile", inp, "--outputFile", out], - cwd=TRT, env=env, capture_output=True, text=True) - dt = time.perf_counter() - t0 - resp = json.load(open(out)).get("responses", []) if os.path.exists(out) else [] - return dt, [x.get("output_text", "").strip() for x in resp], r.returncode - - -def _require_manifest(): - """Fail with an explanation instead of a bare FileNotFoundError. - - The golden manifest is a user-supplied list of samples; no script here produces one. - """ - if not os.path.isfile(MANIFEST): - raise SystemExit( - f"[ERROR] golden manifest not found: {MANIFEST}\n" - f" MANIFEST is required and is not generated by this recipe. It should\n" - f" be a JSON file with a 'samples' list, each entry carrying episode,\n" - f" episode_idx, instruction and images -- the same shape\n" - f" quantize/benchmark_accuracy.py builds from LeRobot episodes.\n" - f" Override the location with WORK_DIR or create the file.") - - -def main(): - _require_manifest() - dev = "cuda" - from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration - from prompt_builder import build_sample_inputs, build_conversation, build_conversation_lookdown - man = json.load(open(MANIFEST)) - t2 = next(s for s in man["samples"] if s.get("turn") == 2) - t1 = next(s for s in man["samples"] if s.get("turn", 1) == 1) - cases = [("turn2/coord", t2), ("turn1/action", t1)] - - print("[1/3] Load model + processor (PyTorch repo path)") - model = Qwen2_5_VLForConditionalGeneration.from_pretrained( - REPKG, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2", - low_cpu_mem_usage=True).to(dev).eval() - proc = AutoProcessor.from_pretrained(REPKG, trust_remote_code=True, min_pixels=MINP, max_pixels=MAXP) - - results = {} - reqs = {} - print("[2/3] PyTorch generate (2 warm-up, 5 timed runs)") - for name, s in cases: - inp = build_sample_inputs(s, proc).to(dev) - S = inp["input_ids"].shape[1] - for _ in range(2): # warm - with torch.no_grad(): - model.generate(**inp, max_new_tokens=MAXNEW, do_sample=False, use_cache=True) - torch.cuda.synchronize() - ts = [] - for _ in range(5): - torch.cuda.synchronize() - t0 = time.perf_counter() - with torch.no_grad(): - out = model.generate(**inp, max_new_tokens=MAXNEW, do_sample=False, use_cache=True) - torch.cuda.synchronize() - ts.append(time.perf_counter() - t0) - ntok = out.shape[1] - S - txt = proc.tokenizer.decode(out[0, S:], skip_special_tokens=True).strip() - results[name] = {"S": S, "ntok": int(ntok), "torch_s": sum(ts) / len(ts), "torch_out": txt} - print(f" {name}: S={S} gen={ntok}tok torch={results[name]['torch_s']*1000:.0f}ms out={txt!r}") - # build llm_inference request (same sample) - turn = s.get("turn", 1) - conv = (build_conversation_lookdown(s["episode_idx"], s["instruction"], s["assistant_turn1"]) - if turn == 2 else build_conversation(s["episode_idx"], s["instruction"])[0]) - conv = conv if isinstance(conv, list) else [conv] - reqs[name] = {"messages": fill_paths(conv, s["images"])} - del model - torch.cuda.empty_cache() - - print("[3/3] FP8 TensorRT llm_inference (load excluded: t_1 vs t_N)") - N = 11 - for name, _ in cases: - t1s, o1, rc1 = run_llm_inference([reqs[name]]) - tNs, oN, rcN = run_llm_inference([reqs[name]] * N) - per = (tNs - t1s) / (N - 1) - results[name]["fp8_s"] = per - results[name]["fp8_out"] = (oN[0] if oN else "") - results[name]["load_s"] = t1s - per - print( - f" {name}: t_1={t1s:.2f}s t_{N}={tNs:.2f}s → per-req={per*1000:.0f}ms (load≈{results[name]['load_s']:.1f}s) out={results[name]['fp8_out']!r}") # noqa: E501 - - print("\n" + "=" * 70) - print(f"{'case':<14}{'S':>6}{'gen':>5}{'PyTorch':>12}{'FP8 TRT':>12}{'speedup':>10}") - print("-" * 70) - for name, _ in cases: - r = results[name] - sp = r["torch_s"] / r["fp8_s"] - print(f"{name:<14}{r['S']:>6}{r['ntok']:>5}{r['torch_s']*1000:>10.0f}ms{r['fp8_s']*1000:>10.0f}ms{sp:>9.2f}x") - report_dir = os.path.join(os.environ.get("WORK_DIR", - os.path.expanduser("~/vln-opt-work")), "reports") - os.makedirs(report_dir, exist_ok=True) # the final line used to crash without this - json.dump(results, open(os.path.join(report_dir, "bench_e2e.json"), "w"), - indent=2, ensure_ascii=False) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/deploy/run_eval_engine.py b/recipes/internvla-n1-dualvln/trt-edgellm/deploy/run_eval_engine.py deleted file mode 100644 index 06aa7f5..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/deploy/run_eval_engine.py +++ /dev/null @@ -1,74 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Run InternNav's closed-loop eval (SR/SPL) with the TensorRT engine instead of PyTorch. - -This is a thin launcher for the VLN team's SIMULATOR machine. It monkeypatches InternNav's policy -registry so `policy_name='internvla_n1'` resolves to the engine-backed policy -(lib/engine_policy.EngineInternVLAN1Net), then hands off to the standard eval entry -`scripts/eval/eval.py`. No InternNav source edit needed. - -Prerequisites (VLN team side): - * InternNav with the simulator installed (Habitat or InternUtopia/Isaac Sim) + the eval episodes. - * TensorRT-Edge-LLM built (llm_inference + plugin) and the engines from this repo on the machine. - -Environment: - export INTERNNAV_ROOT=/path/to/InternNav - export TRT_EDGE_LLM=/path/to/TensorRT-Edge-LLM - export VLN_LLM_ENGINE_DIR=/path/to/engines/system2_llm_fp8 # or system2_llm_base_fp16 - export VLN_VIS_ENGINE_DIR=/path/to/engines/system2_visual - -Run: - python deploy/run_eval_engine.py --config scripts/eval/configs/h1_internvla_n1_async_cfg.py - -The eval computes SR / SPL / NE exactly as the PyTorch run — only the System 2 LLM text generation -is served by the engine (the FP8-quantized decision path). Point `model_path` in the config at your -(retrained) checkpoint; the engines must have been built from that same checkpoint. -""" -import os -import sys - -_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, os.path.join(_REPO, "lib")) - -INTERNNAV_ROOT = os.path.expanduser(os.environ.get("INTERNNAV_PATH", "~/InternNav")) -sys.path.insert(0, INTERNNAV_ROOT) - - -def _install_engine_policy(): - """Route policy_name='internvla_n1' to the engine-backed subclass, in every namespace that - already imported get_policy (the registry module AND the agent that did `from ... import`).""" - import internnav.model as model_mod - from engine_policy import EngineInternVLAN1Net - - orig = model_mod.get_policy - - def get_policy(policy_name): - if policy_name == "internvla_n1": - print("[run_eval_engine] policy 'internvla_n1' -> EngineInternVLAN1Net (TRT engine)") - return EngineInternVLAN1Net - return orig(policy_name) - - model_mod.get_policy = get_policy - # the agent module did `from internnav.model import get_policy` — patch its bound name too - import internnav.agent.internvla_n1_agent as agent_mod - if hasattr(agent_mod, "get_policy"): - agent_mod.get_policy = get_policy - - -def main(): - _install_engine_policy() - # Hand off to the real eval entry, preserving argv (--config ...). - eval_py = os.path.join(INTERNNAV_ROOT, "scripts", "eval", "eval.py") - if not os.path.isfile(eval_py): - sys.exit(f"eval.py not found at {eval_py}; set INTERNNAV_ROOT correctly.") - g = {"__name__": "__main__", "__file__": eval_py} - with open(eval_py) as f: - code = compile(f.read(), eval_py, "exec") - # eval.py uses relative paths like './third_party/...'; run from the InternNav root. - os.chdir(INTERNNAV_ROOT) - exec(code, g) - - -if __name__ == "__main__": - main() diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/diagnose_engine_gap.py b/recipes/internvla-n1-dualvln/trt-edgellm/diagnose_engine_gap.py deleted file mode 100644 index ddbce12..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/diagnose_engine_gap.py +++ /dev/null @@ -1,130 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Isolate the loss that appears only inside the engine. - -NVFP4 measures 0.9786 as fake quant in PyTorch and 0.9310 through the engine. Activation -quantization accounts for 0.009 of the total; the other 0.048 shows up only once the engine -runs, and FP8 has nothing comparable (0.006 in total). This measures that engine-specific -error rather than inferring it from two numbers taken in different environments. - -The reference is what makes it work. Every other check here compares against the -*unquantized* model, which folds quantization error and engine error into one figure. This -one runs the fake-quant model and the engine **in the same process against the same -inputs**, so whatever separates them is the engine alone: kernel selection, block-scale -arithmetic, accumulation order, or a miscompile. - -Read the output as: - - cosine ~1.0 the engine reproduces correct W4A4 and the loss is quantization after all - cosine ~0.95 the engine diverges materially -- the 0.048, localised - cosine << 0.9 the NVFP4 kernel path is badly wrong - -Pass an FP8 engine as a control. FP8 should sit near 1.0; if it does not, suspect the -harness before the NVFP4 path. -""" -import argparse -import os -import sys - -import torch - -_HERE = os.path.dirname(os.path.abspath(__file__)) -sys.path.insert(0, _HERE) -sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "quantize")) # quant_schemes - -import engine_runner as er # noqa: E402 - -PROMPT = ("You are an autonomous navigation assistant. Your task is to go to the kitchen. " - "Where should you go next to stay on track?") - - -def cos(a: torch.Tensor, b: torch.Tensor) -> float: - a = a.double().flatten() - b = b.double().flatten() - return float(a @ b / (a.norm() * b.norm())) - - -def build_fake_quant(base_ckpt: str, scheme: str, strategy: str, n_calib: int, device: str): - import modelopt.torch.quantization as mtq - from quant_schemes import build_quant_config - from transformers import AutoProcessor, AutoTokenizer, Qwen2_5_VLForConditionalGeneration - - tokenizer = AutoTokenizer.from_pretrained(base_ckpt) - processor = AutoProcessor.from_pretrained(base_ckpt, min_pixels=128 * 28 * 28, - max_pixels=2048 * 32 * 32) - model = Qwen2_5_VLForConditionalGeneration.from_pretrained( - base_ckpt, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True).to(device).eval() - - texts = [PROMPT] * n_calib - - def forward_loop(m): - for t in texts: - with torch.no_grad(): - m(tokenizer(t, return_tensors="pt")["input_ids"].to(device)) - - model = mtq.quantize(model, build_quant_config(scheme, strategy), - forward_loop=forward_loop) - return model.eval(), tokenizer, processor - - -def main() -> int: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--repkg_ckpt", required=True) - p.add_argument("--engine_path", required=True) - p.add_argument("--scheme", default="nvfp4_default") - p.add_argument("--strategy", default="s1") - p.add_argument("--num_calib_samples", type=int, default=8) - p.add_argument("--device", default="cuda") - args = p.parse_args() - - print(f"[1/3] fake quant in PyTorch ({args.scheme})") - model, tokenizer, _ = build_fake_quant(args.repkg_ckpt, args.scheme, args.strategy, - args.num_calib_samples, args.device) - - ids = tokenizer(PROMPT, return_tensors="pt")["input_ids"].to(args.device) - inner = model.model - embed = inner.get_input_embeddings() if hasattr(inner, "get_input_embeddings") \ - else model.get_input_embeddings() - embeds = embed(ids) - - # The engine emits hidden states *before* the final norm, while - # output_hidden_states[-1] is the post-norm tensor. Comparing the two reads ~0.49 for - # a known-good FP8 engine, so hook the norm and take its input instead. (Found by - # running the FP8 control, which is why the control is not optional.) - lm = inner.language_model if hasattr(inner, "language_model") else inner - captured = {} - handle = lm.norm.register_forward_hook( - lambda mod, inp, out: captured.update(pre=inp[0].detach())) - with torch.inference_mode(): - inner(inputs_embeds=embeds, use_cache=False) - handle.remove() - h_pt = captured["pre"][0].float().cpu() - print(f" hidden states (pre-norm) {tuple(h_pt.shape)}") - - del model - torch.cuda.empty_cache() - - print(f"[2/3] engine {args.engine_path}") - os.environ["ENGINE_PATH"] = args.engine_path - seq = embeds.shape[1] - pos = torch.arange(seq, device=args.device).view(1, 1, -1).expand(3, 1, -1) - rope = er.build_mrope_table(pos, args.device) - _, hidden = er.run_engine(embeds.half(), rope) - h_eng = hidden[0].float().cpu() - - print("[3/3] compare") - n = min(h_pt.shape[0], h_eng.shape[0]) - per_tok = [cos(h_pt[i], h_eng[i]) for i in range(n)] - print("\n=== engine vs fake quant, same weights and same activation quantization ===") - print(f" full-sequence cosine : {cos(h_pt[:n], h_eng[:n]):.6f}") - print(f" last token : {per_tok[-1]:.6f}") - print(f" worst token : {min(per_tok):.6f} (position {per_tok.index(min(per_tok))})") - print("\n Near 1.0 means the engine reproduces correct W4A4 and the loss lies in") - print(" quantization. Materially below means the divergence is the engine's own.") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/dump_system1_calib.py b/recipes/internvla-n1-dualvln/trt-edgellm/dump_system1_calib.py deleted file mode 100644 index 931e717..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/dump_system1_calib.py +++ /dev/null @@ -1,142 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Capture real System-1 calibration tensors, for quantizing traj_dit and the memory block. - -FP8 in TensorRT is explicit-quantization only, so the ONNX has to carry Q/DQ nodes, which -means a PTQ pass in PyTorch first, which means calibration data. For System 1 that data is -awkward to synthesize: ``traj_dit`` consumes ``z_latents`` produced by System 2 through -``cond_projector``, and the memory block consumes normalized navigation frames. Random -tensors have the wrong scale in both cases, and FP8 calibration is amax-based, so wrong -scale means wrong scale factors. - -This runs the real thing. For each VLN sample it takes real frames, runs System 2's own -``generate_latents`` for a real ``z``, then taps every ``traj_dit`` call the diffusion loop -makes. The result is cached to disk because capturing it runs System 2 on every sample: -re-quantizing with a different config then costs a minute rather than the whole pipeline. - -Run this under the transformers 4.51 environment (Python 3.10 here):: - - INTERNNAV_PATH=~/InternNav PYTHONPATH=~/InternNav \\ - python dump_system1_calib.py --calib_data_root work/calib_scenes \\ - --output_path work/system1_calib.pt -""" -import argparse -import os -import sys - -import torch - -_HERE = os.path.dirname(os.path.abspath(__file__)) -sys.path.insert(0, _HERE) -sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "quantize")) - -import internvla_compat # noqa: E402 - -SEED = 12345 - - -def parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--internvla_ckpt", - default=os.path.expanduser( - os.environ.get("INTERNVLA_CKPT", - "~/InternNav/checkpoints/InternVLA-N1-DualVLN"))) - p.add_argument("--calib_data_root", required=True, - help="LeRobot scene root, as used by quantize/benchmark_accuracy.py") - p.add_argument("--output_path", required=True) - p.add_argument("--num_samples", type=int, default=4, - help="VLN samples to draw. Each contributes num_inference_steps traj_dit " - "batches, so 4 gives 40 -- ample for amax calibration.") - p.add_argument("--num_sample_trajs", type=int, default=32) - p.add_argument("--num_inference_steps", type=int, default=10) - p.add_argument("--num_frames", type=int, default=2) - p.add_argument("--device", default="cuda") - return p.parse_args() - - -def main() -> int: - args = parse_args() - internvla_compat.apply_all(need_system1=True, allow_missing_depth=False) - - import numpy as np - import prompt_builder as pb - from PIL import Image - from benchmark_accuracy import discover_samples - from internnav.model.basemodel.internvla_n1.internvla_n1 import ( - InternVLAN1ForCausalLM, InternVLAN1ModelConfig) - from transformers import AutoProcessor - - print(f"[1/4] Loading {args.internvla_ckpt}") - config = InternVLAN1ModelConfig.from_pretrained(args.internvla_ckpt) - model = InternVLAN1ForCausalLM.from_pretrained( - args.internvla_ckpt, config=config, torch_dtype=torch.bfloat16, - attn_implementation="sdpa", low_cpu_mem_usage=True).to(args.device).eval() - processor = AutoProcessor.from_pretrained(args.internvla_ckpt, - min_pixels=128 * 28 * 28, - max_pixels=2048 * 32 * 32) - - print(f"[2/4] Drawing {args.num_samples} real VLN samples") - samples = discover_samples(args.calib_data_root, args.num_samples, seed=SEED) - if not samples: - print(f"[ERROR] no samples under {args.calib_data_root}") - return 1 - - inner = model.get_model() - dit_batches, image_batches = [], [] - - _real_fwd = inner.traj_dit.forward - - def _tap(x, timestep, z_latents, *a, **kw): - dit_batches.append((x.detach().half().cpu(), - timestep.detach().cpu(), - z_latents.detach().half().cpu())) - return _real_fwd(x, timestep, z_latents, *a, **kw) - - print("[3/4] Running System 2 -> System 1 on each sample") - torch.manual_seed(SEED) - for i, sample in enumerate(samples): - paths = sample["images"][-args.num_frames:] - frames = [np.asarray(Image.open(p).convert("RGB").resize((224, 224))) / 255.0 - for p in paths] - images_dp = torch.from_numpy(np.stack(frames)).unsqueeze(0) - images_dp = images_dp.to(args.device, torch.bfloat16) # [1, T, 224, 224, 3] - - enc = pb.build_sample_inputs(sample, processor) - enc = {k: v.to(args.device) for k, v in enc.items() if isinstance(v, torch.Tensor)} - with torch.no_grad(): - # The real bridge output, not a draw: generate_latents runs the LLM with the - # learned latent queries appended and returns the normalized TRAJ hidden states. - z = model.generate_latents(enc["input_ids"], enc.get("pixel_values"), - enc.get("image_grid_thw")) - - chw = images_dp.permute(0, 1, 4, 2, 3) - norm = ((chw - model._resnet_mean) / model._resnet_std).flatten(0, 1) - image_batches.append(norm.to(torch.bfloat16).float().cpu()) - - inner.traj_dit.forward = _tap - model.generate_traj(z, images_dp, - num_sample_trajs=args.num_sample_trajs, - num_inference_steps=args.num_inference_steps) - inner.traj_dit.forward = _real_fwd - print(f" sample {i + 1}/{len(samples)}: {len(dit_batches)} traj_dit batches") - - payload = { - "dit_batches": dit_batches, - "image_batches": image_batches, - "num_sample_trajs": args.num_sample_trajs, - "num_inference_steps": args.num_inference_steps, - } - print("[4/4] Writing") - os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True) - torch.save(payload, args.output_path) - print(f"\nWrote {args.output_path} " - f"({os.path.getsize(args.output_path) / 1e6:.1f} MB): " - f"{len(dit_batches)} traj_dit batches, {len(image_batches)} image batches") - print("Now run quantize_system1.py (same environment) to PTQ and build.") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/engine_policy.py b/recipes/internvla-n1-dualvln/trt-edgellm/engine_policy.py deleted file mode 100644 index 5b26be1..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/engine_policy.py +++ /dev/null @@ -1,119 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Engine-backed InternVLA-N1 System 2 policy — drop-in for closed-loop / sim eval. - -Subclasses the real ``InternVLAN1Net`` policy and reuses ALL of its logic (prompt building, -two-turn look-down, coordinate/action decode, history) unchanged. The only thing swapped is the -System 2 LLM text generation: ``self.model.generate`` is routed to the TensorRT-Edge-LLM engine -(``llm_inference``) instead of PyTorch. This is the FP8-quantized decision path — the part that -determines navigation, so a closed-loop run of this policy measures the *engine's* SR. - -Everything else stays as the reference policy: - * ``generate_latents`` (the S2→S1 z_latents bridge) stays PyTorch here — validated separately at - cosine ≥ 0.998 vs the engine (verify_system2_latents.py). Set VLN_ENGINE_LATENTS=1 to also route - it through the engine (uses lib/engine_runner.run_engine). - * System 1 (traj_dit/memory) stays as the reference policy configures it (BF16 engines available - via verify_system1.py / lib/trt_torch.py). - -Wiring (see HANDOVER.md): register this class under the eval config's ``policy_name`` (or set the -agent's ``self.policy`` to it), then run ``scripts/eval/eval.py`` in the sim as usual. Requires the -same env as the verify scripts: ``TRT_EDGE_LLM``, ``EDGELLM_PLUGIN_PATH``, and the engine dirs. - -NOTE: this cannot be closed-loop-tested without the Habitat/InternUtopia simulator. It is verified -here at the method level (engine text == direct llm_inference; s2_step returns a valid S2Output). -""" -import json -import os -import subprocess -import tempfile - -import torch - -from internnav.model.basemodel.internvla_n1.internvla_n1_policy import InternVLAN1Net - -# Engine locations (override via env). Defaults match the VLN-Opt repro layout. -TRT_EDGE_LLM = os.path.expanduser(os.environ.get("TRT_EDGE_LLM", "~/modelopt/TensorRT-Edge-LLM")) -_WORK = os.path.expanduser(os.environ.get("WORK_DIR", "~/vln-opt-work")) -LLM_ENGINE_DIR = os.path.expanduser(os.environ.get( - "VLN_LLM_ENGINE_DIR", os.path.join(_WORK, "engines/system2_llm_fp8"))) -VIS_ENGINE_DIR = os.path.expanduser(os.environ.get( - "VLN_VIS_ENGINE_DIR", os.path.join(_WORK, "engines/system2_visual"))) - - -class EngineInternVLAN1Net(InternVLAN1Net): - """InternVLA-N1 policy whose System 2 LLM text generation runs on the TRT engine.""" - - def __init__(self, config): - super().__init__(config) # loads model, processor, all real logic - self._llm_engine_dir = LLM_ENGINE_DIR - self._vis_engine_dir = VIS_ENGINE_DIR - self._inference_bin = os.path.join(TRT_EDGE_LLM, "build/examples/llm/llm_inference") - self._env = dict( - os.environ, - EDGELLM_PLUGIN_PATH=os.environ.get( - "EDGELLM_PLUGIN_PATH", - os.path.join(TRT_EDGE_LLM, "build/libNvInfer_edgellm_plugin.so")), - ) - for p in (self._inference_bin, self._llm_engine_dir, self._vis_engine_dir): - if not os.path.exists(p): - raise FileNotFoundError(f"engine component missing: {p}") - # Route the LLM text generation through the engine. self.model keeps its other methods - # (generate_latents / generate_traj) so the S2→S1 bridge and System 1 stay unchanged. - self._pt_generate = self.model.generate - self.model.generate = self._engine_generate - print(f"[EnginePolicy] S2 LLM text -> engine {os.path.basename(self._llm_engine_dir)} " - f"(visual {os.path.basename(self._vis_engine_dir)})") - - # -- engine-backed replacement for self.model.generate --------------------------------- # - def _engine_generate(self, *args, **kwargs): - """Mirror the reference generate contract: return output_ids = [prompt_ids | generated_ids] - as a LongTensor, so the real s2_step decode (`output_ids[0][prompt_len:]`) and - `generate_latents(output_ids, ...)` work unchanged. Text comes from the TRT engine, driven - by the exact prompt the policy already built in self.conversation_history.""" - input_ids = kwargs.get("input_ids") - if input_ids is None and args: - input_ids = args[0] - dev = input_ids.device - - tmp = tempfile.mkdtemp(prefix="engpol_") - try: - # Rebuild the llm_inference messages from the policy's own conversation history, - # saving each PIL image (already resized by the policy) to a file. - messages, k = [], 0 - for turn in self.conversation_history: - content = [] - for it in turn["content"]: - if it["type"] == "image": - fp = os.path.join(tmp, f"img_{k}.png") - k += 1 - it["image"].save(fp) - content.append({"type": "image", "image": fp}) - else: - content.append({"type": "text", "text": it["text"]}) - messages.append({"role": turn["role"], "content": content}) - - max_new = int(kwargs.get("max_new_tokens", 64) or 64) - in_json = os.path.join(tmp, "in.json") - out_json = os.path.join(tmp, "out.json") - json.dump({"batch_size": 1, "temperature": 0.0, "top_p": 1.0, "top_k": 1, - "max_generate_length": max_new, - "requests": [{"messages": messages}]}, open(in_json, "w")) - subprocess.run( - [self._inference_bin, "--engineDir", self._llm_engine_dir, - "--multimodalEngineDir", self._vis_engine_dir, - "--inputFile", in_json, "--outputFile", out_json], - env=self._env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True) - text = json.load(open(out_json))["responses"][0]["output_text"].strip() - finally: - import shutil - shutil.rmtree(tmp, ignore_errors=True) - - gen_ids = self.tokenizer(text, return_tensors="pt", add_special_tokens=False - ).input_ids.to(dev) - return torch.cat([input_ids, gen_ids], dim=1) - - -def build_engine_policy(config): - """Factory mirroring how the agent builds the policy (`policy(config=...)`).""" - return EngineInternVLAN1Net(config) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/engine_runner.py b/recipes/internvla-n1-dualvln/trt-edgellm/engine_runner.py deleted file mode 100644 index ab79d9a..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/engine_runner.py +++ /dev/null @@ -1,235 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Shared TensorRT engine runner + mRoPE table builder for the System 2 LLM engine. - -Drives the FP8 LLM engine directly from Python: builds the 3D mRoPE cos/sin table from the -reference `position_ids` (per-token, merged by mrope_section (16,24,24), layout [cos64|sin64]), -binds torch GPU buffers to the TensorRT context, and returns (logits, hidden_states) for a -single prefill. Used by the verification scripts. Select an engine via the ENGINE_PATH env var. -""" -import os -import sys -import ctypes -import json -_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, _ROOT) -# TensorRT ships with JetPack outside the venv. -sys.path.append(os.environ.get("SYSTEM_SITE", "/usr/lib/python3.12/dist-packages")) -import tensorrt as trt # noqa: E402 -import torch # noqa: E402 -from PIL import Image # noqa: E402 - - -def _env(name, default): - return os.path.expanduser(os.environ.get(name, default)) - - -WORK_DIR = _env("WORK_DIR", "~/vln-opt-work") -REPKG = _env("REPKG_CKPT", os.path.join(WORK_DIR, "qwen25vl_system2")) -ENGINE = _env("ENGINE_PATH", os.path.join(WORK_DIR, "engines/s1_fp8/llm/llm.engine")) -TRT_EDGELLM_DIR = _env("TRT_EDGELLM_DIR", "~/modelopt/TensorRT-Edge-LLM") -PLUGIN = _env("EDGELLM_PLUGIN_PATH", - os.path.join(TRT_EDGELLM_DIR, "build/libNvInfer_edgellm_plugin.so")) -# The bridge tensors are emitted next to the repackaged checkpoint by -# repackage_system2.py, so the 16 GB original is not needed here. -CKPT = _env("INTERNVLA_CKPT", REPKG) -IMAGE = os.path.expanduser(os.environ.get( - "IMAGE_PATH", os.path.join(TRT_EDGELLM_DIR, - "examples/multimodal/pics/giant_panda.jpeg"))) - -# Qwen2.5-VL-7B / InternVLA-N1 System 2 geometry. Inlined rather than imported: this -# repository has no shared-library convention, each recipe stands alone. -THETA = 1_000_000.0 -HEAD_DIM = 128 -N_LAYERS = 28 -N_KV = 4 -HIDDEN = 3584 -ROPE_MAXPOS = 4096 # must match --maxKVCacheCapacity at build time -MROPE_SECTION = [16, 24, 24] # (T, H, W) -TRT2TORCH = {trt.DataType.HALF: torch.float16, trt.DataType.FLOAT: torch.float32, - trt.DataType.INT32: torch.int32, trt.DataType.INT64: torch.int64, - trt.DataType.BF16: torch.bfloat16} -_ENG = {} - - -def cos(a, b): - a, b = a.flatten().float(), b.flatten().float() - return torch.nn.functional.cosine_similarity(a, b, dim=0).item() - - -def per_tok_cos(a, b): - a, b = a[0].float(), b[0].float() - return torch.nn.functional.cosine_similarity(a, b, dim=-1).mean().item() - - -def build_mrope_table(position_ids, device): - """position_ids [3,1,S] → rope table [1, ROPE_MAXPOS, 128] (rows 0..S-1 are per-token - mRoPE cos/sin merged by mrope_section; layout [cos64|sin64]).""" - S = position_ids.shape[-1] - half = HEAD_DIM // 2 # 64 - zid = torch.arange(half, dtype=torch.float32, device=device) - inv_freq = THETA ** (-2.0 * zid / HEAD_DIM) # [64] - # axis index per freq band: [0]*16 + [1]*24 + [2]*24 - axis = torch.cat([torch.full((MROPE_SECTION[i],), i, device=device) for i in range(3)]) # [64] - pos = position_ids[:, 0, :].float() # [3,S] - # per token s, freq j: angle = pos[axis[j], s] * inv_freq[j] - pos_sel = pos[axis] # [64,S] - ang = pos_sel.T[:, :] * inv_freq[None, :] # [S,64] - c, s = torch.cos(ang), torch.sin(ang) # [S,64] - table = torch.zeros(1, ROPE_MAXPOS, HEAD_DIM, dtype=torch.float32, device=device) - table[0, :S, :half] = c - table[0, :S, half:] = s - return table - - -def load_engine(): - if "eng" not in _ENG: - ctypes.CDLL(PLUGIN, mode=ctypes.RTLD_GLOBAL) - lg = trt.Logger(trt.Logger.ERROR) - trt.init_libnvinfer_plugins(lg, "") - rt = trt.Runtime(lg) - with open(ENGINE, "rb") as f: - _ENG["eng"] = rt.deserialize_cuda_engine(f.read()) - _ENG["rt"] = rt - return _ENG["eng"] - - -def run_engine(embeds_half, rope_table): - S = embeds_half.shape[1] - dev = embeds_half.device - eng = load_engine() - ctx = eng.create_execution_context() - ctx.set_optimization_profile_async(0, torch.cuda.current_stream().cuda_stream) - context_lengths = torch.tensor([S], dtype=torch.int32, device=dev) - kvcache_start = torch.zeros(1, dtype=torch.int32, device=dev) - last_token = torch.tensor([[S - 1]], dtype=torch.int64, device=dev) - kv_cap = ROPE_MAXPOS - kv_cache = [torch.zeros(1, 2, N_KV, kv_cap, HEAD_DIM, dtype=torch.float16, device=dev) - for _ in range(N_LAYERS)] - feed = { - "inputs_embeds": (embeds_half.contiguous(), None), - "rope_rotary_cos_sin": (rope_table.contiguous(), None), - "context_lengths": (context_lengths, None), - "kvcache_start_index": (kvcache_start, (0,)), - "last_token_ids": (last_token, None), - } - for i in range(N_LAYERS): - feed[f"past_key_values_{i}"] = (kv_cache[i], (1, 2, N_KV, kv_cap, HEAD_DIM)) - for name, (t, shp) in feed.items(): - ctx.set_input_shape(name, shp if shp else tuple(t.shape)) - ctx.set_tensor_address(name, t.data_ptr()) - outs = {} - for i in range(eng.num_io_tensors): - n = eng.get_tensor_name(i) - if eng.get_tensor_mode(n) != trt.TensorIOMode.OUTPUT: - continue - if n.startswith("present_key_values_"): - li = int(n.rsplit("_", 1)[1]) - ctx.set_tensor_address(n, kv_cache[li].data_ptr()) - continue - shp = tuple(int(d) for d in ctx.get_tensor_shape(n)) - shp = tuple(S if d < 0 else d for d in shp) - t = torch.empty(shp, dtype=TRT2TORCH[eng.get_tensor_dtype(n)], device=dev) - outs[n] = t - ctx.set_tensor_address(n, t.data_ptr()) - outs["_kv"] = kv_cache - ok = ctx.execute_async_v3(torch.cuda.current_stream().cuda_stream) - torch.cuda.synchronize() - assert ok - return outs["logits"], outs["hidden_states"] - - -def main(): - dev = "cuda" - torch.manual_seed(0) - from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration - print(f"[1/5] Load repackage (transformers) | engine={os.path.basename(ENGINE)}") - model = Qwen2_5_VLForConditionalGeneration.from_pretrained( - REPKG, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2", - low_cpu_mem_usage=True).to(dev).eval() - proc = AutoProcessor.from_pretrained(REPKG, trust_remote_code=True, - min_pixels=128 * 28 * 28, max_pixels=1024 * 28 * 28) - backbone = model.model.language_model if hasattr(model.model, "language_model") else model.model - final_norm = backbone.norm - - msgs = [{"role": "user", "content": [ - {"type": "image", "image": IMAGE}, - {"type": "text", "text": "Please describe the image."}]}] - text = proc.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) - inp = proc(text=[text], images=[Image.open(IMAGE).convert("RGB")], - return_tensors="pt").to(dev) - S = inp["input_ids"].shape[1] - print(f" seq_len={S} (image grid_thw={inp['image_grid_thw'].tolist()})") - - # hook inner LM to capture the reference inputs_embeds + position_ids - cap = {} - - def pre_hook(mod, args, kwargs): - cap["inputs_embeds"] = kwargs.get("inputs_embeds") - cap["position_ids"] = kwargs.get("position_ids") - h1 = backbone.register_forward_pre_hook(pre_hook, with_kwargs=True) - - def norm_hook(mod, i, o): - cap["pre"] = i[0].detach() - cap["post"] = o.detach() - h2 = final_norm.register_forward_hook(norm_hook) - - print("[2/5] Reference forward (WITH image)") - with torch.no_grad(): - ref_out = model(**inp, use_cache=False) - h1.remove() - h2.remove() - embeds = cap["inputs_embeds"] - pos = cap["position_ids"] - if embeds is None: - # fallback: build inputs_embeds externally - raise RuntimeError("inputs_embeds not captured - wrong hook target") - print(f" captured inputs_embeds {tuple(embeds.shape)} position_ids {tuple(pos.shape)}") - print(f" position_ids axis-range T[{pos[0].min()}..{pos[0].max()}] " - f"H[{pos[1].min()}..{pos[1].max()}] W[{pos[2].min()}..{pos[2].max()}]") - ref_pre = cap["pre"].float() - ref_post = cap["post"].float() - ref_logits_last = ref_out.logits[:, -1, :].float() - - print("[3/5] Build merged 3D mRoPE + run engine") - rope = build_mrope_table(pos, dev) - eng_logits, eng_hs = run_engine(embeds.to(torch.float16), rope) - eng_hs = eng_hs.float() - with torch.no_grad(): - eng_post = final_norm(eng_hs.to(torch.bfloat16)).float() - - print("[4/5] Compare hidden_states / logits\n" + "=" * 60) - a, b = eng_hs[0].float(), ref_pre[0].float() - ptc = torch.nn.functional.cosine_similarity(a, b, dim=-1) - print(" per-pos PRE cosine[:8]:", " ".join(f"{ptc[i]:.3f}" for i in range(min(S, 8))), - "... last:", f"{ptc[-1]:.3f}") - print(f" hidden PRE-norm per-token cosine = {per_tok_cos(eng_hs, ref_pre):.6f}") - print(f" hidden POST-norm per-token cosine = {per_tok_cos(eng_post, ref_post):.6f}") - ea, ra = eng_logits[0, 0].argmax().item(), ref_logits_last[0].argmax().item() - print(f" logits last cosine = {cos(eng_logits[0,0], ref_logits_last[0]):.6f} " - f"| argmax eng={ea!r} ref={ra!r} match={ea==ra}") - - print("[5/5] z_latents (the original cond_projector) engine vs reference") - from safetensors import safe_open - idx = json.load(open(os.path.join(CKPT, "model.safetensors.index.json")))["weight_map"] - cp = {} - for k in idx: - if k.startswith("model.cond_projector"): - with safe_open(os.path.join(CKPT, idx[k]), framework="pt") as f: - cp[k.replace("model.cond_projector.", "")] = f.get_tensor(k).float().to(dev) - - def cond_project(x): - x = torch.nn.functional.linear(x, cp["0.weight"], cp.get("0.bias")) - x = torch.nn.functional.gelu(x) - return torch.nn.functional.linear(x, cp["2.weight"], cp.get("2.bias")) - z_eng, z_ref = cond_project(eng_post), cond_project(ref_post) - zc = per_tok_cos(z_eng, z_ref) - print(f" z_latents({z_eng.shape[-1]}) per-token cosine = {zc:.6f}") - print("\n" + ("✅ WITH-IMAGE numeric PASS (z_latents ≥ 0.99)" if zc > 0.99 - else f"z_latents cosine {zc:.4f} < 0.99 - check mRoPE/vision")) - return 0 if zc > 0.99 else 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py b/recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py deleted file mode 100644 index d5aaaa6..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/export_memory_block.py +++ /dev/null @@ -1,123 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Step 5 - export the System 1 memory block (DINOv2 rgb_model + memory_encoder + -rgb_resampler) to ONNX and build the BF16 engine. -""" -import os -import sys -import time -import subprocess -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -sys.path.append("/usr/lib/python3.12/dist-packages") -sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "lib")) -import numpy as np # noqa: E402 -import torch # noqa: E402 -from PIL import Image # noqa: E402 - -ACTIVE = os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) -CKPT = os.path.join(ACTIVE, "checkpoints/InternVLA-N1-DualVLN") -IMG = os.path.expanduser("~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg") -ONNX = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.onnx") -ENG = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.engine") - - -# MemBlock lives in memblock.py; defining it twice is how the two copies drift. -from memblock import MemBlock # noqa: E402 - - -def main(): - dev = "cuda" - # disable fused MHA fastpath (_transformer_encoder_layer_fwd is not ONNX-exportable) - try: - torch.backends.mha.set_fastpath_enabled(False) - except Exception as e: - print(" (mha fastpath toggle:", e, ")") - if ACTIVE not in sys.path: - sys.path.insert(0, ACTIVE) - from internvla_compat import apply_all - apply_all(need_system1=True, allow_missing_depth=True) - from internnav.model.basemodel.internvla_n1.internvla_n1 import ( - InternVLAN1ForCausalLM, InternVLAN1ModelConfig) - print("[1/6] Load full model") - cfg = InternVLAN1ModelConfig.from_pretrained(CKPT) - model = InternVLAN1ForCausalLM.from_pretrained(CKPT, config=cfg, torch_dtype=torch.bfloat16, - attn_implementation="sdpa", low_cpu_mem_usage=True).to(dev).eval() - m = model.get_model() - - print("[2/6] Build input (matching the reference normalization)") - a = np.array(Image.open(IMG).convert("RGB").resize((224, 224))) / 255.0 - tt = torch.from_numpy(a).float() - images_dp = torch.stack([tt, tt]).unsqueeze(0).to(dev) # [1,2,224,224,3] - dtype = torch.bfloat16 - rmean = model._resnet_mean if hasattr(model, "_resnet_mean") else m._resnet_mean - rstd = model._resnet_std if hasattr(model, "_resnet_std") else m._resnet_std - x = images_dp.permute(0, 1, 4, 2, 3) # [1,2,3,224,224] - x = (x - rmean) / rstd - x = x.flatten(0, 1).to(dtype) # [2,3,224,224] - print(f" input {tuple(x.shape)}") - - print("[3/6] Ref base (PyTorch BF16) memory_tokens + latency") - block = MemBlock(m.rgb_model, m.memory_encoder, m.rgb_resampler).eval() - with torch.no_grad(): - ref = block(x).float() - print(f" memory_tokens {tuple(ref.shape)}") - - def lat(fn, warm=3, n=10): - for _ in range(warm): - fn() - torch.cuda.synchronize() - ts = [] - for _ in range(n): - torch.cuda.synchronize() - t0 = time.perf_counter() - fn() - torch.cuda.synchronize() - ts.append(time.perf_counter() - t0) - return sum(ts) / len(ts) * 1000 - with torch.no_grad(): - pt_ms = lat(lambda: block(x)) - - print("[4/6] Export ONNX (FP32 to avoid mixed precision; build BF16 later)") - block_fp32 = MemBlock(m.rgb_model.float(), m.memory_encoder.float(), m.rgb_resampler.float()).eval() - xf = x.float() - os.makedirs(os.path.dirname(ONNX), exist_ok=True) - with torch.inference_mode(): - torch.onnx.export(block_fp32, (xf,), ONNX, input_names=["images"], output_names=["memory_tokens"], - opset_version=19, do_constant_folding=True, export_params=True, dynamo=False) - print(f" {os.path.getsize(ONNX)/1e6:.1f} MB") - - print("[5/6] Build BF16 engine (fixed shape 2x3x224x224)") - cmd = ["/usr/src/tensorrt/bin/trtexec", f"--onnx={ONNX}", f"--saveEngine={ENG}", "--bf16"] # static shape - r = subprocess.run(cmd, capture_output=True, text=True) - print( - f" build {'OK' if os.path.exists(ENG) else 'FAIL'} | {os.path.getsize(ENG)/1e6 if os.path.exists(ENG) else 0:.1f} MB") # noqa: E501 - if not os.path.exists(ENG): - print(" STDERR:", r.stderr[-600:]) - return 1 - - print("[6/6] Parity + latency (TRT vs base)") - # The parity check needs the TensorRT Python bindings, which JetPack ships for - # Python 3.12 only -- while this export needs transformers 4.51, which lives in the - # 3.10 environment. The engine is already built and valid at this point, so a missing - # binding must not turn a successful build into a failure. Run - # verify/verify_system1.py under the 3.12 environment to check parity. - try: - from trt_torch import Engine - except ImportError as exc: - print(f"\n[skip] parity check unavailable in this interpreter: {exc}") - print(" The engine was built successfully. To check it, run") - print(" verify/verify_system1.py under the TensorRT (Python 3.12) environment.") - return 0 - eng = Engine(ENG) - out = eng(images=x.float().contiguous()) - out = (out.get("memory_tokens") if isinstance(out, dict) else out).float() - cos = torch.nn.functional.cosine_similarity(ref.flatten(), out.flatten(), dim=0).item() - trt_ms = lat(lambda: eng(images=x.float().contiguous())) - print(f" parity cos={cos:.5f} rel-L2={(ref-out).norm()/ref.norm():.4f}") - print(f" latency PyTorch {pt_ms:.2f}ms → TRT {trt_ms:.2f}ms = {pt_ms/trt_ms:.2f}x") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/export_traj_dit.py b/recipes/internvla-n1-dualvln/trt-edgellm/export_traj_dit.py deleted file mode 100644 index eef1409..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/export_traj_dit.py +++ /dev/null @@ -1,70 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Step 4 - export the System 1 traj_dit (NextDiT) core to ONNX with a dynamic -z_latents length and build the BF16 engine (trtexec). -""" -import os -import sys -import subprocess -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -sys.path.append("/usr/lib/python3.12/dist-packages") -sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "lib")) -import torch # noqa: E402 -from traj_dit_loader import load_traj_dit # noqa: E402 - -OUT = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "onnx/system1_traj_dit_async.onnx") -ENG = os.path.join( - os.environ.get( - "WORK_DIR", - os.path.expanduser("~/vln-opt-work")), - "onnx/system1_traj_dit_bf16.engine") # noqa: E131 -ZLEN = int(os.environ.get("ZLEN", "36")) # observed z_latents length for the async path - - -def main(): - dev = "cuda" - N, WP, DIM = 32, 32, 384 - print(f"[1/3] Load traj_dit (FP32) | z_latents seq (async, dynamic, opt={ZLEN})") - dit, zdim = load_traj_dit() - dit = dit.to(torch.float32).eval() - x = torch.randn(2 * N, WP, DIM, dtype=torch.float32, device=dev) - ts = torch.ones(2 * N, dtype=torch.int64, device=dev) - z = torch.randn(2 * N, ZLEN, zdim, dtype=torch.float32, device=dev) - - class W(torch.nn.Module): - def __init__(s, d): - super().__init__() - s.d = d - - def forward(s, x, timestep, z_latents): return s.d(x=x, timestep=timestep, z_latents=z_latents) # noqa: E704 - w = W(dit).eval() - with torch.no_grad(): - ref = w(x, ts, z) - print(f" forward OK -> {tuple(ref.shape)}") - os.makedirs(os.path.dirname(OUT), exist_ok=True) - # dynamic: batch (dim0) + z_latents seq (dim1) - dyn = {"x": {0: "batch"}, "timestep": {0: "batch"}, "z_latents": {0: "batch", 1: "zlen"}, "output": {0: "batch"}} - print(f"[2/3] Export ONNX (dynamo=False, opset 19) → {OUT}") - with torch.inference_mode(): - torch.onnx.export(w, (x, ts, z), OUT, input_names=["x", "timestep", "z_latents"], - output_names=["output"], opset_version=19, do_constant_folding=True, - export_params=True, dynamic_axes=dyn, dynamo=False) - print(f" {os.path.getsize(OUT)/1e6:.1f} MB") - B = 2 * N - cmd = ["/usr/src/tensorrt/bin/trtexec", f"--onnx={OUT}", f"--saveEngine={ENG}", "--bf16", - f"--minShapes=x:{B}x{WP}x{DIM},timestep:{B},z_latents:{B}x4x{zdim}", - f"--optShapes=x:{B}x{WP}x{DIM},timestep:{B},z_latents:{B}x{ZLEN}x{zdim}", - f"--maxShapes=x:{B}x{WP}x{DIM},timestep:{B},z_latents:{B}x64x{zdim}"] - print(f"[3/3] Build engine BF16 (z dynamic 4..64, opt {ZLEN})") - print(" " + " ".join(cmd)) - r = subprocess.run(cmd, capture_output=True, text=True) - wanted = [ln for ln in r.stdout.splitlines() - if any(w in ln for w in ("successfully", "PASSED", "FAILED"))] - print(" " + "\n ".join(wanted)[-400:]) - print(" engine:", ENG, os.path.getsize(ENG) / 1e6 if os.path.exists(ENG) else "MISSING", "MB") - return 0 if os.path.exists(ENG) else 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/internvla_compat.py b/recipes/internvla-n1-dualvln/trt-edgellm/internvla_compat.py deleted file mode 100644 index 40632a0..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/internvla_compat.py +++ /dev/null @@ -1,216 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Runtime compatibility patches for loading InternVLA-N1 System 1. - -These patches are applied in-memory (attribute reassignment) and never modify the InternNav -source tree. They fix three issues that otherwise prevent building the System 1 modules: - -1. ``build_depthanythingv2`` loads the DepthAnything-V2 checkpoint from a relative path; the - patch loads it from an absolute path (set via DAV2_CKPT / env), and can tolerate a missing - file since ``from_pretrained`` overwrites ``model.rgb_model.*`` from safetensors afterwards. - -2. ``LuminaNextDiT2DModel._set_gradient_checkpointing`` has an old signature that current - diffusers calls with ``enable=`` / ``gradient_checkpointing_func=`` keywords; the patch - accepts both. - -3. ``build_traj_dit`` must pass ``ffn_dim_multiplier = 2/3`` (SwiGLU convention) so the traj_dit - FFN shape matches the trained DualVLN checkpoint (1024, not 1536). -""" -from __future__ import annotations - -import os - -#: Root of the InternNav checkout to load the model from. -ACTIVE_INTERNNAV = os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) - -#: DepthAnything-V2 checkpoint used by System 1's rgb_model. -DAV2_CKPT = os.path.expanduser(os.environ.get( - "DAV2_CKPT", os.path.join(ACTIVE_INTERNNAV, "checkpoints/depth_anything_v2_vits.pth"))) - -#: FFN multiplier the DualVLN checkpoint was trained with (see patch_traj_dit_ffn). -TRAJ_DIT_FFN_MULTIPLIER = 2.0 / 3.0 - - -def assert_active_tree() -> str: - """Confirm ``internnav`` is imported from ACTIVE_INTERNNAV; fail early otherwise. - - Which tree loads depends on ``cwd`` (``''`` heads ``sys.path``), so a stray ``cd`` can - silently swap the model code. This turns that silent failure into a loud one. - """ - import internnav - - tree = os.path.dirname(os.path.dirname(internnav.__file__)) - if os.path.realpath(tree) != os.path.realpath(ACTIVE_INTERNNAV): - raise RuntimeError( - f"internnav loaded from the wrong tree:\n" - f" actual : {tree}\n" - f" expected: {ACTIVE_INTERNNAV}\n" - f"Usually the cwd sits inside a different InternNav checkout. Re-run elsewhere, " - f"or put {ACTIVE_INTERNNAV} at the front of sys.path." - ) - return tree - - -def patch_depth_anything(allow_missing: bool = False) -> bool: - """Replace ``build_depthanythingv2`` with a version that loads from an absolute path. - - Only reassigns a module attribute in memory. Returns True once patched. - """ - import torch - - import internnav.model.basemodel.internvla_n1.internvla_n1_arch as arch - - if getattr(arch, "_vlnopt_patched", False): - return True - - if not os.path.isfile(DAV2_CKPT): - if not allow_missing: - raise FileNotFoundError( - f"DepthAnything-V2 checkpoint not found: {DAV2_CKPT}\n" - f"Required for System 1 (nextdit_async). Pass allow_missing=True to build the " - f"architecture with random init — acceptable because from_pretrained overwrites " - f"model.rgb_model.* from safetensors immediately afterwards." - ) - print(f"[compat] WARNING: missing {DAV2_CKPT} — rgb_model randomly initialized, " - f"relying on from_pretrained to load weights.") - - def _build_dav2_patched(config): - from internnav.model.encoder.depth_anything.depth_anything_v2.dpt import ( - DepthAnythingV2, - ) - - model_configs = { - "vits": {"encoder": "vits", "features": 64, - "out_channels": [48, 96, 192, 384]} - } - dav2 = DepthAnythingV2(**model_configs["vits"]) - if os.path.isfile(DAV2_CKPT): - dav2.load_state_dict(torch.load(DAV2_CKPT, map_location="cpu")) - return dav2.pretrained - - arch.build_depthanythingv2 = _build_dav2_patched - arch._vlnopt_patched = True - return True - - -def patch_gradient_checkpointing() -> bool: - """Reconcile the ``_set_gradient_checkpointing`` signature between nextdit and diffusers. - - nextdit defines ``(self, module, value)`` but current diffusers calls it with - ``(enable=..., gradient_checkpointing_func=...)``. This fails while building traj_dit, so - it must be patched before any System 1 build path. - """ - try: - from internnav.model.basemodel.internvla_n1.nextdit_traj import ( - LuminaNextDiT2DModel, - ) - except ImportError: - return False - - if getattr(LuminaNextDiT2DModel, "_vlnopt_gc_patched", False): - return True - - def _set_gc_compat(self, module=None, value=False, enable=None, - gradient_checkpointing_func=None): - v = enable if enable is not None else value - if module is not None: - module.gradient_checkpointing = v - else: - for m in self.modules(): - if hasattr(m, "gradient_checkpointing"): - m.gradient_checkpointing = v - - LuminaNextDiT2DModel._set_gradient_checkpointing = _set_gc_compat - LuminaNextDiT2DModel._vlnopt_gc_patched = True - return True - - -def patch_traj_dit_ffn(multiplier: float = TRAJ_DIT_FFN_MULTIPLIER) -> bool: - """Patch ``build_traj_dit`` so the traj_dit FFN matches the checkpoint. - - The DualVLN checkpoint was trained with ``ffn_dim_multiplier = 2/3`` (SwiGLU convention), - giving inner_dim 1024 for dim=384. The stock ``build_traj_dit`` never passes this, leaving - inner_dim at 1536 and causing a state_dict size mismatch. Only feed_forward.linear_1/linear_3 - are affected; norm1.linear is 4*dim on both sides and already matches. - """ - import internnav.model.basemodel.internvla_n1.internvla_n1_arch as arch - - if getattr(arch, "_vlnopt_ffn_patched", False): - return True - - _orig = arch.build_traj_dit - - def _build_traj_dit_patched(config): - from diffusers.schedulers import FlowMatchEulerDiscreteScheduler - - from internnav.model.basemodel.internvla_n1.nextdit_crossattn_traj import ( - NextDiTCrossAttn, NextDiTCrossAttnConfig, - ) - - dit_cfg = NextDiTCrossAttnConfig( - latent_embedding_size=arch.LatentEmbSize, - ffn_dim_multiplier=multiplier, # the only difference from the stock builder - ) - dit = NextDiTCrossAttn(dit_cfg) - return dit, FlowMatchEulerDiscreteScheduler() - - _build_traj_dit_patched.__wrapped__ = _orig - arch.build_traj_dit = _build_traj_dit_patched - arch._vlnopt_ffn_patched = True - return True - - -def apply_all(need_system1: bool = True, allow_missing_depth: bool = False) -> None: - """Convenience entry point: verify the tree and apply the needed patches.""" - tree = assert_active_tree() - print(f"[compat] internnav tree: {tree}") - if need_system1: - # Order matters: patch gradient-checkpointing first, since it fails inside the - # traj_dit builder. - if patch_gradient_checkpointing(): - print("[compat] patched LuminaNextDiT2DModel._set_gradient_checkpointing") - patch_depth_anything(allow_missing=allow_missing_depth) - print(f"[compat] patched build_depthanythingv2 -> {DAV2_CKPT}") - patch_traj_dit_ffn() - print(f"[compat] patched build_traj_dit -> ffn_dim_multiplier={TRAJ_DIT_FFN_MULTIPLIER:.4f}") - # Needed on transformers 5.x regardless of System 1; harmless on 4.x. - patch_config_flattening() - - -def patch_config_flattening() -> bool: - """Re-expose the top-level LLM config fields that transformers 5.x nests. - - InternNav reads ``config.hidden_size``, ``config.num_hidden_layers`` and friends off - the top-level config. transformers 4.x flattened them there; 5.x moves them under - ``text_config``, so constructing the model raises a bare - ``'InternVLAN1ModelConfig' object has no attribute 'hidden_size'``. - - This matters beyond tidiness: System-1 export needs InternNav (transformers 4.51) while - the TensorRT Python bindings ship for 3.12 only, so any script needing *both* -- the - System-1 parity check, for one -- cannot run without reconciling them. Copying the - fields back from text_config is the smaller of the two evils; the alternative is - pinning transformers 4.51 into the TensorRT environment and hoping the edgellm exporter - still works there. - """ - try: - from internnav.model.basemodel.internvla_n1.internvla_n1 import ( - InternVLAN1ModelConfig) - except ImportError: - return False - - _orig = InternVLAN1ModelConfig.from_pretrained.__func__ - - def _from_pretrained(cls, *args, **kwargs): - config = _orig(cls, *args, **kwargs) - inner = getattr(config, "text_config", None) - if inner is not None: - for field in ("hidden_size", "num_hidden_layers", "num_attention_heads", - "num_key_value_heads", "intermediate_size", "rms_norm_eps", - "vocab_size", "max_position_embeddings", "rope_theta"): - if not hasattr(config, field) and hasattr(inner, field): - setattr(config, field, getattr(inner, field)) - return config - - InternVLAN1ModelConfig.from_pretrained = classmethod(_from_pretrained) - print("[compat] re-exposed top-level LLM config fields (transformers 5.x nests them)") - return True diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/investigate_nvfp4.py b/recipes/internvla-n1-dualvln/trt-edgellm/investigate_nvfp4.py deleted file mode 100644 index 44520bd..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/investigate_nvfp4.py +++ /dev/null @@ -1,310 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Find out why NVFP4 keeps the text fluent but collapses the System 1 bridge. - -NVFP4 quantizes, exports and generates coherent text, yet z_latents cosine against the -reference falls to 0.647 where FP8 holds 0.99. That is not a paradox, and this script does -not treat it as one -- the two readouts have very different sensitivity: - -* Tokens are ``argmax(lm_head(final_norm(h)))`` over a 152k vocabulary. RMSNorm removes - per-token scale entirely and argmax only cares about ranking, so a hidden state can be - badly distorted and still select the same token. -* z_latents are ``cond_projector(final_norm(h[-4:]))`` -- two linears with a GELU between, - consumed as a continuous 4x768 vector by a diffusion model. GELU is not scale-invariant - and cosine over 3072 numbers has no ranking slack. - -So the expectation is that the hidden states are genuinely damaged and argmax is simply -tolerant. The job is to locate the damage and decide whether it is recoverable. - -Stages, cheapest and most decisive first: - - weights How much error does NVFP4 put into the weights, versus FP8? Pure checkpoint - arithmetic, no forward pass. If NVFP4 weight error is only modestly worse than - FP8's while the bridge is 50x worse, the collapse is not raw weight precision - and the later stages matter. Runs in seconds. - layers Per-layer hidden-state cosine against the unquantized reference. Smooth decay - means a global capacity limit and nothing local to exclude; a knee at specific - layers names them. - channels At the last layer, how concentrated the error is. If masking the top ~1% of - channels by magnitude restores cosine above 0.99, this is outlier clipping - under per-16-element block scaling, and AWQ or a targeted exclusion should fix - it. If the error is spread evenly, no scaling trick will help. - bridge End-to-end z_latents for each checkpoint, the number that actually decides. - -Decision rule, fixed in advance so the investigation is bounded: any variant reaching -z_latents >= 0.99 is promoted to a supported scheme. A best of 0.95-0.99 stays experimental -with the number recorded. If nothing clears 0.95, record NVFP4 as rejected on this model -and stop -- do not keep hunting. -""" -import argparse -import json -import os -import sys - -import numpy as np -import torch - -_HERE = os.path.dirname(os.path.abspath(__file__)) -sys.path.insert(0, _HERE) -sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "quantize")) - - -def cos(a: torch.Tensor, b: torch.Tensor) -> float: - """Cosine in float64. - - float32 accumulation over tensors this size silently returns values above 1.0, which - is how a broken measurement can look like a good one. - """ - a = a.double().flatten() - b = b.double().flatten() - return float(a @ b / (a.norm() * b.norm())) - - -def rel_err(ref: torch.Tensor, other: torch.Tensor) -> float: - ref = ref.double() - return float((other.double() - ref).norm() / ref.norm()) - - -def _has_awq_scales(model_path: str) -> bool: - """True if the checkpoint carries AWQ pre-quant scales.""" - import glob - from safetensors import safe_open - for shard in sorted(glob.glob(os.path.join(model_path, "*.safetensors"))): - with safe_open(shard, framework="pt") as f: - if any(k.endswith("pre_quant_scale") for k in f.keys()): - return True - return False - - -def stage_weights(args, report: dict) -> None: - """Compare per-projection weight error, NVFP4 vs FP8, against the unquantized weights.""" - from load_quantized import dequantize_state_dict - from safetensors import safe_open - import glob - - print("\n=== stage: weights — how much error does each format put in the weights? ===") - - keep = ("layers.0.", "layers.13.", "layers.27.") - base = {} - for shard in sorted(glob.glob(os.path.join(args.repkg_ckpt, "*.safetensors"))): - if os.path.basename(shard) == "bridge.safetensors": - continue - with safe_open(shard, framework="pt") as f: - for k in f.keys(): - if k.endswith(".weight") and any(x in k for x in keep) and "proj" in k: - base[k] = f.get_tensor(k) - - rows = {} - for label, path in (("fp8", args.fp8_ckpt), ("nvfp4", args.nvfp4_ckpt)): - if not path: - continue - # AWQ redistributes scale between each layernorm and the projections that read it - # (measured on this checkpoint: layernorm weights move by 0.012x to 91x while - # plain NVFP4 leaves them at exactly 1.0). Comparing its linear weights against - # the reference's linear weights is therefore meaningless -- the invariant is the - # composition, not either factor. Only the end-to-end bridge stage is valid there. - if _has_awq_scales(path): - print(f" {label:6s} [skip] AWQ checkpoint: per-layer weight comparison is not " - f"meaningful, see the bridge stage") - rows[label] = {"skipped": "awq_rescales_layernorms"} - continue - state = dequantize_state_dict(path) - errs, coss = [], [] - for k, ref in base.items(): - if k in state: - errs.append(rel_err(ref, state[k])) - coss.append(cos(ref, state[k])) - rows[label] = {"n": len(errs), - "rel_err_mean": float(np.mean(errs)), - "cos_mean": float(np.mean(coss))} - print(f" {label:6s} over {len(errs)} projections: " - f"rel-err {100 * np.mean(errs):.2f}% cos {np.mean(coss):.6f}") - - if "fp8" in rows and "nvfp4" in rows: - ratio = rows["nvfp4"]["rel_err_mean"] / max(rows["fp8"]["rel_err_mean"], 1e-12) - print(f"\n NVFP4 weight error is {ratio:.1f}x FP8's.") - print(" Compare that against the bridge gap in the 'bridge' stage: if the bridge") - print(" degrades far more than this ratio, raw weight precision is not the cause.") - rows["nvfp4_over_fp8"] = ratio - report["weights"] = rows - - -def _load(path, device): - from load_quantized import load_for_eval - return load_for_eval(path, device=device) - - -def stage_layers(args, report: dict) -> None: - """Per-layer hidden-state cosine against the unquantized model, on one real prompt.""" - print("\n=== stage: layers — where does the divergence start? ===") - - prompt = ("You are an autonomous navigation assistant. Your task is to go to the " - "kitchen. Where should you go next to stay on track?") - - hiddens = {} - for label, path in (("ref", args.repkg_ckpt), ("fp8", args.fp8_ckpt), - ("nvfp4", args.nvfp4_ckpt)): - if not path: - continue - model, processor, _ = _load(path, args.device) - enc = processor.tokenizer(prompt, return_tensors="pt").to(args.device) - with torch.inference_mode(): - out = model(**enc, output_hidden_states=True) - hiddens[label] = [h[0, -1].float().cpu() for h in out.hidden_states] - del model - torch.cuda.empty_cache() - - if "ref" not in hiddens: - report["layers"] = {"status": "no reference checkpoint"} - return - - table = {} - for label in ("fp8", "nvfp4"): - if label not in hiddens: - continue - per_layer = [cos(r, q) for r, q in zip(hiddens["ref"], hiddens[label])] - table[label] = per_layer - drops = np.diff(per_layer) - worst = int(np.argmin(drops)) + 1 if len(drops) else -1 - print(f" {label:6s} layer 0 {per_layer[0]:.6f} -> " - f"mid {per_layer[len(per_layer) // 2]:.6f} -> " - f"final {per_layer[-1]:.6f}") - print(f" biggest single-layer drop at layer {worst} " - f"({drops.min():.6f})" if len(drops) else "") - report["layers"] = table - print("\n A smooth decline means accumulation and nothing local to exclude;") - print(" a sharp knee names the layers worth excluding from NVFP4.") - - -def stage_channels(args, report: dict) -> None: - """Is the final-layer error concentrated in a few high-magnitude channels?""" - print("\n=== stage: channels — is the error carried by outliers? ===") - - if not args.nvfp4_ckpt: - report["channels"] = {"status": "no nvfp4 checkpoint"} - return - - prompt = ("You are an autonomous navigation assistant. Your task is to go to the " - "kitchen. Where should you go next to stay on track?") - hs = {} - for label, path in (("ref", args.repkg_ckpt), ("nvfp4", args.nvfp4_ckpt)): - model, processor, _ = _load(path, args.device) - enc = processor.tokenizer(prompt, return_tensors="pt").to(args.device) - with torch.inference_mode(): - out = model(**enc, output_hidden_states=True) - hs[label] = out.hidden_states[-1][0, -1].float().cpu() - del model - torch.cuda.empty_cache() - - ref, q = hs["ref"], hs["nvfp4"] - err = (q - ref).abs() - order = torch.argsort(ref.abs(), descending=True) - total = float((err ** 2).sum()) - - result = {"baseline_cos": cos(ref, q), "share_by_topk": {}, "cos_masking_topk": {}} - print(f" baseline final-layer cosine: {result['baseline_cos']:.6f}") - for k in (1, 4, 16, 36, 128): - idx = order[:k] - share = float((err[idx] ** 2).sum()) / max(total, 1e-30) - mask = torch.ones_like(ref, dtype=torch.bool) - mask[idx] = False - result["share_by_topk"][k] = share - result["cos_masking_topk"][k] = cos(ref[mask], q[mask]) - print(f" top {k:4d} channels by |ref|: carry {100 * share:5.1f}% of squared error" - f" | cosine with them masked out: {result['cos_masking_topk'][k]:.6f}") - - report["channels"] = result - best = max(result["cos_masking_topk"].values()) - if best > 0.99: - print("\n Masking a small number of channels restores the signal: this is outlier") - print(" clipping under per-16 block scaling. AWQ or a targeted exclusion is the fix.") - else: - print("\n The error is spread across channels, not carried by a few outliers.") - print(" No amount of scaling or exclusion will recover it -- this is a capacity limit.") - - -def stage_bridge(args, report: dict) -> None: - """End-to-end z_latents per checkpoint -- the number that decides.""" - print("\n=== stage: bridge — z_latents, the acceptance metric ===") - from safetensors import safe_open - - bridge_path = os.path.join(args.repkg_ckpt, "bridge.safetensors") - if not os.path.isfile(bridge_path): - report["bridge"] = {"status": f"bridge.safetensors not found under {args.repkg_ckpt}"} - print(f" [skip] {report['bridge']['status']}") - return - with safe_open(bridge_path, framework="pt") as f: - bt = {k: f.get_tensor(k) for k in f.keys()} - cond = {k.replace("model.cond_projector.", ""): v.float().to(args.device) - for k, v in bt.items() if "cond_projector" in k} - - prompt = ("You are an autonomous navigation assistant. Your task is to go to the " - "kitchen. Where should you go next to stay on track?") - z = {} - for label, path in (("ref", args.repkg_ckpt), ("fp8", args.fp8_ckpt), - ("nvfp4", args.nvfp4_ckpt)): - if not path: - continue - model, processor, _ = _load(path, args.device) - enc = processor.tokenizer(prompt, return_tensors="pt").to(args.device) - with torch.inference_mode(): - out = model(**enc, output_hidden_states=True) - h = out.hidden_states[-1][0, -4:].float() - x = torch.nn.functional.linear(h, cond["0.weight"], cond.get("0.bias")) - x = torch.nn.functional.gelu(x) - z[label] = torch.nn.functional.linear(x, cond["2.weight"], cond.get("2.bias")).cpu() - del model - torch.cuda.empty_cache() - - table = {} - for label in ("fp8", "nvfp4"): - if label in z: - table[label] = cos(z["ref"], z[label]) - print(f" {label:6s} z_latents cosine vs reference: {table[label]:.6f}") - report["bridge"] = table - - if "nvfp4" in table: - verdict = ("supported" if table["nvfp4"] >= 0.99 - else "experimental" if table["nvfp4"] >= 0.95 else "rejected") - print(f"\n Verdict for NVFP4 on this model by the stated rule: {verdict}") - report["verdict"] = verdict - - -def parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--repkg_ckpt", required=True, - help="Repackaged System 2 checkpoint (the unquantized reference)") - p.add_argument("--nvfp4_ckpt", default=None) - p.add_argument("--fp8_ckpt", default=None) - p.add_argument("--work_dir", default=os.path.expanduser("~/vln-opt-work")) - p.add_argument("--device", default="cuda") - p.add_argument("--stage", default="all", - choices=["all", "weights", "layers", "channels", "bridge"]) - return p.parse_args() - - -def main() -> int: - args = parse_args() - out = os.path.join(args.work_dir, "nvfp4_investigation.json") - report = {} - if os.path.isfile(out): - with open(out) as f: - report = json.load(f) - - stages = {"weights": stage_weights, "layers": stage_layers, - "channels": stage_channels, "bridge": stage_bridge} - todo = list(stages) if args.stage == "all" else [args.stage] - for name in todo: - stages[name](args, report) - - os.makedirs(args.work_dir, exist_ok=True) - with open(out, "w") as f: - json.dump(report, f, indent=2) - print(f"\nWrote {out}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/memblock.py b/recipes/internvla-n1-dualvln/trt-edgellm/memblock.py deleted file mode 100644 index 021f52b..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/memblock.py +++ /dev/null @@ -1,25 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""System 1 memory-encoding block, wrapped as a single exportable module. - -Mirrors the reference `generate_traj` (nextdit_async) memory path: - rgb_model.get_intermediate_layers -> memory_encoder -> concat -> rgb_resampler -Input : images_dp_norm [T, 3, 224, 224] (T frames, ImageNet-normalized) -Output: memory_tokens [1, 32, 768] -""" -import torch - - -class MemBlock(torch.nn.Module): - def __init__(self, rgb_model, memory_encoder, rgb_resampler): - super().__init__() - self.rgb = rgb_model - self.me = memory_encoder - self.rr = rgb_resampler - - def forward(self, x): # x = [T, 3, 224, 224] - feat = self.rgb.get_intermediate_layers(x)[0].unflatten(dim=0, sizes=(1, -1)) # [1, T, Np, 384] - f = feat.flatten(1, 2) # [1, T*Np, 384] - mf = self.me(f) - mf = torch.cat([f, mf], dim=-1) - return self.rr(mf) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/quantize_system1.py b/recipes/internvla-n1-dualvln/trt-edgellm/quantize_system1.py deleted file mode 100644 index 0b840aa..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/quantize_system1.py +++ /dev/null @@ -1,262 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""FP8 PTQ for System 1: quantize traj_dit and the memory block, export ONNX, build engines. - -System 1 ships BF16 by default. This is the experiment that asks whether it should not. - -The route differs from System 2's. System 2 goes through TensorRT-Edge-LLM, which consumes a -quantized HF checkpoint and inserts the scaling itself. System 1 is a plain -``torch.onnx.export`` into ``trtexec``, and **TensorRT supports FP8 only through explicit -quantization** -- there is no implicit FP8 calibration the way there is for INT8. So the Q/DQ -nodes have to be in the ONNX, which means a ModelOpt PTQ pass in PyTorch first. - -Calibration comes from ``dump_system1_calib.py`` rather than from random tensors, and that -matters more here than usual: FP8 scale factors are amax-based, and both of these modules -consume tensors whose scale is set by something upstream -- ``z_latents`` by System 2 through -``cond_projector``, the frames by the ResNet normalization. A synthetic draw has the wrong -amax and produces wrong scales, quietly. - -The calibration bundle is read from disk rather than regenerated because capturing it runs -System 2 on every sample; caching it means re-quantizing with a different config costs a -minute instead of the whole pipeline. - -Run under the transformers 4.51 environment (Python 3.10 here), which has both ModelOpt and -a working InternNav:: - - INTERNNAV_PATH=~/InternNav PYTHONPATH=~/InternNav \\ - python quantize_system1.py --calib_path work/system1_calib.pt --out_dir work/onnx_fp8 -""" -import argparse -import os -import subprocess -import sys - -import torch - -_HERE = os.path.dirname(os.path.abspath(__file__)) -sys.path.insert(0, _HERE) - -import internvla_compat # noqa: E402 -from memblock import MemBlock # noqa: E402 - -TRTEXEC = os.environ.get("TRTEXEC", "/usr/src/tensorrt/bin/trtexec") - - -def parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--internvla_ckpt", - default=os.path.expanduser( - os.environ.get("INTERNVLA_CKPT", - "~/InternNav/checkpoints/InternVLA-N1-DualVLN"))) - p.add_argument("--calib_path", required=True, - help="Bundle from dump_system1_calib.py") - p.add_argument("--out_dir", required=True) - p.add_argument("--components", default="traj_dit,memory", - help="Comma-separated subset of traj_dit,memory") - p.add_argument("--zlen", type=int, default=36, - help="Optimization profile z_latents length (32 memory tokens + 4 TRAJ)") - p.add_argument("--exclude_memory", default="nn.Conv2d", - help="Comma-separated exclusions for the memory block: 'nn.X' matches a " - "module class, anything else is a name glob. The default leaves " - "Conv2d alone because the legacy ONNX exporter cannot infer a " - "convolution kernel shape through Q/DQ and dies with 'convolution " - "for kernel of unknown shape'. In DepthAnythingV2 that is the " - "patch_embed projection -- one layer of a ViT, so almost no compute " - "is left behind.") - p.add_argument("--skip_build", action="store_true", - help="Export ONNX only, do not call trtexec") - p.add_argument("--device", default="cuda") - return p.parse_args() - - -def fp8_config(exclude: str = ""): - """FP8_DEFAULT_CFG with name patterns disabled.""" - import copy - - import modelopt.torch.quantization as mtq - - # In ModelOpt 0.44 quant_cfg is an ordered *list* of rules, not a dict, and later rules - # win -- so exclusions go on the end. Entries starting with "nn." match by module class - # (the form the built-in BatchNorm exclusions use); anything else matches by name glob. - cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG) - for pat in (e.strip() for e in exclude.split(",")): - if not pat: - continue - rule = {"quantizer_name": "*", "enable": False} - if pat.startswith("nn."): - rule["parent_class"] = pat - else: - rule["quantizer_name"] = pat - cfg["quant_cfg"].append(rule) - return cfg - - -def count_qdq(onnx_path: str) -> tuple[int, int]: - """Q/DQ node counts. Zero means the PTQ pass did not survive the export, and the - engine below would silently be BF16 with extra steps. - - ModelOpt emits TensorRT's own ``trt::TRT_FP8QuantizeLinear`` for FP8 rather than the - standard ONNX ``QuantizeLinear``, so counting only the latter reports a correctly - quantized graph as unquantized.""" - import onnx - - model = onnx.load(onnx_path, load_external_data=False) - ops = [n.op_type for n in model.graph.node] - q = ops.count("QuantizeLinear") + ops.count("TRT_FP8QuantizeLinear") - dq = ops.count("DequantizeLinear") + ops.count("TRT_FP8DequantizeLinear") - return q, dq - - -def build(onnx_path: str, engine_path: str, shapes: list[str]) -> bool: - cmd = [TRTEXEC, f"--onnx={onnx_path}", f"--saveEngine={engine_path}", "--fp8", "--bf16"] - cmd += shapes - print(" " + " ".join(cmd)) - r = subprocess.run(cmd, capture_output=True, text=True) - if not os.path.exists(engine_path): - tail = "\n ".join((r.stdout + r.stderr).strip().splitlines()[-12:]) - print(f" BUILD FAILED\n {tail}") - return False - print(f" engine {os.path.getsize(engine_path) / 1e6:.1f} MB") - return True - - -def quantize_traj_dit(inner, calib, args) -> str: - import modelopt.torch.quantization as mtq - - dit = inner.traj_dit.eval() - batches = calib["dit_batches"] - dev = args.device - - def forward_loop(m): - for x, ts, z in batches: - with torch.no_grad(): - m(x=x.to(dev, torch.bfloat16), timestep=ts.to(dev), - z_latents=z.to(dev, torch.bfloat16)) - - print(f"[traj_dit] PTQ FP8 over {len(batches)} real batches") - dit = mtq.quantize(dit, fp8_config(), forward_loop=forward_loop) - - # Export in FP32: the RoPE freqs_cis are float32 and a mixed-precision trace fails. - # Precision is re-established by the Q/DQ nodes plus trtexec --bf16. - dit = dit.to(torch.float32).eval() - x, ts, z = batches[0] - x, ts, z = x.to(dev).float(), ts.to(dev), z.to(dev).float() - - class Wrap(torch.nn.Module): - def __init__(self, d): - super().__init__() - self.d = d - - def forward(self, x, timestep, z_latents): - return self.d(x=x, timestep=timestep, z_latents=z_latents) - - w = Wrap(dit).eval() - onnx_path = os.path.join(args.out_dir, "system1_traj_dit_fp8.onnx") - dyn = {"x": {0: "batch"}, "timestep": {0: "batch"}, - "z_latents": {0: "batch", 1: "zlen"}, "output": {0: "batch"}} - print(f"[traj_dit] Export ONNX -> {onnx_path}") - with torch.inference_mode(): - torch.onnx.export(w, (x, ts, z), onnx_path, - input_names=["x", "timestep", "z_latents"], - output_names=["output"], opset_version=19, - do_constant_folding=True, export_params=True, - dynamic_axes=dyn, dynamo=False) - q, dq = count_qdq(onnx_path) - print(f" {os.path.getsize(onnx_path) / 1e6:.1f} MB, " - f"{q} QuantizeLinear / {dq} DequantizeLinear") - if q == 0: - print(" [WARN] no Q/DQ in the graph -- the engine will not be FP8") - - if args.skip_build: - return onnx_path - B, WP, DIM = x.shape[0], x.shape[1], x.shape[2] - zdim = z.shape[-1] - engine = os.path.join(args.out_dir, "system1_traj_dit_fp8.engine") - print("[traj_dit] Build engine") - build(onnx_path, engine, [ - f"--minShapes=x:{B}x{WP}x{DIM},timestep:{B},z_latents:{B}x4x{zdim}", - f"--optShapes=x:{B}x{WP}x{DIM},timestep:{B},z_latents:{B}x{args.zlen}x{zdim}", - f"--maxShapes=x:{B}x{WP}x{DIM},timestep:{B},z_latents:{B}x64x{zdim}"]) - return onnx_path - - -def quantize_memory(model, inner, calib, args) -> str: - import modelopt.torch.quantization as mtq - - dev = args.device - block = MemBlock(inner.rgb_model, inner.memory_encoder, inner.rgb_resampler).eval() - batches = calib["image_batches"] - - def forward_loop(m): - for imgs in batches: - with torch.no_grad(): - m(imgs.to(dev, torch.bfloat16)) - - print(f"[memory] PTQ FP8 over {len(batches)} real frame batches" - + (f", excluding {args.exclude_memory}" if args.exclude_memory else "")) - block = mtq.quantize(block, fp8_config(args.exclude_memory), forward_loop=forward_loop) - block = block.to(torch.float32).eval() - - imgs = batches[0].to(dev).float() - onnx_path = os.path.join(args.out_dir, "system1_memory_fp8.onnx") - print(f"[memory] Export ONNX -> {onnx_path}") - with torch.inference_mode(): - torch.onnx.export(block, (imgs,), onnx_path, input_names=["images"], - output_names=["memory_tokens"], opset_version=19, - do_constant_folding=True, export_params=True, - dynamic_axes={"images": {0: "frames"}}, dynamo=False) - q, dq = count_qdq(onnx_path) - print(f" {os.path.getsize(onnx_path) / 1e6:.1f} MB, " - f"{q} QuantizeLinear / {dq} DequantizeLinear") - if q == 0: - print(" [WARN] no Q/DQ in the graph -- the engine will not be FP8") - - if args.skip_build: - return onnx_path - T, C, H, W = imgs.shape - engine = os.path.join(args.out_dir, "system1_memory_fp8.engine") - print("[memory] Build engine") - build(onnx_path, engine, [ - f"--minShapes=images:1x{C}x{H}x{W}", - f"--optShapes=images:{T}x{C}x{H}x{W}", - f"--maxShapes=images:8x{C}x{H}x{W}"]) - return onnx_path - - -def main() -> int: - args = parse_args() - wanted = {c.strip() for c in args.components.split(",") if c.strip()} - # The memory block's QFormer takes PyTorch's fused MHA fast path, and - # aten::_transformer_encoder_layer_fwd has no ONNX symbolic. Same toggle the BF16 - # export uses. - try: - torch.backends.mha.set_fastpath_enabled(False) - except Exception as exc: # pragma: no cover - print(f" (mha fastpath toggle unavailable: {exc})") - os.makedirs(args.out_dir, exist_ok=True) - internvla_compat.apply_all(need_system1=True, allow_missing_depth=False) - - from internnav.model.basemodel.internvla_n1.internvla_n1 import ( - InternVLAN1ForCausalLM, InternVLAN1ModelConfig) - - calib = torch.load(args.calib_path, map_location="cpu", weights_only=False) - print(f"Loading {args.internvla_ckpt}") - config = InternVLAN1ModelConfig.from_pretrained(args.internvla_ckpt) - model = InternVLAN1ForCausalLM.from_pretrained( - args.internvla_ckpt, config=config, torch_dtype=torch.bfloat16, - attn_implementation="sdpa", low_cpu_mem_usage=True).to(args.device).eval() - inner = model.get_model() - - if "traj_dit" in wanted: - quantize_traj_dit(inner, calib, args) - if "memory" in wanted: - quantize_memory(model, inner, calib, args) - print("\nDone. Verify with verify/compare_system1_engines.py --engine_dir " - f"{args.out_dir}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/scripts/03_export_build_system2.sh b/recipes/internvla-n1-dualvln/trt-edgellm/scripts/03_export_build_system2.sh deleted file mode 100644 index b0d6d78..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/scripts/03_export_build_system2.sh +++ /dev/null @@ -1,145 +0,0 @@ -#!/usr/bin/env bash -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -# -# Export a System-2 checkpoint to ONNX and build the TensorRT-Edge-LLM engines. -# -# Works on both a quantized checkpoint and the unquantized one; pass --no_quantization -# for the latter to get the FP16 fidelity reference. -# -# Two details here are load-bearing and must not be "cleaned up": -# -# * CausalLM.emit_hidden_states = True before the export. The System 2 -> System 1 -# bridge reads the last-layer hidden states of the trajectory tokens; without this the -# engine emits logits only and the bridge cannot be evaluated at all. -# -# * __LUNOWUD=-peep:fc_h_fusion=off on the engine build. TensorRT 10.13 miscompiles -# Myelin's horizontal fusion of the gate/up projections on sm_110 at batch 1, and an -# FP16 engine built without this emits fluent-looking gibberish. TensorRT-Edge-LLM -# already disables the fusion, but only for TensorRT >= 10.15, so 10.13 falls through -# the gap. FP8 happens to dodge the same bug because its Q/DQ nodes break the fusion -# pattern, which is why FP8 once looked mandatory -- it is not. -# -# The visual encoder is sized for multi-image VLN prompts (~1764 image tokens across 9-10 -# frames). The single-image demo default of 512 max image tokens cannot hold one. - -set -euo pipefail - -MODEL_PATH="" -ENGINE_DIR="" -ONNX_DIR="" -TRT_EDGELLM_DIR="${TRT_EDGELLM_DIR:-$HOME/modelopt/TensorRT-Edge-LLM}" -NO_QUANTIZATION=0 -SKIP_VISUAL=0 -MAX_BATCH_SIZE="${MAX_BATCH_SIZE:-1}" -MAX_INPUT_LEN="${MAX_INPUT_LEN:-3072}" -MAX_KV_CACHE="${MAX_KV_CACHE:-4096}" -LUNOWUD_WAR="${LUNOWUD_WAR:--peep:fc_h_fusion=off}" - -while [[ $# -gt 0 ]]; do - case "$1" in - --model_path) MODEL_PATH="$2"; shift 2 ;; - --engine_dir) ENGINE_DIR="$2"; shift 2 ;; - --onnx_dir) ONNX_DIR="$2"; shift 2 ;; - --trt_edgellm_dir) TRT_EDGELLM_DIR="$2"; shift 2 ;; - --no_quantization) NO_QUANTIZATION=1; shift ;; - --skip_visual) SKIP_VISUAL=1; shift ;; - --max_batch_size) MAX_BATCH_SIZE="$2"; shift 2 ;; - --max_input_len) MAX_INPUT_LEN="$2"; shift 2 ;; - --max_kv_cache) MAX_KV_CACHE="$2"; shift 2 ;; - -h|--help) - sed -n '3,25p' "$0" | sed 's/^# \{0,1\}//' - exit 0 ;; - *) echo "[ERROR] unknown argument: $1" >&2; exit 1 ;; - esac -done - -[[ -n "$MODEL_PATH" ]] || { echo "[ERROR] --model_path is required" >&2; exit 1; } -[[ -n "$ENGINE_DIR" ]] || { echo "[ERROR] --engine_dir is required" >&2; exit 1; } -[[ -d "$MODEL_PATH" ]] || { echo "[ERROR] model path not found: $MODEL_PATH" >&2; exit 1; } - -ONNX_DIR="${ONNX_DIR:-${ENGINE_DIR%/}_onnx}" -PLUGIN="$TRT_EDGELLM_DIR/build/libNvInfer_edgellm_plugin.so" -LLM_BUILD="$TRT_EDGELLM_DIR/build/examples/llm/llm_build" -VISUAL_BUILD="$TRT_EDGELLM_DIR/build/examples/multimodal/visual_build" - -for f in "$PLUGIN" "$LLM_BUILD"; do - [[ -f "$f" ]] || { echo "[ERROR] not found: $f" >&2 - echo " set TRT_EDGELLM_DIR or build TensorRT-Edge-LLM first" >&2 - exit 1; } -done - -export EDGELLM_PLUGIN_PATH="$PLUGIN" -export __LUNOWUD="${__LUNOWUD:+$__LUNOWUD }$LUNOWUD_WAR" - -echo "==============================================================" -echo " InternVLA-N1 System 2 — ONNX export and engine build" -echo "==============================================================" -echo " Model path: $MODEL_PATH" -echo " Quantized: $([[ $NO_QUANTIZATION -eq 1 ]] && echo 'no (FP16 reference)' || echo yes)" -echo " ONNX dir: $ONNX_DIR" -echo " Engine dir: $ENGINE_DIR" -echo " TensorRT-Edge: $TRT_EDGELLM_DIR" -echo " maxBatchSize: $MAX_BATCH_SIZE" -echo " maxInputLen: $MAX_INPUT_LEN" -echo " maxKVCacheCap: $MAX_KV_CACHE" -echo " __LUNOWUD: $__LUNOWUD" -echo "==============================================================" - -mkdir -p "$ONNX_DIR" "$ENGINE_DIR" - -echo -echo "[1/3] Exporting ONNX (emit_hidden_states=True for the System 1 bridge)..." -python - "$MODEL_PATH" "$ONNX_DIR" <<'PYEOF' -import sys -from tensorrt_edgellm.models.default.modeling_default import CausalLM -# The bridge to System 1 reads the last-layer hidden states, not just logits. -CausalLM.emit_hidden_states = True -from tensorrt_edgellm.scripts.export import main -sys.argv = ["tensorrt-edgellm-export", sys.argv[1], sys.argv[2]] -sys.exit(main()) -PYEOF - -echo -echo "[2/3] Building the LLM engine..." -mkdir -p "$ENGINE_DIR/llm" -"$LLM_BUILD" \ - --onnxDir "$ONNX_DIR/llm" \ - --engineDir "$ENGINE_DIR/llm" \ - --maxBatchSize "$MAX_BATCH_SIZE" \ - --maxInputLen "$MAX_INPUT_LEN" \ - --maxKVCacheCapacity "$MAX_KV_CACHE" - -# llm_bench reads its engine configuration from base_config.json, while llm_build writes -# config.json. Same content, different name; without this copy the benchmark cannot open -# the engine it was just handed. -cp "$ENGINE_DIR/llm/config.json" "$ENGINE_DIR/llm/base_config.json" - -if [[ $SKIP_VISUAL -eq 0 && -d "$ONNX_DIR/visual" ]]; then - echo - echo "[3/3] Building the visual engine..." - [[ -f "$VISUAL_BUILD" ]] || { echo "[ERROR] not found: $VISUAL_BUILD" >&2; exit 1; } - # visual_build appends its own "visual" component directory under --engineDir, so - # pass the parent: giving it $ENGINE_DIR/visual yields $ENGINE_DIR/visual/visual and - # every consumer that expects $ENGINE_DIR/visual/visual.engine then misses it. - # - # A VLN prompt carries 9-10 frames and roughly 1764 image tokens. The single-image - # demo default of 512 cannot hold one. - "$VISUAL_BUILD" \ - --onnxDir "$ONNX_DIR/visual" \ - --engineDir "$ENGINE_DIR" \ - --minImageTokens 4 \ - --maxImageTokens 4096 \ - --maxImageTokensPerImage 1024 -else - echo - echo "[3/3] Skipping the visual engine." -fi - -echo -echo "==============================================================" -echo " Build complete" -echo "==============================================================" -du -sh "$ENGINE_DIR"/* 2>/dev/null | sed 's/^/ /' -echo " Engines: $ENGINE_DIR" -echo " ONNX kept at $ONNX_DIR (safe to delete once the engines verify)" diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/scripts/04_export_system1.sh b/recipes/internvla-n1-dualvln/trt-edgellm/scripts/04_export_system1.sh deleted file mode 100755 index 1b9e1bf..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/scripts/04_export_system1.sh +++ /dev/null @@ -1,64 +0,0 @@ -#!/usr/bin/env bash -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -# -# Export System 1 -- the NextDiT diffusion head and the memory block -- to BF16 engines. -# -# Upstream InternNav ships no ONNX or TensorRT path at all, so this is entirely the -# recipe's. Both exporters need InternNav importable, unlike everything under System 2. -# -# System 1 stays BF16 on purpose. FP8 was measured (trt-edgellm/quantize_system1.py) and -# costs 6x the waypoint deviation to save 0.7% of deployed weights and 1.7% of a planning -# step -- see the README. -# -# The two exporters write into $WORK_DIR/onnx, which is also where verify_system1.py looks; -# --engine_dir is linked to those files rather than holding copies, so both paths stay valid -# without a second 176 MB on disk. - -set -euo pipefail - -HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -ROOT="$(dirname "$HERE")" -INTERNVLA_CKPT="${INTERNVLA_CKPT:-$HOME/InternNav/checkpoints/InternVLA-N1-DualVLN}" -INTERNNAV_PATH="${INTERNNAV_PATH:-$HOME/InternNav}" -WORK_DIR="${WORK_DIR:-$HOME/vln-opt-work}" -ENGINE_DIR="" -COMPONENTS="traj_dit,memory" - -while [[ $# -gt 0 ]]; do - case "$1" in - --internvla_ckpt) INTERNVLA_CKPT="$2"; shift 2 ;; - --internnav_path) INTERNNAV_PATH="$2"; shift 2 ;; - --work_dir) WORK_DIR="$2"; shift 2 ;; - --engine_dir) ENGINE_DIR="$2"; shift 2 ;; - --components) COMPONENTS="$2"; shift 2 ;; - -h|--help) sed -n '2,17p' "$0"; exit 0 ;; - *) echo "[ERROR] unknown argument: $1" >&2; exit 2 ;; - esac -done - -[[ -d "$INTERNNAV_PATH" ]] || { echo "[ERROR] InternNav not found: $INTERNNAV_PATH" >&2; exit 1; } -[[ -d "$INTERNVLA_CKPT" ]] || { echo "[ERROR] checkpoint not found: $INTERNVLA_CKPT" >&2; exit 1; } - -export INTERNNAV_PATH INTERNVLA_CKPT WORK_DIR -# The exporters import both InternNav and the recipe's own modules. -export PYTHONPATH="$INTERNNAV_PATH:$ROOT${PYTHONPATH:+:$PYTHONPATH}" - -if [[ ",$COMPONENTS," == *",traj_dit,"* ]]; then - echo "== traj_dit ==" - python -u "$ROOT/export_traj_dit.py" -fi -if [[ ",$COMPONENTS," == *",memory,"* ]]; then - echo "== memory block ==" - python -u "$ROOT/export_memory_block.py" -fi - -if [[ -n "$ENGINE_DIR" ]]; then - mkdir -p "$ENGINE_DIR" - for f in "$WORK_DIR"/onnx/system1_*.engine; do - [[ -e "$f" ]] || continue - ln -sfn "$f" "$ENGINE_DIR/$(basename "$f")" - done - echo "Engines linked into: $ENGINE_DIR" -fi -ls -lh "$WORK_DIR"/onnx/system1_*.engine 2>/dev/null || true diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/scripts/05_verify.sh b/recipes/internvla-n1-dualvln/trt-edgellm/scripts/05_verify.sh deleted file mode 100755 index 43e1054..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/scripts/05_verify.sh +++ /dev/null @@ -1,83 +0,0 @@ -#!/usr/bin/env bash -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -# -# Run the acceptance gates. -# -# Default (System 2): the z_latents bridge -- the last-layer hidden states of the 4 TRAJ -# tokens through the host-side norm and cond_projector, engine against a PyTorch reference. -# This is the number that decides whether a scheme ships. Text fluency is not sufficient -# evidence: NVFP4 stays fluent and fails this at 0.931. -# -# --system1: trajectory parity for the diffusion head and memory block. This one runs in TWO -# stages under TWO interpreters, and that is not incidental. InternNav targets transformers -# 4.x while the TensorRT bindings ship for Python 3.12 where transformers is 5.x, so no -# single environment has both. Stage A writes the PyTorch reference's inputs and outputs to -# a .pt; stage B feeds the engines those same tensors. Set PYTHON_PT and PYTHON_TRT to the -# two interpreters -- if they are left at the default the stages run in whatever is active, -# which works only if one environment happens to satisfy both. - -set -euo pipefail - -HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -ROOT="$(dirname "$HERE")" -SYSTEM1=0 -ENGINE_DIR="" -ENGINE_SUFFIX="${ENGINE_SUFFIX:-bf16}" -REPKG_CKPT="${REPKG_CKPT:-$HOME/vln-opt-work/qwen25vl_system2}" -CALIB_DATA_ROOT="${CALIB_DATA_ROOT:-$HOME/vln-opt-work/calib_scenes}" -INTERNVLA_CKPT="${INTERNVLA_CKPT:-$HOME/InternNav/checkpoints/InternVLA-N1-DualVLN}" -INTERNNAV_PATH="${INTERNNAV_PATH:-$HOME/InternNav}" -WORK_DIR="${WORK_DIR:-$HOME/vln-opt-work}" -VLN=0 -BENCH_ITERS="${BENCH_ITERS:-0}" -# The System 2 gate needs transformers and the TensorRT bindings in one interpreter; the -# System 1 gate cannot have both and splits across PYTHON_PT / PYTHON_TRT below. -PYTHON="${PYTHON:-python}" - -while [[ $# -gt 0 ]]; do - case "$1" in - --system1) SYSTEM1=1; shift ;; - --vln) VLN=1; shift ;; - --engine_dir) ENGINE_DIR="$2"; shift 2 ;; - --engine_suffix) ENGINE_SUFFIX="$2"; shift 2 ;; - --repkg_ckpt) REPKG_CKPT="$2"; shift 2 ;; - --calib_data_root) CALIB_DATA_ROOT="$2"; shift 2 ;; - --internvla_ckpt) INTERNVLA_CKPT="$2"; shift 2 ;; - --internnav_path) INTERNNAV_PATH="$2"; shift 2 ;; - --work_dir) WORK_DIR="$2"; shift 2 ;; - --bench_iters) BENCH_ITERS="$2"; shift 2 ;; - -h|--help) sed -n '2,19p' "$0"; exit 0 ;; - *) echo "[ERROR] unknown argument: $1" >&2; exit 2 ;; - esac -done - -export WORK_DIR REPKG_CKPT INTERNNAV_PATH INTERNVLA_CKPT CALIB_DATA_ROOT -export PYTHONPATH="$ROOT${PYTHONPATH:+:$PYTHONPATH}" -export EDGELLM_PLUGIN_PATH="${EDGELLM_PLUGIN_PATH:-$HOME/modelopt/Hung-TRT-Edge-LLM/build/libNvInfer_edgellm_plugin.so}" - -if [[ "$SYSTEM1" -eq 1 ]]; then - : "${ENGINE_DIR:=$WORK_DIR/onnx}" - PYTHON_PT="${PYTHON_PT:-python}" - PYTHON_TRT="${PYTHON_TRT:-python}" - REF="${REFERENCE_PATH:-$WORK_DIR/out/system1_reference.pt}" - - echo "== stage A: PyTorch reference (needs InternNav; $PYTHON_PT) ==" - PYTHONPATH="$INTERNNAV_PATH:$PYTHONPATH" "$PYTHON_PT" -u "$ROOT/verify/dump_system1_reference.py" \ - --internvla_ckpt "$INTERNVLA_CKPT" --output_path "$REF" - - echo "== stage B: engines (needs TensorRT; $PYTHON_TRT) ==" - ARGS=(--reference_path "$REF" --engine_dir "$ENGINE_DIR" --engine_suffix "$ENGINE_SUFFIX") - [[ "$BENCH_ITERS" -gt 0 ]] && ARGS+=(--bench_iters "$BENCH_ITERS") - exec "$PYTHON_TRT" -u "$ROOT/verify/compare_system1_engines.py" "${ARGS[@]}" -fi - -[[ -n "$ENGINE_DIR" ]] || { echo "[ERROR] --engine_dir is required" >&2; exit 2; } -ENGINE_PATH="${ENGINE_PATH:-$ENGINE_DIR/llm/llm.engine}" -[[ -f "$ENGINE_PATH" ]] || { echo "[ERROR] engine not found: $ENGINE_PATH" >&2; exit 1; } -export ENGINE_PATH - -if [[ "$VLN" -eq 1 ]]; then - exec "$PYTHON" -u "$ROOT/verify/verify_latents_vln.py" -fi -exec "$PYTHON" -u "$ROOT/verify/verify_latents.py" diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/scripts/06_measure.sh b/recipes/internvla-n1-dualvln/trt-edgellm/scripts/06_measure.sh deleted file mode 100644 index 7ff4c7f..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/scripts/06_measure.sh +++ /dev/null @@ -1,55 +0,0 @@ -#!/usr/bin/env bash -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -# -# Measure prefill/decode latency and bridge fidelity for every built engine, and write the -# numbers where run_matrix.py can find them. -# -# llm_bench reads base_config.json while llm_build writes config.json, so the copy below is -# not optional -- without it the benchmark cannot open an engine it was just handed. - -set -euo pipefail - -WORK_DIR="${WORK_DIR:-$HOME/vln-opt-work}" -ENGINE_DIR="${ENGINE_DIR:-$WORK_DIR/engines}" -TRT_EDGELLM_DIR="${TRT_EDGELLM_DIR:-$HOME/modelopt/Hung-TRT-Edge-LLM}" -INPUT_LEN="${INPUT_LEN:-1024}" -PAST_KV_LEN="${PAST_KV_LEN:-1024}" -OUT="${OUT:-$WORK_DIR/out/latency.json}" - -LLM_BENCH="$TRT_EDGELLM_DIR/build/examples/llm/llm_bench" -[[ -x "$LLM_BENCH" ]] || { echo "[ERROR] not found: $LLM_BENCH" >&2; exit 1; } -export EDGELLM_PLUGIN_PATH="$TRT_EDGELLM_DIR/build/libNvInfer_edgellm_plugin.so" - -# Map engine directory -> the checkpoint name run_matrix.py keys on. -declare -A CKPT_OF=( - [base_fp16]=qwen25vl_system2 - [s1_fp8]=qwen25vl_s1_fp8 - [s1_nvfp4]=qwen25vl_s1_nvfp4 -) - -mkdir -p "$(dirname "$OUT")" -echo "{" > "$OUT" -first=1 - -for name in "${!CKPT_OF[@]}"; do - engine="$ENGINE_DIR/$name/llm" - [[ -f "$engine/llm.engine" ]] || { echo " [skip] $name: no engine"; continue; } - [[ -f "$engine/base_config.json" ]] || cp "$engine/config.json" "$engine/base_config.json" - - echo " measuring $name ..." - prefill=$("$LLM_BENCH" --engineDir "$engine" --mode prefill --inputLen "$INPUT_LEN" 2>&1 \ - | grep -oE "E2E Time \(actual performance\): [0-9.]+" | tail -1 | grep -oE "[0-9.]+$" || echo "") - decode=$("$LLM_BENCH" --engineDir "$engine" --mode decode --pastKVLen "$PAST_KV_LEN" 2>&1 \ - | grep -oE "E2E Time \(actual performance\): [0-9.]+" | tail -1 | grep -oE "[0-9.]+$" || echo "") - - [[ $first -eq 1 ]] || echo "," >> "$OUT" - first=0 - printf ' "%s": {"prefill_ms": %s, "decode_ms": %s}' \ - "${CKPT_OF[$name]}" "${prefill:-null}" "${decode:-null}" >> "$OUT" - echo " prefill ${prefill:-n/a} ms | decode ${decode:-n/a} ms" -done - -echo "" >> "$OUT" -echo "}" >> "$OUT" -echo "Wrote $OUT" diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/traj_dit_loader.py b/recipes/internvla-n1-dualvln/trt-edgellm/traj_dit_loader.py deleted file mode 100644 index 40674d4..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/traj_dit_loader.py +++ /dev/null @@ -1,109 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Load the InternVLA-N1 System 1 traj_dit (NextDiT) core, and export it to ONNX. - -The traj_dit is exported in FP32 (its RoPE freqs_cis are float32; FP32 export avoids -mixed-precision trace errors). Precision is set to BF16 at engine-build time. The surrounding -host-side ops (action encoder/decoder, cond_projector, pos_encoding, CFG blend, scheduler) -stay in PyTorch. Legacy exporter (dynamo=False, opset 19). - -Export shapes: - input : x [2N, 32, 384] · timestep [2N] int64 · z_latents [2N, Z, 768] - output: noise_pred [2N, 32, 384] -""" -import os -import sys - -import torch - -_LIB = os.path.dirname(os.path.abspath(__file__)) -sys.path.insert(0, _LIB) -sys.path.insert(0, os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav"))) - -MODEL = os.environ.get("INTERNVLA_CKPT", os.path.join( - os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")), - "checkpoints/InternVLA-N1-DualVLN")) -OUT = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), - "onnx/system1_traj_dit.onnx") - - -def load_traj_dit(): - """Return (traj_dit module, LatentEmbSize). The rgb_model (DepthAnything) is stubbed - out since only the traj_dit is needed here.""" - from internvla_compat import patch_gradient_checkpointing, patch_traj_dit_ffn - - patch_gradient_checkpointing() - patch_traj_dit_ffn() # requires diffusers 0.33.1 -> inner_dim 1024 - import internnav.model.basemodel.internvla_n1.internvla_n1_arch as _arch - - class _DepthStub(torch.nn.Module): - def forward(self, *a, **k): - raise RuntimeError("rgb_model stub") - - _arch.build_depthanythingv2 = lambda config: _DepthStub() - - from internnav.model.basemodel.internvla_n1.internvla_n1 import ( - InternVLAN1ForCausalLM, InternVLAN1ModelConfig, - ) - - config = InternVLAN1ModelConfig.from_pretrained(MODEL) - model = InternVLAN1ForCausalLM.from_pretrained( - MODEL, config=config, torch_dtype=torch.float32, low_cpu_mem_usage=True) - model.to("cuda").eval() - return model.get_model().traj_dit, _arch.LatentEmbSize - - -class TrajDiTWrapper(torch.nn.Module): - """Adapt positional args to keyword call (torch.onnx.export passes positionally).""" - - def __init__(self, dit): - super().__init__() - self.dit = dit - - def forward(self, x, timestep, z_latents): - return self.dit(x=x, timestep=timestep, z_latents=z_latents) - - -def main(): - N, WP, DIM = 32, 32, 384 # num_sample_trajs, waypoints, dim - dev = "cuda" - print("[1/3] Load traj_dit (FP32)") - dit, zdim = load_traj_dit() - dit = dit.to(torch.float32).eval() - print(f" z_latents dim (LatentEmbSize) = {zdim}") - - x = torch.randn(2 * N, WP, DIM, dtype=torch.float32, device=dev) - timestep = torch.ones(2 * N, dtype=torch.int64, device=dev) - z_latents = torch.randn(2 * N, 4, zdim, dtype=torch.float32, device=dev) - wrapped = TrajDiTWrapper(dit).eval() - - with torch.no_grad(): - ref = wrapped(x, timestep, z_latents) - print(f"[2/3] Forward OK -> output {tuple(ref.shape)} (expected [{2*N},{WP},{DIM}])") - - os.makedirs(os.path.dirname(OUT), exist_ok=True) - # batch (2N) dynamic for flexible num_sample_trajs; seq (32) & n_query (4) static - dyn = {"x": {0: "batch"}, "timestep": {0: "batch"}, - "z_latents": {0: "batch"}, "output": {0: "batch"}} - print(f"[3/3] Export ONNX (dynamo=False, opset 19) -> {OUT}") - with torch.inference_mode(): - torch.onnx.export( - wrapped, (x, timestep, z_latents), OUT, - input_names=["x", "timestep", "z_latents"], - output_names=["output"], - opset_version=19, do_constant_folding=True, export_params=True, - dynamic_axes=dyn, dynamo=False) - print(f" Done. {OUT} ({os.path.getsize(OUT) / 1e6:.1f} MB)") - - try: - import onnx - onnx.checker.check_model(onnx.load(OUT)) - print(" onnx.checker: valid") - except Exception as e: - print(f" onnx.checker error: {e}") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/trt_torch.py b/recipes/internvla-n1-dualvln/trt-edgellm/trt_torch.py deleted file mode 100644 index fef4848..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/trt_torch.py +++ /dev/null @@ -1,230 +0,0 @@ -# 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. - -# -# 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. - -"""TensorRT Engine wrapper for GR00T inference. - -Loads serialized TRT engines, manages input/output tensor bindings, and -executes inference. Supports dynamic shapes and BF16/FP16/FP32 dtypes. -""" - -import ctypes -import os - -import tensorrt as trt -import torch - - -def torch_type(trt_type): - """Convert TensorRT data type to PyTorch equivalent.""" - mapping = { - trt.float32: torch.float32, - trt.float16: torch.float16, - trt.bfloat16: torch.bfloat16, - trt.int8: torch.int8, - trt.int32: torch.int32, - trt.bool: torch.bool, - trt.uint8: torch.uint8, - trt.int64: torch.int64, - } - if trt_type in mapping: - return mapping[trt_type] - - raise TypeError( - f"Could not resolve TensorRT datatype to an equivalent PyTorch datatype. {trt_type}" - ) - - -class Engine(object): - """TensorRT engine wrapper for loading and executing inference.""" - - def __init__(self, file, plugins=[]): - super().__init__() - - self._closed = False - self.execution_context = None - self.handle = None - - self.logger = trt.Logger(trt.Logger.ERROR) - trt.init_libnvinfer_plugins(self.logger, "") - - self.plugins = [ctypes.CDLL(plugin, ctypes.RTLD_GLOBAL) for plugin in plugins] - self.file = file - self.load(file) - - self.print() - - def close(self): - """Release the execution context, then the engine handle. - - TensorRT requires the context be dropped before the engine. - Idempotent so it composes safely with ``__del__``. - """ - if self._closed: - return - self._closed = True - self.execution_context = None - self.handle = None - - def __del__(self): - # tensorrt / CUDA context may already be gone at shutdown; swallow - # so the traceback doesn't surface as a spurious error on exit. - try: - self.close() - except Exception: - pass - - def print(self): - """Display engine details (inputs/outputs) on rank 0 only.""" - if int(os.getenv("LOCAL_RANK", -1)) not in [0, -1]: - return - - print("============= TRT Engine Detail =============") - print(f"Engine file: {self.file}") - print(f"Inputs: {len(self.in_meta)}") - for ib, item in enumerate(self.in_meta): - tensor_name, shape, dtype = item[:3] - print(f" {ib}. {tensor_name}: {'x'.join(map(str, shape))} [{dtype}]") - - print(f"Outputs: {len(self.out_meta)}") - for ib, item in enumerate(self.out_meta): - tensor_name, shape, dtype = item[:3] - print(f" {ib}. {tensor_name}: {'x'.join(map(str, shape))} [{dtype}]") - print("=============================================") - - def load(self, file): - """Deserialize and load a TensorRT engine from file.""" - runtime = trt.Runtime(self.logger) - - try: - with open(file, "rb") as f: - self.handle = runtime.deserialize_cuda_engine(f.read()) - assert self.handle is not None, ( - f"Failed to deserialize the cuda engine from file: {file}" - ) - - self.execution_context = self.handle.create_execution_context() - self.meta, self.in_meta, self.out_meta = [], [], [] - for tensor_name in self.handle: - shape = self.handle.get_tensor_shape(tensor_name) - dtype = torch_type(self.handle.get_tensor_dtype(tensor_name)) - if self.handle.get_tensor_mode(tensor_name) == trt.TensorIOMode.INPUT: - self.in_meta.append([tensor_name, shape, dtype]) - else: - self.out_meta.append([tensor_name, shape, dtype]) - except BaseException: - # Roll back a half-loaded engine (e.g. create_execution_context - # failed after the handle was set) so it can't leak GPU memory. - self.close() - raise - - def __call__(self, *args, **inputs): - return self.forward(*args, **inputs) - - def dtype_of(self, tensor_name: str) -> torch.dtype: - """Return the expected PyTorch dtype for a named input tensor.""" - for name, _shape, dtype in self.in_meta: - if name == tensor_name: - return dtype - raise KeyError(f"Input tensor '{tensor_name}' not found in engine.") - - def set_runtime_tensor_shape(self, name, shape): - """Set runtime input shape for dynamic dimensions.""" - self.execution_context.set_input_shape(name, shape) - - def forward(self, *args, **kwargs): - """Execute TRT inference with the given input tensors. - - Accepts both positional and keyword arguments. - Returns a dict of output tensors by default, or a list if return_list=True. - """ - return_list = kwargs.pop("return_list", False) - reference_tensors = [] - stream = torch.cuda.current_stream() - - # Process positional arguments - for iarg, x in enumerate(args): - name, shape, dtype = self.in_meta[iarg] - runtime_shape = self.execution_context.get_tensor_shape(name) - assert isinstance(x, torch.Tensor), f"Unsupported tensor type: {type(x)}" - assert runtime_shape == x.shape, f"Invalid input shape: {runtime_shape} != {x.shape}" - assert dtype == x.dtype, ( - f"Invalid tensor dtype, expected dtype is {dtype}, but got {x.dtype}" - ) - assert x.is_cuda, f"Invalid tensor device, expected device is cuda, but got {x.device}" - x = x.cuda().contiguous() - self.execution_context.set_tensor_address(name, x.data_ptr()) - reference_tensors.append(x) - - # Process keyword arguments - for name, shape, dtype in self.in_meta: - if name not in kwargs: - continue - - runtime_shape = self.execution_context.get_tensor_shape(name) - x = kwargs[name] - assert isinstance(x, torch.Tensor), f"Unsupported tensor[{name}] type: {type(x)}" - assert runtime_shape == x.shape, ( - f"Invalid input[{name}] shape: {x.shape}, but the expected shape is: {runtime_shape}" - ) - assert dtype == x.dtype, ( - f"Invalid tensor[{name}] dtype, expected dtype is {dtype}, but got {x.dtype}" - ) - assert x.is_cuda, ( - f"Invalid tensor[{name}] device, expected device is cuda, but got {x.device}" - ) - x = x.cuda().contiguous() - self.execution_context.set_tensor_address(name, x.data_ptr()) - reference_tensors.append(x) - - # Allocate output tensors - for item in self.out_meta: - name = item[0] - runtime_shape = self.execution_context.get_tensor_shape(name) - output_tensor = torch.zeros( - *runtime_shape, dtype=item[2], device=reference_tensors[0].device - ) - self.execution_context.set_tensor_address(name, output_tensor.data_ptr()) - reference_tensors.append(output_tensor) - - # Execute - self.execution_context.execute_async_v3(stream.cuda_stream) - stream.synchronize() - assert len(reference_tensors) == len(self.in_meta) + len(self.out_meta), ( - f"Invalid input tensors. The expected I/O tensors are " - f"{len(self.in_meta) + len(self.out_meta)}, but got {len(reference_tensors)}" - ) - - if return_list: - return [ - reference_tensors[len(self.in_meta) + i] for i, item in enumerate(self.out_meta) - ] - else: - return { - item[0]: reference_tensors[len(self.in_meta) + i] - for i, item in enumerate(self.out_meta) - } diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/compare_system1_engines.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/compare_system1_engines.py deleted file mode 100644 index 77ef552..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/compare_system1_engines.py +++ /dev/null @@ -1,241 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Stage B of the System-1 parity check: run the engines on the reference's own inputs. - -Reads the ``.pt`` written by ``dump_system1_reference.py`` and pushes the *same* tensors -through the memory and traj_dit engines, so the two halves never need to share an -interpreter — which they cannot, since InternNav wants transformers 4.x and the TensorRT -bindings are Python 3.12 only. - -Two comparisons, deliberately separate: - -* ``memory_tokens`` — the memory block alone (DepthAnythingV2 + MemoryEncoder + QFormer). -* the full trajectory — memory block plus the diffusion loop over traj_dit. - -Checking the intermediate first means a mismatch localises to one engine instead of only -appearing at the end. - -Run this under the TensorRT environment (Python 3.12 here):: - - EDGELLM_PLUGIN_PATH=... python verify/compare_system1_engines.py \\ - --reference_path work/system1_reference.pt --engine_dir work/onnx -""" -import argparse -import os -import sys -import time - -import numpy as np -import torch - -_HERE = os.path.dirname(os.path.abspath(__file__)) -sys.path.insert(0, os.path.dirname(_HERE)) -sys.path.append(os.environ.get("SYSTEM_SITE", "/usr/lib/python3.12/dist-packages")) - -from trt_torch import Engine # noqa: E402 - - -def timeit(fn, iters: int, warmup: int = 3) -> float: - """Mean wall-clock milliseconds, synchronized on both sides.""" - for _ in range(warmup): - fn() - torch.cuda.synchronize() - t0 = time.perf_counter() - for _ in range(iters): - fn() - torch.cuda.synchronize() - return (time.perf_counter() - t0) * 1000.0 / iters - - -def cos(a: torch.Tensor, b: torch.Tensor) -> float: - a = a.double().flatten() - b = b.double().flatten() - return float(a @ b / (a.norm() * b.norm())) - - -def out_of(result, key: str) -> torch.Tensor: - if isinstance(result, dict): - return result[key] - return result[0] - - -def parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--reference_path", required=True) - p.add_argument("--engine_dir", required=True, - help="Directory holding the two System-1 engines") - p.add_argument("--engine_suffix", default="bf16", - help="Precision tag in the engine filenames: system1_traj_dit_.engine") - p.add_argument("--guidance_scale", type=float, default=1.0) - p.add_argument("--device", default="cuda") - p.add_argument("--loop_dtype", default="bfloat16", - help="Host-side dtype for the diffusion loop. generate_traj runs it in " - "bfloat16; fp32 here diverges from the reference even with correct " - "engines, because the sampler amplifies the difference.") - p.add_argument("--bench_iters", type=int, default=0, - help="If > 0, also time each engine and the whole trajectory. Run on an " - "idle GPU -- a shared device reads 40-60%% high.") - p.add_argument("--gate", type=float, default=0.99, - help="Minimum trajectory cosine to report PASS") - return p.parse_args() - - -def main() -> int: - args = parse_args() - ref = torch.load(args.reference_path, map_location="cpu", weights_only=False) - dev = args.device - - tag = args.engine_suffix - mem_path = os.path.join(args.engine_dir, f"system1_memory_{tag}.engine") - dit_path = os.path.join(args.engine_dir, f"system1_traj_dit_{tag}.engine") - for path in (mem_path, dit_path): - if not os.path.isfile(path): - print(f"[ERROR] engine not found: {path}") - return 1 - - print("[1/3] memory block") - mem = Engine(mem_path) - # The memory engine was built from MemBlock, which takes [T, C, H, W]. - images = ref["images_chw"].to(dev).float().contiguous() - # The FP8 memory engine is built with a dynamic frame count; the BF16 one is static. - # Setting the shape is a no-op for the static engine and required for the dynamic one. - try: - mem.set_runtime_tensor_shape("images", tuple(images.shape)) - except Exception: - pass - tokens = out_of(mem(images=images), "memory_tokens").float().cpu() - timing = {} - if args.bench_iters: - timing["memory_ms"] = timeit(lambda: mem(images=images), args.bench_iters) - mem.close() - mem_cos = cos(ref["memory_tokens"], tokens) - print(f" memory_tokens cosine vs PyTorch: {mem_cos:.6f}") - - print("[2/3] traj_dit, single forward on the reference's own tensors") - dit_probe = Engine(dit_path) - st = ref["dit_step"] - dit_probe.set_runtime_tensor_shape("z_latents", tuple(st["z_latents"].shape)) - one = out_of(dit_probe(x=st["x"].to(dev).contiguous(), - timestep=st["timestep"].to(dev).to(torch.int64).contiguous(), - z_latents=st["z_latents"].to(dev).contiguous()), "output") - if args.bench_iters: - timing["traj_dit_step_ms"] = timeit( - lambda: dit_probe(x=st["x"].to(dev).contiguous(), - timestep=st["timestep"].to(dev).to(torch.int64).contiguous(), - z_latents=st["z_latents"].to(dev).contiguous()), - args.bench_iters) - dit_probe.close() - step_cos = cos(st["output"], one.float().cpu()) - print(f" traj_dit single-step cosine: {step_cos:.6f}") - - print("[3/4] diffusion loop over traj_dit") - from diffusers.schedulers import FlowMatchEulerDiscreteScheduler - - steps = int(ref["num_inference_steps"]) - n_traj = int(ref["num_sample_trajs"]) - dit = Engine(dit_path) - - enc_w = ref["action_encoder"]["weight"].to(dev) - enc_b = ref["action_encoder"]["bias"].to(dev) - dec_w = ref["action_decoder"]["weight"].to(dev) - dec_b = ref["action_decoder"]["bias"].to(dev) - - # cond_projector is the System 2 -> System 1 bridge: Linear, GELU, Linear mapping the - # 3584-wide hidden state down to the 768 traj_dit consumes. It stays on the host, so - # apply it here from the weights the reference stage saved. - cp = {k: v.to(dev) for k, v in ref["cond_projector"].items()} - z = ref["z"].to(dev) - z = torch.nn.functional.linear(z, cp["0.weight"], cp.get("0.bias")) - z = torch.nn.functional.gelu(z, approximate="tanh") - z = torch.nn.functional.linear(z, cp["2.weight"], cp.get("2.bias")) - # generate_traj runs classifier-free guidance: the conditioning is [null, real] and the - # latents are duplicated, so the engine's batch is 2 * num_sample_trajs. That is why the - # traj_dit engine was built at 64 rather than 32. - hidden = torch.cat([tokens.to(dev), z], dim=1) # [1, 36, 768] - cond = torch.cat([torch.zeros_like(hidden), hidden], 0) # [2, 36, 768] - cond = cond.repeat_interleave(n_traj, dim=0).contiguous() # [2*ns, 36, 768] - pos_embed = ref["pos_embed"].to(dev) - - # Does the loop's own conditioning match what the reference actually fed traj_dit? If - # not, a trajectory mismatch is this reimplementation's, not the engine's. - ref_cond = ref["dit_step"]["z_latents"] - print(f" z_latents vs the reference's own: {cos(ref_cond, cond.float().cpu()):.6f}") - - scheduler = FlowMatchEulerDiscreteScheduler() - scheduler.set_timesteps(steps, sigmas=np.linspace(1.0, 1 / steps, steps)) - - # The reference's own starting noise, captured at stage A. Redrawing it here would give - # a different valid trajectory and look like an engine failure. - latents = ref["init_latents"].to(dev) - dt = getattr(torch, args.loop_dtype) - latents = latents.to(dt) - cond, pos_embed = cond.to(dt), pos_embed.to(dt) - enc_w, enc_b, dec_w, dec_b = (t.to(dt) for t in (enc_w, enc_b, dec_w, dec_b)) - dit.set_runtime_tensor_shape("z_latents", tuple(cond.shape)) - - def run_loop(latents): - """One full sampling run. Mirrors generate_traj line for line.""" - # The scheduler carries step_index across calls, so benchmarking a second run - # walks off the end of its sigma table. Reset it per run. - scheduler.set_timesteps(steps, sigmas=np.linspace(1.0, 1 / steps, steps)) - for t in scheduler.timesteps: - feats = torch.nn.functional.linear(latents, enc_w, enc_b) + pos_embed - feats = feats.repeat(2, 1, 1) - if hasattr(scheduler, "scale_model_input"): - feats = scheduler.scale_model_input(feats, t) - ts = t.to(dev).expand(feats.shape[0]).to(torch.int64).contiguous() - pred = out_of(dit(x=feats.float().contiguous(), timestep=ts, - z_latents=cond.float().contiguous()), "output").to(dt) - pred = torch.nn.functional.linear(pred, dec_w, dec_b) - uncond, condit = pred.chunk(2) - pred = uncond + args.guidance_scale * (condit - uncond) - latents = scheduler.step(pred, t, latents).prev_sample - return latents - - latents = run_loop(latents) - if args.bench_iters: - start = ref["init_latents"].to(dev).to(dt) - timing["trajectory_ms"] = timeit(lambda: run_loop(start), args.bench_iters) - dit.close() - traj = latents.float().cpu() - - print("[4/4] compare") - ref_traj = ref["trajectory"] - if traj.shape != ref_traj.shape: - print(f" shape mismatch: engine {tuple(traj.shape)} vs " - f"reference {tuple(ref_traj.shape)}") - print(" The diffusion loop here reimplements generate_traj; if the shapes") - print(" disagree, the reimplementation is wrong, not the engines.") - return 1 - - traj_cos = cos(ref_traj, traj) - l2 = float((traj - ref_traj).norm() / ref_traj.norm()) - # Cosine says how aligned the trajectories are; it does not say whether a robot would - # end up somewhere else. Per-waypoint Euclidean deviation is in the trajectory's own - # units (metres) and is the number to judge a scheme on. - dev = (traj - ref_traj).norm(dim=-1) - reach = ref_traj.norm(dim=-1).mean() - print("=" * 58) - print(f" memory_tokens cosine : {mem_cos:.6f}") - print(f" traj_dit single step : {step_cos:.6f}") - print(f" trajectory cosine : {traj_cos:.6f}") - print(f" trajectory rel-L2 : {l2:.4f}") - print(f" waypoint deviation : mean {dev.mean():.4f} / median " - f"{dev.median():.4f} / p95 {dev.flatten().quantile(0.95):.4f} " - f"(reference waypoint reach {reach:.4f})") - ok = traj_cos >= args.gate - if timing: - print("\n --- latency, mean over " - f"{args.bench_iters} iterations (idle GPU assumed) ---") - print(f" memory block engine : {timing['memory_ms']:.2f} ms") - print(f" traj_dit, one step : {timing['traj_dit_step_ms']:.2f} ms") - print(f" full trajectory : {timing['trajectory_ms']:.2f} ms " - f"({steps} steps, {n_traj} samples)") - print(f" {'PASS' if ok else 'BELOW GATE'} (gate {args.gate})") - return 0 if ok else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/dump_system1_reference.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/dump_system1_reference.py deleted file mode 100644 index 87506e4..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/dump_system1_reference.py +++ /dev/null @@ -1,183 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Stage A of the System-1 parity check: capture the PyTorch reference to disk. - -The check cannot run in one interpreter. InternNav targets transformers 4.x, while the -TensorRT Python bindings ship for Python 3.12 where transformers is 5.x — under which -InternNav fails on `config.hidden_size`, then on `apply_chunking_to_forward`, and so on -without a natural end. - -Splitting it across the two environments avoids that entirely: this stage runs the PyTorch -reference where InternNav works and writes its inputs and outputs to a `.pt` file; -`compare_system1_engines.py` reads that file where TensorRT works. Nothing needs both at -once, and the comparison stays exact because the *same* tensors cross the boundary — the -engines are fed the reference's own inputs rather than regenerated ones. - -Run this under the transformers 4.51 environment (Python 3.10 here):: - - INTERNNAV_PATH=~/InternNav INTERNVLA_CKPT=~/InternNav/checkpoints/InternVLA-N1-DualVLN \\ - python verify/dump_system1_reference.py --output_path work/system1_reference.pt -""" -import argparse -import os -import sys - -import numpy as np -import torch - -_HERE = os.path.dirname(os.path.abspath(__file__)) -sys.path.insert(0, os.path.dirname(_HERE)) - -import internvla_compat # noqa: E402 - -SEED = 12345 - - -def parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--internvla_ckpt", - default=os.path.expanduser( - os.environ.get("INTERNVLA_CKPT", - "~/InternNav/checkpoints/InternVLA-N1-DualVLN"))) - p.add_argument("--output_path", required=True) - p.add_argument("--num_sample_trajs", type=int, default=32) - p.add_argument("--num_inference_steps", type=int, default=10) - p.add_argument("--num_frames", type=int, default=2) - p.add_argument("--device", default="cuda") - return p.parse_args() - - -def main() -> int: - args = parse_args() - internvla_compat.apply_all(need_system1=True, allow_missing_depth=False) - - from internnav.model.basemodel.internvla_n1.internvla_n1 import ( - InternVLAN1ForCausalLM, InternVLAN1ModelConfig) - - print(f"[1/3] Loading {args.internvla_ckpt}") - config = InternVLAN1ModelConfig.from_pretrained(args.internvla_ckpt) - model = InternVLAN1ForCausalLM.from_pretrained( - args.internvla_ckpt, config=config, torch_dtype=torch.bfloat16, - attn_implementation="sdpa", low_cpu_mem_usage=True).to(args.device).eval() - - print("[2/3] Building fixed inputs") - torch.manual_seed(SEED) - np.random.seed(SEED) - inner = model.get_model() - z_dim = inner.latent_queries.shape[-1] - # The bridge output System 1 consumes. Random but seeded, so the engine stage sees - # exactly these numbers rather than its own draw. - z = torch.randn(1, 4, z_dim, dtype=torch.bfloat16, device=args.device) - # generate_traj permutes (0, 1, 4, 2, 3), so it wants [B, T, H, W, C] and produces - # [B, T, C, H, W] internally. The memory engine, built from MemBlock, takes the - # already-permuted [T, C, H, W] -- both forms are saved so stage B feeds each the - # layout it expects. - images_dp = torch.randn(1, args.num_frames, 224, 224, 3, - dtype=torch.bfloat16, device=args.device) - - # generate_traj draws its starting noise with randn_tensor part-way through the - # function, after forwards that may or may not touch the global RNG. Reseeding in stage B - # therefore does not reproduce it -- and a different draw yields a different but equally - # valid trajectory, which reads as a broken engine (cosine ~0.3). Capture the actual - # tensor instead of trying to redraw it. - from internnav.model.basemodel.internvla_n1 import internvla_n1 as _n1 - _real_randn = _n1.randn_tensor - captured = {} - - def _capture(*a, **kw): - out = _real_randn(*a, **kw) - captured.setdefault("latents", out.detach().float().cpu()) - return out - - _n1.randn_tensor = _capture - - # Capture one real traj_dit call, inputs and output together. The diffusion loop in - # stage B is a reimplementation, so a trajectory mismatch on its own cannot tell a bad - # engine from a bad reimplementation. A single forward on the module's own tensors can. - _dit = inner.traj_dit - _real_fwd = _dit.forward - step = {} - - def _tap(x, timestep, z_latents, *a, **kw): - out = _real_fwd(x, timestep, z_latents, *a, **kw) - step.setdefault("x", x.detach().float().cpu()) - step.setdefault("timestep", timestep.detach().cpu()) - step.setdefault("z_latents", z_latents.detach().float().cpu()) - step.setdefault("output", out.detach().float().cpu()) - return out - - _dit.forward = _tap - - print("[3/3] Reference generate_traj") - torch.manual_seed(SEED) - np.random.seed(SEED) - with torch.no_grad(): - traj = model.generate_traj(z, images_dp, - num_sample_trajs=args.num_sample_trajs, - num_inference_steps=args.num_inference_steps) - _n1.randn_tensor = _real_randn - _dit.forward = _real_fwd - traj = traj.float().cpu() - if "latents" not in captured: - print("[ERROR] randn_tensor was never called -- the capture hook missed the draw") - return 1 - print(f" trajectory {tuple(traj.shape)}") - - # The intermediate the engines replace, so a mismatch can be localised to the memory - # block or to the diffusion head rather than only showing up at the end. - with torch.no_grad(): - # generate_traj normalizes with the ResNet statistics before the memory block, and - # MemBlock -- the module that was exported to ONNX -- does not, so it expects the - # already-normalized tensor. Feeding raw pixels here makes memory_tokens disagree - # with what generate_traj actually used (measured 0.315) while both halves look - # internally consistent. - chw = images_dp.permute(0, 1, 4, 2, 3) - # The statistics are fp32 buffers, so the division promotes; generate_traj casts - # back with .to(dtype) before rgb_model and this must match. - images_chw = ((chw - model._resnet_mean) / model._resnet_std) - images_chw = images_chw.flatten(0, 1).to(torch.bfloat16) - # Use MemBlock itself rather than re-deriving the sequence: it is exactly what was - # exported to ONNX, so any mismatch downstream is the engine's, not a difference in - # how the reference was computed. - from memblock import MemBlock - block = MemBlock(inner.rgb_model, inner.memory_encoder, - inner.rgb_resampler).to(args.device).eval() - memory_tokens = block(images_chw).float().cpu() - # SinusoidalPositionalEncoding stays on the host, and generate_traj adds it to the - # action features every diffusion step. Dump the evaluated tensor rather than its - # weights so stage B does not have to reimplement the encoding. - wp = traj.shape[-2] - pos_ids = torch.arange(wp, device=args.device).reshape(1, -1) - pos_embed = inner.pos_encoding(pos_ids).float().cpu() - - payload = { - "seed": SEED, - "z": z.float().cpu(), - "images_dp": images_dp.float().cpu(), # [B, T, H, W, C], generate_traj form - "images_chw": images_chw.float().cpu(), # [T, C, H, W], memory-engine form - "memory_tokens": memory_tokens, - "pos_embed": pos_embed, - "init_latents": captured["latents"], - "dit_step": step, - "trajectory": traj, - "num_sample_trajs": args.num_sample_trajs, - "num_inference_steps": args.num_inference_steps, - "cond_projector": {k: v.float().cpu() - for k, v in inner.cond_projector.state_dict().items()}, - "action_encoder": {k: v.float().cpu() - for k, v in inner.action_encoder.state_dict().items()}, - "action_decoder": {k: v.float().cpu() - for k, v in inner.action_decoder.state_dict().items()}, - } - os.makedirs(os.path.dirname(os.path.abspath(args.output_path)), exist_ok=True) - torch.save(payload, args.output_path) - size_mb = os.path.getsize(args.output_path) / 1e6 - print(f"\nWrote {args.output_path} ({size_mb:.1f} MB)") - print("Now run verify/compare_system1_engines.py under the TensorRT environment.") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_accuracy.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_accuracy.py deleted file mode 100644 index 2078e2f..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_accuracy.py +++ /dev/null @@ -1,161 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Verify System 2 accuracy against the documented reference agent -(InternVLAN1AsyncAgent) run verbatim on sample data: branch decision + pixel goal, -compared frame-by-frame for both PyTorch and the FP8 TensorRT engine. -""" -import os -import sys -import json -import subprocess -import glob -_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, _R) -sys.path.insert(0, os.path.join(_R, "lib")) -import numpy as np # noqa: E402 -import torch # noqa: E402 -from PIL import Image # noqa: E402 -from engine_runner import ENGINE # noqa: E402 - -ACTIVE = os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) -CKPT = os.path.join(ACTIVE, "checkpoints/InternVLA-N1-DualVLN") -TRT = os.environ.get("TRT_EDGE_LLM", os.path.expanduser("~/TensorRT-Edge-LLM")) -ENG_LLM = os.path.dirname(ENGINE) -ENG_VIS = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "engines/system2_visual") -SCRATCH = os.path.expanduser(os.environ.get('VLN_OPT_OUT', '~/vln-opt-work/out')) -# camera_intrinsic matches the demo (inference_only_demo cell 15) -INTR = np.array([[386.5, 0, 328.9, 0], [0, 386.5, 244, 0], [0, 0, 1, 0], [0, 0, 0, 1]]) - - -class Args: - device = "cuda:0" - model_path = CKPT - model_path_original = CKPT - resize_w = 384 - resize_h = 384 - num_history = 8 - plan_step_gap = 4 - - -def main(): - if ACTIVE not in sys.path: - sys.path.insert(0, ACTIVE) - from internvla_compat import apply_all - apply_all(need_system1=True, allow_missing_depth=True) - import importlib.util - _s = importlib.util.spec_from_file_location( - "iar", os.path.join(ACTIVE, "internnav/agent/internvla_n1_agent_realworld.py")) - _m = importlib.util.module_from_spec(_s) - _s.loader.exec_module(_m) - Agent = _m.InternVLAN1AsyncAgent - - print("[1/4] Load the documented reference agent") - agent = Agent(Args()) - dbg = os.path.join(SCRATCH, "verb_dbg") - os.makedirs(dbg, exist_ok=True) - IMG_DIR = os.path.join(SCRATCH, "verb_imgs") - os.makedirs(IMG_DIR, exist_ok=True) - - cap = {} - orig = agent.model.generate - - def hooked(*a, **kw): - cap["fired"] = True - cap["input_ids"] = kw.get("input_ids") - return orig(*a, **kw) - agent.model.generate = hooked - - scenes = sorted(g for g in glob.glob(os.path.join(ACTIVE, "assets/realworld_sample_data*")) - if os.path.isdir(g)) - samples = [] - for scene in scenes: - sname = os.path.basename(scene) - instr = open(os.path.join(scene, "instruction.txt")).read().strip() - # documented demo order: sorted debug_raw_*.jpg (look_down frames interleaved) - rgb_paths = sorted(glob.glob(os.path.join(scene, "debug_raw_*.jpg"))) - print(f"[2/4] {sname} | instr={instr!r} | {len(rgb_paths)} frames (docs loop)") - agent.reset() - agent.save_dir = dbg - for p in rgb_paths: - look_down = ('look_down' in p) - rgb = np.asarray(Image.open(p).convert('RGB')) - depth = 10 * np.ones((rgb.shape[0], rgb.shape[1]), np.float32) # docs: fill depth - pose = np.eye(4) - cap.clear() - try: - with torch.no_grad(): - agent.step(rgb, depth, pose, instr, intrinsic=INTR, look_down=look_down) - except Exception as e: - print(f" {os.path.basename(p)}: step err {type(e).__name__}: {e}") - continue - if not cap.get("fired"): - continue # S2 did not run on this frame (buffer) - skip - ref_out = agent.llm_output.strip() - msgs, ii = [], 0 - for turn in agent.conversation_history: - content = [] - for it in turn["content"]: - if it["type"] == "image": - fp = os.path.join(IMG_DIR, f"s{len(samples)}_{ii}.png") - it["image"].save(fp) - content.append({"type": "image", "image": fp}) - ii += 1 - else: - content.append({"type": "text", "text": it["text"]}) - msgs.append({"role": turn["role"], "content": content}) - samples.append(dict(scene=sname, img=os.path.basename(p), look_down=look_down, - ref=ref_out, S=int(cap["input_ids"].shape[1]) if cap.get( - "input_ids") is not None else -1, - messages=msgs)) - print(f" {os.path.basename(p):32} look_down={look_down} S={samples[-1]['S']} ref={ref_out!r}") - - in_json = os.path.join(SCRATCH, "verb_in.json") - out_json = os.path.join(SCRATCH, "verb_out.json") - json.dump({"batch_size": 1, "temperature": 0.0, "top_p": 1.0, "top_k": 1, "max_generate_length": 128, - "requests": [{"messages": s["messages"]} for s in samples]}, open(in_json, "w")) - del agent - torch.cuda.empty_cache() - - print(f"[3/4] Engine {os.path.basename(os.path.dirname(ENGINE))} through llm_inference ({len(samples)} req)") - env = dict(os.environ, EDGELLM_PLUGIN_PATH=f"{TRT}/build/libNvInfer_edgellm_plugin.so") - r = subprocess.run([f"{TRT}/build/examples/llm/llm_inference", "--engineDir", ENG_LLM, - "--multimodalEngineDir", ENG_VIS, "--inputFile", in_json, "--outputFile", out_json], - cwd=TRT, env=env, capture_output=True, text=True) - resp = json.load(open(out_json)).get("responses", []) if os.path.exists(out_json) else [] - print(f" exit={r.returncode}, {len(resp)} responses") - - print("[4/4] Compare engine against the reference verbatim\n" + "=" * 66) - import re - - def xy(t): - d = [int(c) for c in re.findall(r"\d+", t)] - return (d[0], d[1]) if len(d) >= 2 else None - ex = br = 0 - l2 = [] - for i, s in enumerate(samples): - fp = resp[i].get("output_text", "").strip() if i < len(resp) else "" - s["eng"] = fp - rb = "coord" if any(c.isdigit() for c in s["ref"]) else "action" - fb = "coord" if any(c.isdigit() for c in fp) else "action" - ex += s["ref"] == fp - br += rb == fb - a, b = xy(s["ref"]), xy(fp) - if a and b: - l2.append(((a[0] - b[0])**2 + (a[1] - b[1])**2)**.5) - print(f" [{'OK' if s['ref']==fp else 'DIFF':4}] {s['img']:32} ref={s['ref']!r:18} eng={fp[:40]!r}") - n = len(samples) - print("\n" + "=" * 66) - print(f" exact-match : {ex}/{n} = {ex/n*100:.1f}%") - print(f" branch-agree : {br}/{n} = {br/n*100:.1f}%") - if l2: - print( - f" coord-L2 px : mean={sum(l2)/len(l2):.1f} median={sorted(l2)[len(l2)//2]:.1f} max={max(l2):.0f} (n={len(l2)})") # noqa: E501 - tag = os.path.basename(os.path.dirname(ENGINE)).replace("edgellm_engines_", "") - json.dump(samples, open(os.path.join(os.environ.get("VLN_OPT_OUT", os.path.expanduser("~/vln-opt-work/out")), f"verbatim_{tag}.json"), "w"), # noqa: E501 - indent=2, ensure_ascii=False) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_e2e_agent.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_e2e_agent.py deleted file mode 100644 index 2b4920a..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_e2e_agent.py +++ /dev/null @@ -1,220 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""End-to-end verification: run the InternVLA agent with the TensorRT engines -(FP8 LLM for the bridge + BF16 System 1 engines) and compare its outputs, frame by -frame, against the pure-PyTorch agent on the documented sample data. -""" -import os -import sys -import glob -_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, _R) -sys.path.insert(0, os.path.join(_R, "lib")) -sys.path.append("/usr/lib/python3.12/dist-packages") -sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "lib")) -import numpy as np # noqa: E402 -import torch # noqa: E402 -from PIL import Image # noqa: E402 -from engine_runner import build_mrope_table, run_engine # LLM engine harness # noqa: E402 - -ACTIVE = os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) -CKPT = os.path.join(ACTIVE, "checkpoints/InternVLA-N1-DualVLN") -REPKG = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "qwen25vl_system2") -TRAJDIT = os.path.join( - os.environ.get( - "WORK_DIR", - os.path.expanduser("~/vln-opt-work")), - "onnx/system1_traj_dit_async_bf16.engine") # noqa: E131 -MEM = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.engine") -SCRATCH = os.path.expanduser(os.environ.get('VLN_OPT_OUT', '~/vln-opt-work/out')) -INTR = np.array([[386.5, 0, 328.9, 0], [0, 386.5, 244, 0], [0, 0, 1, 0], [0, 0, 0, 1]]) -TRAJ_TOKEN_INDEX = 151667 -IMAGE_TOKEN_INDEX = 151655 -SEED = 12345 - - -def out_of(o, k): - if isinstance(o, dict): - return o.get(k, next(iter(o.values()))) - return o[0] if isinstance(o, (list, tuple)) else o - - -def main(): - dev = "cuda" - try: - torch.backends.mha.set_fastpath_enabled(False) - except Exception: - pass - if ACTIVE not in sys.path: - sys.path.insert(0, ACTIVE) - from internvla_compat import apply_all - apply_all(need_system1=True, allow_missing_depth=True) - import importlib.util - _s = importlib.util.spec_from_file_location("iar", os.path.join( - ACTIVE, "internnav/agent/internvla_n1_agent_realworld.py")) - _m = importlib.util.module_from_spec(_s) - _s.loader.exec_module(_m) - Agent = _m.InternVLAN1AsyncAgent - from diffusers.schedulers import FlowMatchEulerDiscreteScheduler - from diffusers.utils.torch_utils import randn_tensor - from trt_torch import Engine - - class A: - device = "cuda:0" - model_path = CKPT - model_path_original = CKPT - resize_w = 384 - resize_h = 384 - num_history = 8 - plan_step_gap = 4 - print("[1/5] Load agent + engines") - agent = Agent(A()) - agent.save_dir = os.path.join(SCRATCH, "trtagent_dbg") - os.makedirs(agent.save_dir, exist_ok=True) - model = agent.model - m = model.get_model() - mem_eng = Engine(MEM) - dit_eng = Engine(TRAJDIT) - lm = m.language_model if hasattr(m, "language_model") else m - final_norm = lm.norm - NQ = model.get_n_query() - rmean = model._resnet_mean - rstd = model._resnet_std - - # ---- FP8 LLM generate_latents (latent_query + mRoPE logic, run on the engine) ---- - def trt_generate_latents(input_ids, pixel_values, image_grid_thw): - with torch.no_grad(): - te = m.embed_tokens(input_ids) - ie = model.visual(pixel_values.type(model.visual.dtype), grid_thw=image_grid_thw) - te[input_ids == IMAGE_TOKEN_INDEX] = ie.to(te.dtype)[:(input_ids == IMAGE_TOKEN_INDEX).sum(), :] - lq = m.latent_queries.repeat(te.shape[0], 1, 1) - inputs_embeds = torch.cat([te, lq], dim=1) - ids_traj = torch.cat([input_ids, torch.tensor([[TRAJ_TOKEN_INDEX] * NQ], device=dev)], dim=1) - position_ids, _ = model.get_rope_index(ids_traj, image_grid_thw) - rope = build_mrope_table(position_ids, dev) - _, eng_hs = run_engine(inputs_embeds.to(torch.float16), rope) - eng_pre = eng_hs[:, -NQ:, :].float() - return final_norm(eng_pre.to(torch.bfloat16)) # hidden [1,NQ,3584] (as generate_latents) - - # ---- S1 generate_traj through the engines (async logic matches the baseline) ---- - def trt_generate_traj(traj_latents, images_dp, depths_dp=None, predict_step_nums=32, - guidance_scale=1.0, num_inference_steps=10, num_sample_trajs=32): - torch.manual_seed(SEED) - np.random.seed(SEED) - dtype = traj_latents.dtype - with torch.no_grad(): - tl = m.cond_projector(traj_latents) - xdp = images_dp.permute(0, 1, 4, 2, 3) - xdp = ((xdp - rmean) / rstd).flatten(0, 1) - memory_tokens = out_of(mem_eng(images=xdp.float().contiguous()), "memory_tokens").to(dtype) - hs = torch.cat([memory_tokens, tl], dim=1) - hs_in = torch.cat([torch.zeros_like(hs), hs], 0) - bs = tl.shape[0] - latents = randn_tensor((bs * num_sample_trajs, predict_step_nums, 3), - generator=None, device=dev, dtype=dtype) - sch = FlowMatchEulerDiscreteScheduler() - sch.set_timesteps( - num_inference_steps, # noqa: E122 - sigmas=np.linspace( # noqa: E122 - 1.0, - 1 / num_inference_steps, - num_inference_steps)) # noqa: E131 - hs_in = hs_in.repeat_interleave(num_sample_trajs, dim=0) - dit_eng.set_runtime_tensor_shape("z_latents", tuple(hs_in.shape)) - for t in sch.timesteps: - lf = m.action_encoder(latents) - pid = torch.arange(lf.shape[1]).reshape(1, -1).repeat(bs * num_sample_trajs, 1).to(dev) - lf = lf + m.pos_encoding(pid) - lmi = lf.repeat(2, 1, 1) - if hasattr(sch, "scale_model_input"): - lmi = sch.scale_model_input(lmi, t) - tt = t.unsqueeze(0).expand(lmi.shape[0]).to(dev, torch.long) - npd = out_of( - dit_eng( # noqa: E122 - x=lmi.float().contiguous(), - timestep=tt.to( - torch.int64).contiguous(), - z_latents=hs_in.float().contiguous()), # noqa: E131 - "output").to(dtype) # noqa: E122 - npd = m.action_decoder(npd) - unc, cnd = npd.chunk(2) - npd = unc + guidance_scale * (cnd - unc) - latents = sch.step(npd, t, latents).prev_sample - return latents - - # PyTorch reference generate_traj: fixed seed for a fair comparison - base_traj = model.generate_traj - - def seeded_base_traj(*a, **k): - torch.manual_seed(SEED) - np.random.seed(SEED) - return base_traj(*a, **k) - - scene = sorted(g for g in glob.glob(os.path.join(ACTIVE, "assets/realworld_sample_data*")) if os.path.isdir(g))[0] - instr = open(os.path.join(scene, "instruction.txt")).read().strip() - rgbs = sorted(glob.glob(os.path.join(scene, "debug_raw_*.jpg")))[:40] - print(f"[2/5] scene {os.path.basename(scene)} | {len(rgbs)} frames | instr={instr[:50]!r}...") - - def run(tag): - agent.reset() - agent.save_dir = os.path.join(SCRATCH, "trtagent_dbg") - outs = [] - for p in rgbs: - ld = ('look_down' in p) - rgb = np.asarray(Image.open(p).convert('RGB')) - depth = 10 * np.ones(rgb.shape[:2], np.float32) - pose = np.eye(4) - try: - with torch.no_grad(): - o = agent.step(rgb, depth, pose, instr, intrinsic=INTR, look_down=ld) - except Exception as e: - print(f" {os.path.basename(p)}: {type(e).__name__}: {e}") - continue - traj = o.output_trajectory - act = o.output_action - if traj is not None: - outs.append(("traj", os.path.basename(p), np.asarray(traj))) - elif act is not None: - outs.append(("act", os.path.basename(p), list(act))) - return outs - - print("[3/5] Run PyTorch agent (reference)") - model.generate_traj = seeded_base_traj - ref = run("pytorch") - print(f" {len(ref)} outputs") - - print("[4/5] Run TensorRT agent (FP8 LLM latents + S1 engines)") - model.generate_latents = trt_generate_latents - model.generate_traj = trt_generate_traj - trt = run("trt") - print(f" {len(trt)} outputs") - - print("[5/5] Compare e2e TRT agent vs PyTorch\n" + "=" * 56) - n = min(len(ref), len(trt)) - mact = 0 - nact = 0 - trajerr = [] - for i in range(n): - rk, rf, rv = ref[i] - tk, tf, tv = trt[i] - if rk == "act" and tk == "act": - nact += 1 - mact += (rv == tv) - elif rk == "traj" and tk == "traj": - e = np.linalg.norm(np.asarray(rv) - np.asarray(tv), axis=-1) - trajerr.append(float(np.mean(e))) - print( - f" outputs: pytorch={len(ref)} trt={len(trt)} (type match {sum(1 for i in range(n) if ref[i][0]==trt[i][0])}/{n})") # noqa: E501 - if nact: - print(f" action match: {mact}/{nact}") - if trajerr: - import statistics - print( - f" trajectory per-wp L2 (m): mean={statistics.mean(trajerr):.4f} max={max(trajerr):.4f} (n={len(trajerr)})") # noqa: E501 - print(" Agent TRT runs end-to-end and matches PyTorch." if n > 0 else " no output") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_engine_policy.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_engine_policy.py deleted file mode 100644 index d7d6ca5..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_engine_policy.py +++ /dev/null @@ -1,142 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Method-level verification of the engine-backed policy adapter (no simulator needed). - -Two checks: - 1. Import + subclass: EngineInternVLAN1Net subclasses the real InternVLAN1Net and only overrides - the LLM text generation (self.model.generate). - 2. Engine-generate roundtrip: drive the adapter's `_engine_generate` mechanism with a real VLN - look-down conversation (built via lib/prompt_builder) and confirm it returns token ids that - decode to the SAME text as a direct `llm_inference` run on the same messages. - -Closed-loop SR itself needs the sim (Habitat/InternUtopia) — that is the VLN team's step; this only -proves the engine is correctly wired into the policy's generate contract. -""" -import glob -import json -import os -import subprocess -import sys -import tempfile - -_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, _R) -sys.path.insert(0, os.path.join(_R, "lib")) - -import numpy as np # noqa: E402 -from PIL import Image # noqa: E402 -# prompt_builder is the single source of truth for the VLN prompt and lives in the -# quantize path, so calibration and verification cannot drift apart. Walk up to the -# recipe root rather than counting directory levels -- these scripts sit at two -# different depths. -_d = os.path.dirname(os.path.abspath(__file__)) -while _d != "/" and not os.path.isdir(os.path.join(_d, "quantize")): - _d = os.path.dirname(_d) -sys.path.insert(0, os.path.join(_d, "quantize")) -import prompt_builder as pb # noqa: E402 - -TRT = os.path.expanduser(os.environ.get("TRT_EDGE_LLM", "~/modelopt/TensorRT-Edge-LLM")) -LLM_DIR = os.path.expanduser(os.environ.get( - "VLN_LLM_ENGINE_DIR", "~/vln-opt-work/repro/engines/system2_llm_fp8_vlncalib")) -VIS_DIR = os.path.expanduser(os.environ.get( - "VLN_VIS_ENGINE_DIR", - os.path.join(os.environ.get("ENGINE_DIR", - os.path.expanduser("~/vln-opt-work/engines")), "s1_fp8/visual"))) -REPKG = os.path.expanduser(os.environ.get("REPKG", "~/vln-opt-work/repro/qwen25vl_system2")) -DATA = os.path.expanduser(os.environ.get("VLN_CALIB_DATA", "~/vln-opt-work/probe_heldout")) -env = dict(os.environ, EDGELLM_PLUGIN_PATH=os.path.join(TRT, "build/libNvInfer_edgellm_plugin.so")) - - -def check_import(): - print("[1/2] Import + subclass check") - try: - from engine_policy import EngineInternVLAN1Net - from internnav.model.basemodel.internvla_n1.internvla_n1_policy import InternVLAN1Net - except Exception as e: # noqa: BLE001 - print(f" import FAILED: {type(e).__name__}: {e}") - return False - ok = issubclass(EngineInternVLAN1Net, InternVLAN1Net) - # the only override on the class body is __init__ (which patches generate) + _engine_generate - overrides = set(EngineInternVLAN1Net.__dict__) - {"__init__", "__doc__", "__module__"} - print(f" subclass of InternVLAN1Net: {ok}") - print(f" class adds: {sorted(overrides)} (expect just _engine_generate)") - return ok and overrides <= {"_engine_generate"} - - -def one_lookdown_messages(tmp): - """Build one real VLN look-down conversation (turn-1 '↓', turn-2 tilted view) as llm_inference - messages + save images, using the same prompt_builder the policy uses.""" - meta = sorted(glob.glob(os.path.join(DATA, "**", "meta", "episodes.jsonl"), recursive=True))[0] - scene = os.path.dirname(os.path.dirname(meta)) - eps = [json.loads(l) for l in open(meta) if l.strip()] # noqa: E741 - ep = [e for e in eps if e.get("length", 0) > 20 and e.get("tasks")][0] - t = ep["length"] // 2 - lvl = f"{scene}/videos/chunk-000/observation.images.rgb.125cm_0deg" - tilt = f"{scene}/videos/chunk-000/observation.images.rgb.125cm_30deg" - hist = np.unique(np.linspace(0, t - 1, pb.NUM_HISTORY, dtype=np.int32)).tolist() - srcs = [f"{lvl}/episode_{ep['episode_index']:06d}_{h}.jpg" for h in hist] - srcs += [f"{lvl}/episode_{ep['episode_index']:06d}_{t}.jpg", - f"{tilt}/episode_{ep['episode_index']:06d}_{t}.jpg"] - conv = pb.build_conversation_lookdown(t, ep["tasks"][0], "↓") - paths, k = [], 0 - for turn in conv: - for it in turn["content"]: - if it["type"] == "image": - d = os.path.join(tmp, f"i{k}.png") - # the real policy resizes to (resize_w, resize_h)=384 before inference - Image.open(srcs[k]).convert("RGB").resize( - (pb.RESIZE_W, pb.RESIZE_H)).save(d) - it["image"] = d - k += 1 - paths.append(d) - # to llm_inference messages (images already file paths) - return conv - - -def run_llm_inference(messages, tmp): - in_json = os.path.join(tmp, "in.json") - out_json = os.path.join(tmp, "out.json") - json.dump({"batch_size": 1, "temperature": 0.0, "top_p": 1.0, "top_k": 1, - "max_generate_length": 32, "requests": [{"messages": messages}]}, open(in_json, "w")) - subprocess.run([os.path.join(TRT, "build/examples/llm/llm_inference"), - "--engineDir", LLM_DIR, "--multimodalEngineDir", VIS_DIR, - "--inputFile", in_json, "--outputFile", out_json], - env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True) - return json.load(open(out_json))["responses"][0]["output_text"].strip() - - -def check_roundtrip(): - print("[2/2] Engine-generate roundtrip (adapter mechanism vs direct llm_inference)") - from transformers import AutoTokenizer - tok = AutoTokenizer.from_pretrained(REPKG, use_fast=True) - tmp = tempfile.mkdtemp(prefix="verifyeng_") - try: - conv = one_lookdown_messages(tmp) - # (a) direct engine text - direct = run_llm_inference(conv, tmp) - # (b) the adapter's roundtrip: text -> ids -> decode (must reproduce `direct`) - gen_ids = tok(direct, return_tensors="pt", add_special_tokens=False).input_ids - roundtrip = tok.decode(gen_ids[0], skip_special_tokens=True).strip() - finally: - import shutil - shutil.rmtree(tmp, ignore_errors=True) - print(f" engine output text : {direct!r}") - print(f" tokenizer roundtrip: {roundtrip!r}") - has_coord = any(c.isdigit() for c in direct) - ok = roundtrip == direct - print(f" roundtrip exact : {ok} | output is a pixel-goal/coord: {has_coord}") - return ok - - -def main(): - a = check_import() - b = check_roundtrip() - print("\n" + ("PASS — adapter wires the engine into the policy generate contract correctly." - if (a and b) else "CHECK — see failures above.")) - print("Closed-loop SR must still be run in the simulator (VLN team). See HANDOVER.md.") - return 0 if (a and b) else 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents.py deleted file mode 100644 index 5779602..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents.py +++ /dev/null @@ -1,202 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Verify System 2 numeric fidelity: z_latents cosine of the FP8 LLM engine vs the -PyTorch reference, using the real latent-query bridge path. -""" -import os -import sys -import json -_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, _R) -sys.path.insert(0, _R) -import torch # noqa: E402 -from PIL import Image # noqa: E402 -from engine_runner import (REPKG, ENGINE, CKPT, build_mrope_table, run_engine, # noqa: E402 - cos, per_tok_cos) - -TRAJ_TOKEN_INDEX = 151667 -IMAGE_TOKEN_INDEX = 151655 -N_QUERY = 4 -IMAGE = os.path.expanduser(os.environ.get( - "IMAGE_PATH", "~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg")) - - -def load_ckpt_tensor(prefix): - """Read bridge tensors, preferring the small bridge.safetensors next to the checkpoint. - - repackage_system2.py sets latent_queries and cond_projector aside in a ~25 MB - bridge.safetensors precisely so this check does not have to reopen the 16 GB original. - Falling back to a full sharded checkpoint keeps the script usable against one. - """ - from safetensors import safe_open - - bridge = os.path.join(CKPT, "bridge.safetensors") - if os.path.isfile(bridge): - with safe_open(bridge, framework="pt") as f: - return {k: f.get_tensor(k) for k in f.keys() if k.startswith(prefix)} - - index = os.path.join(CKPT, "model.safetensors.index.json") - if not os.path.isfile(index): - raise FileNotFoundError( - f"neither bridge.safetensors nor model.safetensors.index.json under {CKPT}; " - f"point INTERNVLA_CKPT at a repackaged checkpoint or the original") - idx = json.load(open(index))["weight_map"] - out = {} - for k in idx: - if k.startswith(prefix): - with safe_open(os.path.join(CKPT, idx[k]), framework="pt") as f: - out[k] = f.get_tensor(k) - return out - - -def rope_index(model, input_ids, image_grid_thw, image_token_id): - """Call get_rope_index across the two transformers signatures. - - transformers 4.x took (input_ids, image_grid_thw). 5.x inserted a required - mm_token_type_ids argument -- 0 text, 1 image, 2 video -- and moved the grids to - keywords. Passing the 4.x form to a 5.x model does not raise on arity; it fails later - inside with 'NoneType is not an iterator', which is a confusing way to learn this. - """ - import inspect - - fn = model_attr(model, "get_rope_index") - params = inspect.signature(fn).parameters - if "mm_token_type_ids" in params: - mm = (input_ids == image_token_id).to(torch.int32) - return fn(input_ids, mm, image_grid_thw=image_grid_thw) - return fn(input_ids, image_grid_thw) - - -def model_attr(model, name): - """Fetch an attribute or bound method from the model or its inner model. - - transformers 5.x moved several Qwen2.5-VL helpers (get_rope_index, visual, ...) from - the ForConditionalGeneration wrapper down onto the inner Qwen2_5_VLModel. Probing both - keeps this working on either layout instead of pinning a transformers version. - """ - for owner in (model, getattr(model, "model", None)): - if owner is not None and hasattr(owner, name): - return getattr(owner, name) - raise AttributeError(f"neither the model nor its inner model has {name!r}") - - -def get_visual(model): - """Return the vision tower across transformers layouts. - - Older versions expose it as ``model.visual``; newer ones nest it under - ``model.model.visual``. Probing beats pinning a transformers version here. - """ - for owner in (model, getattr(model, "model", None)): - vis = getattr(owner, "visual", None) if owner is not None else None - if vis is not None: - return vis - raise AttributeError("no vision tower found on this model " - "(looked at .visual and .model.visual)") - - -def visual_embeds(model, pixel_values, grid_thw): - """Return merged image embeddings at LLM hidden width, across transformers layouts. - - In transformers 5.x get_image_features returns a BaseModelOutputWithPooling whose - ``pooler_output`` holds the *merged* embeddings (split per image), while - ``last_hidden_state`` is the pre-merger tensor at vision width -- 1280 here against the - LLM's 3584. Reading the wrong field fails loudly on the shape, but only after the - vision tower has already run, so unwrap explicitly. - """ - if hasattr(model, "get_image_features"): - out = model.get_image_features(pixel_values.type(get_visual(model).dtype), grid_thw) - feats = getattr(out, "pooler_output", out) - else: - vis = get_visual(model) - out = vis(pixel_values.type(vis.dtype), grid_thw=grid_thw) - feats = getattr(out, "last_hidden_state", out) - if isinstance(feats, (list, tuple)): - feats = torch.cat([f.reshape(-1, f.shape[-1]) for f in feats], dim=0) - return feats.reshape(-1, feats.shape[-1]) - - -def main(): - dev = "cuda" - from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration - print(f"[1/6] Load repackage + processor | engine={os.path.basename(ENGINE)}") - model = Qwen2_5_VLForConditionalGeneration.from_pretrained( - REPKG, torch_dtype=torch.bfloat16, # flash-attn is not available on Jetson; sdpa is the supported path and - # is what the deployed agent uses too. - attn_implementation=os.environ.get("ATTN_IMPL", "sdpa"), - low_cpu_mem_usage=True).to(dev).eval() - proc = AutoProcessor.from_pretrained(REPKG, trust_remote_code=True, - min_pixels=128 * 28 * 28, max_pixels=1024 * 28 * 28) - inner = model.model # Qwen2_5_VLModel (takes inputs_embeds) - lm = inner.language_model if hasattr(inner, "language_model") else inner - final_norm = lm.norm - - print("[2/6] latent_queries + cond_projector from the InternVLA checkpoint") - lq = load_ckpt_tensor("model.latent_queries")["model.latent_queries"].to(dev).to(torch.bfloat16) - cpw = load_ckpt_tensor("model.cond_projector") - cp = {k.replace("model.cond_projector.", ""): v.float().to(dev) for k, v in cpw.items()} - print(f" latent_queries {tuple(lq.shape)} (n_query={lq.shape[1]}) cond keys {sorted(cp.keys())}") - assert lq.shape[1] == N_QUERY - - print("[3/6] Build image + instruction input, append 4 TRAJ tokens (as generate_latents)") - msgs = [{"role": "user", "content": [ - {"type": "image", "image": IMAGE}, - {"type": "text", "text": "Go straight then stop at the green plant."}]}] - text = proc.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) - enc = proc(text=[text], images=[Image.open(IMAGE).convert("RGB")], return_tensors="pt").to(dev) - input_ids = enc["input_ids"] - grid = enc["image_grid_thw"] - # embed + scatter image + append latent_queries - with torch.no_grad(): - text_embeds = model.get_input_embeddings()(input_ids) # [1,S,3584] - image_embeds = visual_embeds(model, enc["pixel_values"], grid) - image_idx = (input_ids == IMAGE_TOKEN_INDEX) - text_embeds[image_idx] = image_embeds.to(text_embeds.dtype)[: image_idx.sum(), :] - inputs_embeds = torch.cat([text_embeds, lq.repeat(text_embeds.shape[0], 1, 1)], dim=1) - ids_traj = torch.cat( - [input_ids, torch.tensor([[TRAJ_TOKEN_INDEX] * N_QUERY], device=dev)], dim=1) - S = inputs_embeds.shape[1] - position_ids, _ = rope_index(model, ids_traj, grid, IMAGE_TOKEN_INDEX) - print(f" seq_len={S} (+{N_QUERY} traj) grid={grid.tolist()} pos {tuple(position_ids.shape)}") - - print("[4/6] Reference forward (inner LM); capture pre/post norm at the last 4 positions") - cap = {} - h = final_norm.register_forward_hook(lambda m, i, o: cap.update(pre=i[0].detach(), post=o.detach())) - with torch.no_grad(): - inner(inputs_embeds=inputs_embeds, position_ids=position_ids, - output_hidden_states=True, return_dict=True) - h.remove() - ref_pre = cap["pre"][:, -N_QUERY:, :].float() - ref_post = cap["post"][:, -N_QUERY:, :].float() - - print("[5/6] Engine forward (3D mRoPE); take hidden[-4:] pre-norm -> host norm") - rope = build_mrope_table(position_ids, dev) - _, eng_hs = run_engine(inputs_embeds.to(torch.float16), rope) - eng_pre = eng_hs[:, -N_QUERY:, :].float() - with torch.no_grad(): - eng_post = final_norm(eng_pre.to(torch.bfloat16)).float() - - print("[6/6] Compare z_latents (System1 traj_dit input)\n" + "=" * 60) - - def cond_project(x): - x = torch.nn.functional.linear(x, cp["0.weight"], cp.get("0.bias")) - x = torch.nn.functional.gelu(x) - return torch.nn.functional.linear(x, cp["2.weight"], cp.get("2.bias")) - z_ref, z_eng = cond_project(ref_post), cond_project(eng_post) - ptc = torch.nn.functional.cosine_similarity(eng_pre[0], ref_pre[0], dim=-1) - print(" per-query PRE cosine:", " ".join(f"{ptc[i]:.4f}" for i in range(N_QUERY))) - print(f" hidden PRE-norm cosine = {per_tok_cos(eng_pre, ref_pre):.6f}") - print(f" hidden POST-norm cosine = {per_tok_cos(eng_post, ref_post):.6f}") - zc = per_tok_cos(z_eng, z_ref) - zc_flat = cos(z_eng, z_ref) - l2 = (z_eng - z_ref).norm() / z_ref.norm() - print(f" z_latents({z_eng.shape[-1]}) per-query cosine = {zc:.6f} | flat = {zc_flat:.6f} " - f"| rel-L2 = {l2:.4f}") - ok = zc > 0.99 - print("\n" + (f"FAITHFUL e2e bridge PASS - z_latents match (cosine {zc:.4f})" - if ok else f"⚠️ z_latents cosine {zc:.4f} < 0.99")) - return 0 if ok else 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents_vln.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents_vln.py deleted file mode 100644 index f928de0..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_latents_vln.py +++ /dev/null @@ -1,219 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Verify System 2 z_latents fidelity on the REAL VLN input distribution. - -Same z_latents bridge comparison as ``verify_system2_latents.py`` (FP8 engine vs -PyTorch reference), but instead of a single caption image it drives the model with -genuine InternData-N1 VLN-CE steps (multi-image history + navigation instruction + -4 appended TRAJ tokens), averaged over several samples. This is the probe that -actually reflects deployment, so it is the fair test of whether VLN-distribution -FP8 calibration helps — a single-image caption probe is out-of-distribution. - -Select the engine with ENGINE_PATH; select the calibration/probe data with -VLN_CALIB_DATA. Uses a fixed seed so every engine sees identical inputs. -""" -import os -import json -import sys - -_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, _R) -sys.path.insert(0, _R) -sys.path.insert(0, os.path.join(_R, "build", "quantize")) - -import torch # noqa: E402 - -from engine_runner import (REPKG, ENGINE, CKPT, build_mrope_table, run_engine, # noqa: E402 - cos, per_tok_cos) - -TRAJ_TOKEN_INDEX = 151667 -IMAGE_TOKEN_INDEX = 151655 -N_QUERY = 4 -N_SAMPLES = int(os.environ.get("VLN_PROBE_SAMPLES", "12")) -DATA = os.path.expanduser(os.environ.get( - "VLN_CALIB_DATA", - "~/vln-opt-work/calib_scenes")) # output of build/00_fetch_calib_scenes.sh; override via env - - -def load_ckpt_tensor(prefix): - """Read bridge tensors, preferring the small bridge.safetensors next to the checkpoint. - - repackage_system2.py sets latent_queries and cond_projector aside in a ~25 MB - bridge.safetensors precisely so this check does not have to reopen the 16 GB original. - Falling back to a full sharded checkpoint keeps the script usable against one. - """ - from safetensors import safe_open - - bridge = os.path.join(CKPT, "bridge.safetensors") - if os.path.isfile(bridge): - with safe_open(bridge, framework="pt") as f: - return {k: f.get_tensor(k) for k in f.keys() if k.startswith(prefix)} - - index = os.path.join(CKPT, "model.safetensors.index.json") - if not os.path.isfile(index): - raise FileNotFoundError( - f"neither bridge.safetensors nor model.safetensors.index.json under {CKPT}; " - f"point INTERNVLA_CKPT at a repackaged checkpoint or the original") - idx = json.load(open(index))["weight_map"] - out = {} - for k in idx: - if k.startswith(prefix): - with safe_open(os.path.join(CKPT, idx[k]), framework="pt") as f: - out[k] = f.get_tensor(k) - return out - - -def rope_index(model, input_ids, image_grid_thw, image_token_id): - """Call get_rope_index across the two transformers signatures. - - transformers 4.x took (input_ids, image_grid_thw). 5.x inserted a required - mm_token_type_ids argument -- 0 text, 1 image, 2 video -- and moved the grids to - keywords. Passing the 4.x form to a 5.x model does not raise on arity; it fails later - inside with 'NoneType is not an iterator', which is a confusing way to learn this. - """ - import inspect - - fn = model_attr(model, "get_rope_index") - params = inspect.signature(fn).parameters - if "mm_token_type_ids" in params: - mm = (input_ids == image_token_id).to(torch.int32) - return fn(input_ids, mm, image_grid_thw=image_grid_thw) - return fn(input_ids, image_grid_thw) - - -def model_attr(model, name): - """Fetch an attribute or bound method from the model or its inner model. - - transformers 5.x moved several Qwen2.5-VL helpers (get_rope_index, visual, ...) from - the ForConditionalGeneration wrapper down onto the inner Qwen2_5_VLModel. Probing both - keeps this working on either layout instead of pinning a transformers version. - """ - for owner in (model, getattr(model, "model", None)): - if owner is not None and hasattr(owner, name): - return getattr(owner, name) - raise AttributeError(f"neither the model nor its inner model has {name!r}") - - -def get_visual(model): - """Return the vision tower across transformers layouts. - - Older versions expose it as ``model.visual``; newer ones nest it under - ``model.model.visual``. Probing beats pinning a transformers version here. - """ - for owner in (model, getattr(model, "model", None)): - vis = getattr(owner, "visual", None) if owner is not None else None - if vis is not None: - return vis - raise AttributeError("no vision tower found on this model " - "(looked at .visual and .model.visual)") - - -def visual_embeds(model, pixel_values, grid_thw): - """Return merged image embeddings at LLM hidden width, across transformers layouts. - - In transformers 5.x get_image_features returns a BaseModelOutputWithPooling whose - ``pooler_output`` holds the *merged* embeddings (split per image), while - ``last_hidden_state`` is the pre-merger tensor at vision width -- 1280 here against the - LLM's 3584. Reading the wrong field fails loudly on the shape, but only after the - vision tower has already run, so unwrap explicitly. - """ - if hasattr(model, "get_image_features"): - out = model.get_image_features(pixel_values.type(get_visual(model).dtype), grid_thw) - feats = getattr(out, "pooler_output", out) - else: - vis = get_visual(model) - out = vis(pixel_values.type(vis.dtype), grid_thw=grid_thw) - feats = getattr(out, "last_hidden_state", out) - if isinstance(feats, (list, tuple)): - feats = torch.cat([f.reshape(-1, f.shape[-1]) for f in feats], dim=0) - return feats.reshape(-1, feats.shape[-1]) - - -def main(): - dev = "cuda" - from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration - from calibration import vln_calib_dataloader - - print(f"[1/4] Load repackage + processor | engine={os.path.basename(ENGINE)}") - model = Qwen2_5_VLForConditionalGeneration.from_pretrained( - REPKG, torch_dtype=torch.bfloat16, # flash-attn is not available on Jetson; sdpa is the supported path and - # is what the deployed agent uses too. - attn_implementation=os.environ.get("ATTN_IMPL", "sdpa"), - low_cpu_mem_usage=True).to(dev).eval() - proc = AutoProcessor.from_pretrained(REPKG, trust_remote_code=True, - min_pixels=128 * 28 * 28, max_pixels=1024 * 28 * 28) - inner = model.model - lm = inner.language_model if hasattr(inner, "language_model") else inner - final_norm = lm.norm - - lq = load_ckpt_tensor("model.latent_queries")["model.latent_queries"].to(dev).to(torch.bfloat16) - cpw = load_ckpt_tensor("model.cond_projector") - cp = {k.replace("model.cond_projector.", ""): v.float().to(dev) for k, v in cpw.items()} - assert lq.shape[1] == N_QUERY - - def cond_project(x): - x = torch.nn.functional.linear(x, cp["0.weight"], cp.get("0.bias")) - x = torch.nn.functional.gelu(x) - return torch.nn.functional.linear(x, cp["2.weight"], cp.get("2.bias")) - - print(f"[2/4] Build {N_SAMPLES} real VLN inputs (multi-image + instruction) | data={DATA}") - batches = vln_calib_dataloader(proc, data_root=DATA, num_samples=N_SAMPLES, seed=0) - - print("[3/4] For each: reference LM forward vs engine forward at the 4 TRAJ positions") - hid_pre, hid_post, zc_list, zflat_list = [], [], [], [] - for bi, enc in enumerate(batches): - input_ids = enc["input_ids"].to(dev) - grid = enc["image_grid_thw"].to(dev) - pixel_values = enc["pixel_values"].to(dev) - with torch.no_grad(): - text_embeds = model.get_input_embeddings()(input_ids) - image_embeds = model.visual(pixel_values.type(model.visual.dtype), grid_thw=grid) - image_idx = (input_ids == IMAGE_TOKEN_INDEX) - text_embeds[image_idx] = image_embeds.to(text_embeds.dtype)[: image_idx.sum(), :] - inputs_embeds = torch.cat([text_embeds, lq.repeat(text_embeds.shape[0], 1, 1)], dim=1) - ids_traj = torch.cat( - [input_ids, torch.tensor([[TRAJ_TOKEN_INDEX] * N_QUERY], device=dev)], dim=1) - position_ids, _ = rope_index(model, ids_traj, grid, IMAGE_TOKEN_INDEX) - - cap = {} - h = final_norm.register_forward_hook( - lambda m, i, o: cap.update(pre=i[0].detach(), post=o.detach())) - with torch.no_grad(): - inner(inputs_embeds=inputs_embeds, position_ids=position_ids, - output_hidden_states=True, return_dict=True) - h.remove() - ref_pre = cap["pre"][:, -N_QUERY:, :].float() - ref_post = cap["post"][:, -N_QUERY:, :].float() - - rope = build_mrope_table(position_ids, dev) - _, eng_hs = run_engine(inputs_embeds.to(torch.float16), rope) - eng_pre = eng_hs[:, -N_QUERY:, :].float() - with torch.no_grad(): - eng_post = final_norm(eng_pre.to(torch.bfloat16)).float() - - z_ref, z_eng = cond_project(ref_post), cond_project(eng_post) - hp = per_tok_cos(eng_pre, ref_pre) - hpo = per_tok_cos(eng_post, ref_post) - zc = per_tok_cos(z_eng, z_ref) - zf = cos(z_eng, z_ref) - hid_pre.append(hp) - hid_post.append(hpo) - zc_list.append(zc) - zflat_list.append(zf) - print(f" #{bi:2d} imgs={grid.shape[0]:2d} seq={input_ids.shape[1]:4d} " - f"| hidPRE={hp:.5f} hidPOST={hpo:.5f} z={zc:.5f}") - - import statistics as st - n = len(zc_list) - print("\n[4/4] Mean over VLN inputs\n" + "=" * 60) - print(f" samples : {n}") - print(f" hidden PRE-norm : {st.mean(hid_pre):.6f} (min {min(hid_pre):.6f})") - print(f" hidden POST-norm : {st.mean(hid_post):.6f} (min {min(hid_post):.6f})") - print(f" z_latents perq : {st.mean(zc_list):.6f} (min {min(zc_list):.6f})") - print(f" z_latents flat : {st.mean(zflat_list):.6f}") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_pixelgoal_gt.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_pixelgoal_gt.py deleted file mode 100644 index 8c30f37..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_pixelgoal_gt.py +++ /dev/null @@ -1,217 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Offline task-level metric: System 2 pixel-goal L2 vs dataset ground truth. - -This is the metric InternVLA §4.2 uses for System 2: at a look-down goal frame the model -predicts the next waypoint's pixel `(u, v)`; the dataset stores the projected ground-truth -`goal.` (produced by the official label generator -`scripts/dataset_converters/internvla_labels.py::project_world_point`) in the original 640x480 -frame. Training targets the raw `f"{u} {v}"` (dataset line 1217, no rescaling), so predictions -and GT live in the same 640x480 space and L2 is direct. - -Teacher-forced two-turn look-down (matching the agent): turn 1 on the level view is forced to -"↓" (idx2actions[5]); turn 2 appends the tilted look-down view and the model emits `u v`. -The conversation is built once (lib/prompt_builder) and fed to BOTH PyTorch and the engines, so -they see identical prompts. - -Self-validation gate: PyTorch's own L2 vs GT must be small (paper regime, a few → tens of px). If -it is large, the coordinate mapping is wrong and NO number is reported — only cosine stands. - -Select the engine via ENGINE_PATH; select data via VLN_CALIB_DATA; pick the pitch via VLN_PITCH. -""" -import os -import sys -import glob -import json -import re -import subprocess -import tempfile - -_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, _R) -sys.path.insert(0, os.path.join(_R, "lib")) - -import numpy as np # noqa: E402 -from PIL import Image # noqa: E402 -# prompt_builder is the single source of truth for the VLN prompt and lives in the -# quantize path, so calibration and verification cannot drift apart. Walk up to the -# recipe root rather than counting directory levels -- these scripts sit at two -# different depths. -_d = os.path.dirname(os.path.abspath(__file__)) -while _d != "/" and not os.path.isdir(os.path.join(_d, "quantize")): - _d = os.path.dirname(_d) -sys.path.insert(0, os.path.join(_d, "quantize")) -import prompt_builder as pb # noqa: E402 - -TRT = os.path.expanduser(os.environ.get("TRT_EDGE_LLM", "~/modelopt/TensorRT-Edge-LLM")) -REPKG = os.path.expanduser(os.environ.get("REPKG", "~/vln-opt-work/repro/qwen25vl_system2")) -VIS = os.path.expanduser(os.environ.get( - "VIS_ENG", - os.path.join(os.environ.get("ENGINE_DIR", - os.path.expanduser("~/vln-opt-work/engines")), "s1_fp8/visual"))) -DATA = os.path.expanduser(os.environ.get("VLN_CALIB_DATA", "~/vln-opt-work/probe_heldout")) -# Look-down config: level history (0deg) + tilted goal view. GT lives at the tilted pitch. -LEVEL_KEY = "observation.images.rgb.125cm_0deg" -# One or more tilted look-down pitches (comma-separated). GT goal lives at each pitch. -PITCHES = [p.strip() for p in os.environ.get("VLN_PITCH", "125cm_30deg,60cm_30deg").split(",")] -LOOKDOWN_TOKEN = "↓" # idx2actions[5] -MAX_SAMPLES = int(os.environ.get("N", "0")) or None # None = all goal frames -ENGINE = os.path.expanduser(os.environ.get("ENGINE_PATH", "")) or None -env = dict(os.environ, EDGELLM_PLUGIN_PATH=f"{TRT}/build/libNvInfer_edgellm_plugin.so") -DIGITS = re.compile(r"\d+") - - -def frame_path(scene, key, ep_idx, fr): - return os.path.join(scene, "videos", "chunk-000", key, f"episode_{ep_idx:06d}_{fr}.jpg") - - -def collect_goal_frames(): - """Yield (scene_dir, pitch, ep_idx, t, instruction, gt_uv) for every populated goal frame, - across all configured pitches.""" - import pandas as pd - out = [] - for meta in glob.glob(os.path.join(DATA, "**", "meta", "episodes.jsonl"), recursive=True): - scene = os.path.dirname(os.path.dirname(meta)) - eps = {e["episode_index"]: e for e in - (json.loads(l) for l in open(meta) if l.strip())} # noqa: E741 - for pq in sorted(glob.glob(os.path.join(scene, "data", "chunk-000", "*.parquet"))): - df = pd.read_parquet(pq) - ep_idx = int(df["episode_index"].iloc[0]) - ep = eps.get(ep_idx) - if not ep or not ep.get("tasks"): - continue - for pitch in PITCHES: - gcol = f"goal.{pitch}" - if gcol not in df.columns or not os.path.isdir( - os.path.join(scene, "videos", "chunk-000", f"observation.images.rgb.{pitch}")): - continue - goals = np.stack(df[gcol].values) - for t in range(len(df)): - u, v = int(goals[t][0]), int(goals[t][1]) - if u < 0 or t == 0: - continue - out.append((scene, pitch, ep_idx, t, ep["tasks"][0], (u, v))) - return out - - -def build_conv(instruction, t, img_paths): - """One look-down conversation, images filled in order (history+current level, then tilt).""" - conv = pb.build_conversation_lookdown(t, instruction, LOOKDOWN_TOKEN) - j = 0 - for turn in conv: - for it in turn["content"]: - if it["type"] == "image" and it.get("image") is None: - it["image"] = img_paths[j] - j += 1 - assert j == len(img_paths), f"image count {j} != {len(img_paths)}" - return conv - - -def pred_uv(text): - d = [int(x) for x in DIGITS.findall(text)] - return (d[0], d[1]) if len(d) >= 2 else None # raw "u v" in 640x480 (policy.py decode) - - -def run_engine(conv, tmp): - js = {"batch_size": 1, "temperature": 0.0, "top_p": 1.0, "top_k": 1, - "max_generate_length": 16, "requests": [{"messages": conv}]} - inp = os.path.join(tmp, "in.json") - out = os.path.join(tmp, "out.json") - json.dump(js, open(inp, "w")) - subprocess.run([f"{TRT}/build/examples/llm/llm_inference", - "--engineDir", os.path.dirname(ENGINE), "--multimodalEngineDir", VIS, - "--inputFile", inp, "--outputFile", out], - env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True) - return json.load(open(out))["responses"][0]["output_text"].strip() - - -def main(): - frames = collect_goal_frames() - if not frames: - print(f"No populated goal frames for pitches {PITCHES} under {DATA}", file=sys.stderr) - return 1 - import random - random.Random(0).shuffle(frames) - if MAX_SAMPLES: - frames = frames[:MAX_SAMPLES] - sq = os.environ.get("VLN_SQUARE") == "1" - print(f"[data] {len(frames)} goal frames | pitches={PITCHES} | " - f"preproc={'384-square (deploy)' if sq else 'aspect-preserving (official)'} | engine=" - f"{os.path.basename(os.path.dirname(ENGINE)) if ENGINE else 'PyTorch-only'}") - - from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration - import torch - model = Qwen2_5_VLForConditionalGeneration.from_pretrained( - REPKG, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2", - low_cpu_mem_usage=True).to("cuda").eval() - proc = AutoProcessor.from_pretrained(REPKG, trust_remote_code=True, - min_pixels=128 * 28 * 28, max_pixels=1024 * 28 * 28) - - pt_l2, eng_l2 = [], [] - tmp = tempfile.mkdtemp(prefix="pgoal_") - for scene, PITCH, ep_idx, t, instr, (gu, gv) in frames: - TILT_KEY = f"observation.images.rgb.{PITCH}" - hist = np.unique(np.linspace(0, t - 1, pb.NUM_HISTORY, dtype=np.int32)).tolist() - # history + current from the level view; look-down frame from the tilted view - src = [frame_path(scene, LEVEL_KEY, ep_idx, h) for h in hist] - src.append(frame_path(scene, LEVEL_KEY, ep_idx, t)) - src.append(frame_path(scene, TILT_KEY, ep_idx, t)) - if not all(os.path.isfile(p) for p in src): - continue - # Feed native frames (aspect-preserving); the processor's min/max_pixels sizes them as in - # training. Forcing a 384x384 square would distort the 4:3 frame and shift the goal. GT - # stays in native 640x480, matching the model's output space. Set VLN_SQUARE=1 to override. - if os.environ.get("VLN_SQUARE") == "1": - paths = [] - for i, p in enumerate(src): - d = os.path.join(tmp, f"{i}.jpg") - Image.open(p).convert("RGB").resize((pb.RESIZE_W, pb.RESIZE_H)).save(d, quality=95) - paths.append(d) - else: - paths = src - conv = build_conv(instr, t, paths) - - imgs = [Image.open(p).convert("RGB") for p in paths] - text = proc.apply_chat_template(conv, tokenize=False, add_generation_prompt=True) - enc = proc(text=[text], images=imgs, return_tensors="pt").to("cuda") - with torch.no_grad(): - g = model.generate(**enc, max_new_tokens=16, do_sample=False) - pt_txt = proc.tokenizer.decode(g[0, enc["input_ids"].shape[1]:], skip_special_tokens=True) - puv = pred_uv(pt_txt) - if puv: - pt_l2.append(np.hypot(puv[0] - gu, puv[1] - gv)) - - if ENGINE: - euv = pred_uv(run_engine(conv, tmp)) - if euv: - eng_l2.append(np.hypot(euv[0] - gu, euv[1] - gv)) - - def ci(a, stat=np.median, n=1000, seed=0): - a = np.asarray(a) - rng = np.random.default_rng(seed) - bs = [stat(rng.choice(a, len(a), replace=True)) for _ in range(n)] - return np.percentile(bs, 2.5), np.percentile(bs, 97.5) - - def report(a, tag): - a = np.asarray(a) - mlo, mhi = ci(a, np.median) - alo, ahi = ci(a, np.mean) - print(f" {tag:13}: n={len(a)} median={np.median(a):.1f} [{mlo:.1f},{mhi:.1f}] " - f"mean={a.mean():.1f} [{alo:.1f},{ahi:.1f}] max={a.max():.0f}") - - print("\n=== pixel-goal L2 vs GT (640x480, held-out, 95% CI bootstrap) ===") - report(pt_l2, "PyTorch BF16") - # Validation gate uses the aspect-preserving PyTorch run; at 384-square the absolute is - # expectedly inflated (deploy resize) so the mapping is validated only in the native run. - if not sq: - med = np.median(pt_l2) - print(f" [gate] PyTorch median-L2 (aspect-preserving) = {med:.1f}px " - f"({'PASS <60, mapping valid' if med < 60 else 'CHECK'})") - if ENGINE and eng_l2: - report(eng_l2, os.path.basename(os.path.dirname(ENGINE))) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_system1.py b/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_system1.py deleted file mode 100644 index 381d623..0000000 --- a/recipes/internvla-n1-dualvln/trt-edgellm/verify/verify_system1.py +++ /dev/null @@ -1,142 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026 VinRobotics -# SPDX-License-Identifier: BSD-3-Clause -"""Verify System 1: wire the traj_dit + memory TensorRT engines into generate_traj and -compare the resulting trajectory and latency against the PyTorch base (same noise seed). -""" -import os -import sys -import time -_R = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, _R) -sys.path.insert(0, os.path.join(_R, "lib")) -sys.path.append("/usr/lib/python3.12/dist-packages") -sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "lib")) -import numpy as np # noqa: E402 -import torch # noqa: E402 -from PIL import Image # noqa: E402 - -ACTIVE = os.environ.get("INTERNNAV_PATH", os.path.expanduser("~/InternNav")) -CKPT = os.path.join(ACTIVE, "checkpoints/InternVLA-N1-DualVLN") -IMG = os.path.expanduser("~/modelopt/TensorRT-Edge-LLM/examples/multimodal/pics/giant_panda.jpeg") -TRAJDIT = os.path.join( - os.environ.get( - "WORK_DIR", - os.path.expanduser("~/vln-opt-work")), - "onnx/system1_traj_dit_bf16.engine") # noqa: E131 -MEM = os.path.join(os.environ.get("WORK_DIR", os.path.expanduser("~/vln-opt-work")), "onnx/system1_memory_bf16.engine") -SEED = 12345 - - -def out_of(o, key): - if isinstance(o, dict): - return o.get(key, next(iter(o.values()))) - return o[0] if isinstance(o, (list, tuple)) else o - - -def main(): - dev = "cuda" - try: - torch.backends.mha.set_fastpath_enabled(False) - except Exception: - pass - if ACTIVE not in sys.path: - sys.path.insert(0, ACTIVE) - from internvla_compat import apply_all - apply_all(need_system1=True, allow_missing_depth=True) - from internnav.model.basemodel.internvla_n1.internvla_n1 import ( - InternVLAN1ForCausalLM, InternVLAN1ModelConfig) - from diffusers.schedulers import FlowMatchEulerDiscreteScheduler - from diffusers.utils.torch_utils import randn_tensor - from trt_torch import Engine - print("[1/4] Load model + engines") - cfg = InternVLAN1ModelConfig.from_pretrained(CKPT) - model = InternVLAN1ForCausalLM.from_pretrained(CKPT, config=cfg, torch_dtype=torch.bfloat16, - attn_implementation="sdpa", low_cpu_mem_usage=True).to(dev).eval() - m = model.get_model() - mem_eng = Engine(MEM) - dit_eng = Engine(TRAJDIT) - NQ = model.get_n_query() - z = torch.randn(1, NQ, cfg.hidden_size, device=dev, dtype=torch.bfloat16) - a = np.array(Image.open(IMG).convert("RGB").resize((224, 224))) / 255.0 - tt = torch.from_numpy(a).float() - images_dp = torch.stack([tt, tt]).unsqueeze(0).to(dev) - rmean = model._resnet_mean - rstd = model._resnet_std - - def base_gen(): - torch.manual_seed(SEED) - np.random.seed(SEED) - with torch.no_grad(): - return model.generate_traj(z, images_dp, num_sample_trajs=32, num_inference_steps=10).float().cpu() - - def trt_gen(steps=10, ns=32, guidance_scale=1.0, predict_step_nums=32): - """Replicate generate_traj (async) with the memory + traj_dit engines.""" - torch.manual_seed(SEED) - np.random.seed(SEED) - dtype = z.dtype - with torch.no_grad(): - traj_latents = m.cond_projector(z) # [1,4,768] - # --- memory block through the engine --- - xdp = images_dp.permute(0, 1, 4, 2, 3) - xdp = ((xdp - rmean) / rstd).flatten(0, 1) # [2,3,224,224] - memory_tokens = out_of(mem_eng(images=xdp.float().contiguous()), "memory_tokens").to(dtype) # [1,32,768] - hidden_states = torch.cat([memory_tokens, traj_latents], dim=1) # [1,36,768] - hs_null = torch.zeros_like(hidden_states) - hs_input = torch.cat([hs_null, hidden_states], 0) # [2,36,768] - bs = traj_latents.shape[0] - latents = randn_tensor((bs * ns, predict_step_nums, 3), generator=None, device=dev, dtype=dtype) - sch = FlowMatchEulerDiscreteScheduler() - sigmas = np.linspace(1.0, 1 / steps, steps) - sch.set_timesteps(steps, sigmas=sigmas) - hs_input = hs_input.repeat_interleave(ns, dim=0) # [2*ns,36,768] - dit_eng.set_runtime_tensor_shape("z_latents", tuple(hs_input.shape)) - for t in sch.timesteps: - lf = m.action_encoder(latents) - pos_ids = torch.arange(lf.shape[1]).reshape(1, -1).repeat(bs * ns, 1).to(dev) - lf = lf + m.pos_encoding(pos_ids) - lmi = lf.repeat(2, 1, 1) - if hasattr(sch, "scale_model_input"): - lmi = sch.scale_model_input(lmi, t) - tt_ = t.unsqueeze(0).expand(lmi.shape[0]).to(dev, torch.long) - np_ = out_of(dit_eng(x=lmi.float().contiguous(), timestep=tt_.to(torch.int64).contiguous(), - z_latents=hs_input.float().contiguous()), "output").to(dtype) - np_ = m.action_decoder(np_) - unc, cnd = np_.chunk(2) - np_ = unc + guidance_scale * (cnd - unc) - latents = sch.step(np_, t, latents).prev_sample - return latents.float().cpu() - - def lat(fn, warm=2, n=5): - for _ in range(warm): - fn() - torch.cuda.synchronize() - ts = [] - for _ in range(n): - torch.cuda.synchronize() - t0 = time.perf_counter() - fn() - torch.cuda.synchronize() - ts.append(time.perf_counter() - t0) - return sum(ts) / len(ts) * 1000 - - print("[2/4] Base PyTorch generate_traj") - tref = base_gen() - pt_ms = lat(base_gen) - print(f" {tuple(tref.shape)} | {pt_ms:.1f}ms = {1000/pt_ms:.1f}Hz") - print("[3/4] Full-TRT S1 (memory+traj_dit engines)") - ttrt = trt_gen() - trt_ms = lat(trt_gen) - print(f" {tuple(ttrt.shape)} | {trt_ms:.1f}ms = {1000/trt_ms:.1f}Hz") - print("[4/4] Compare\n" + "=" * 56) - d = (tref - ttrt).norm(dim=-1) - endp = (tref[:, -1] - ttrt[:, -1]).norm(dim=-1) - cos = torch.nn.functional.cosine_similarity(tref.flatten(), ttrt.flatten(), dim=0).item() - print(f" Parity: per-wp L2 mean={d.mean():.4f} max={d.max():.4f} | endpoint={endp.mean():.4f} | cos={cos:.5f}") - print( - f" Latency S1: PyTorch {pt_ms:.1f}ms ({1000/pt_ms:.1f}Hz) → full-TRT {trt_ms:.1f}ms ({1000/trt_ms:.1f}Hz) = {pt_ms/trt_ms:.2f}x") # noqa: E501 - return 0 - - -if __name__ == "__main__": - sys.exit(main()) From 6efb73c8aca6745fff5245bdf8450f890084a71e Mon Sep 17 00:00:00 2001 From: hungho77 Date: Tue, 25 Aug 2026 02:14:46 +0700 Subject: [PATCH 30/30] fix(internvla-n1-dualvln): drop the dead experimental_models link, point at PR #193 The file link resolved against TensorRT-Edge-LLM's main branch, where PR #193 hasn't landed yet, so it 404s. Link the PR itself instead of a path that only exists on its branch. --- recipes/internvla-n1-dualvln/README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/recipes/internvla-n1-dualvln/README.md b/recipes/internvla-n1-dualvln/README.md index 388f0c5..3d5d872 100644 --- a/recipes/internvla-n1-dualvln/README.md +++ b/recipes/internvla-n1-dualvln/README.md @@ -66,9 +66,9 @@ System 1 (the trajectory expert) is not part of this checkpoint's quantization BF16, built separately with `trtexec` — and the async runtime (`internvla_n1_dual_system_inference` / `internvla_n1_dual_system_server`) that drives both systems is not part of this repo either; it ships with TensorRT-Edge-LLM. See -[`experimental_models/internvla_n1/README.md`](https://github.com/NVIDIA/TensorRT-Edge-LLM/blob/main/experimental_models/internvla_n1/README.md) -in that repo (or the PR branch, until it merges) for the full export → build → run flow and the -resident-server protocol for driving it from a Python agent. +[NVIDIA/TensorRT-Edge-LLM#193](https://github.com/NVIDIA/TensorRT-Edge-LLM/pull/193) for the +full export → build → run flow and the resident-server protocol for driving it from a Python +agent. ## Why calibration text needs no special tokens here