ubr-physical-ai/Cosmos3-Edge-INT4-AWQ

Cosmos3-Edge INT4-AWQ (Orin-targeted)

INT4 AWQ quantisation of nvidia/Cosmos3-Edge, produced for a Jetson Orin Nano deployment by Team UBR Stack (UB Robotics) during the NVIDIA / OpenHackathons / Oracle Open Models Codefest 2026.

This checkpoint has never been executed. See Validation status — read it before using this.

The Hub's auto-derived 8-bit tag is wrong — this is a 4-bit model. ModelOpt writes quant_algo: "W4A16_AWQ" with quant_method: "modelopt" and no bits field, so the Hub falls back to the safetensors dtype, which is U8 because 4-bit weights are packed into uint8. config.json is left exactly as ModelOpt wrote it — it is what the Edge-LLM toolchain consumes — so the tag cannot be corrected without risking the real loader. A consumer inferring bit-width the same way will make the same mistake. The authority is hf_quant_config.json: W4A16_AWQ, group size 128.

Read this before you run it: thinking is OFF by default, and that breaks structured output

onnx/llm/processed_chat_template.json ships both generation prompts:

"generation_prompt":          "<|im_start|>assistant\n<think></think>",   // default — thinking OFF
"generation_prompt_thinking": "<|im_start|>assistant\n<think>\n"          // thinking ON

The C++ runtime selects between them per request via request.enableThinking (tokenizer.cpp:1078, llmRankRuntime.cpp:86). The default is thinking-off, and on structured-output tasks that default produces malformed JSON.

Measured on a Jetson Orin Nano over a 24-item grounded benchmark (structured JSON with required fields, enums and nesting), by the team running the board:

setting	outputs that parse
baked default (thinking off)	11 / 24
enableThinking = true	24 / 24

Set enableThinking on the request if you are asking for JSON, a schema, or anything with syntax. Nothing needs rebuilding — the engine you already have supports both, and the prompts are in the file.

Do not read the default's output as the model being bad at JSON. It is a setting. Our own evaluation missed this because our benchmark answers are a single word or number, where there is no syntax to get wrong: on 200 paired spatial-QA items thinking changed nothing measurable (p = 0.60). That null result does not transfer to structured output, and this table is why.

Which runtime loads this

Not transformers, and not vLLM. A W4A16_AWQ checkpoint has GPTQ-packed weights whose shapes do not match the FP16 originals, so AutoModel.from_pretrained raises a shape-mismatch error; vLLM's ModelOpt backend accepts only FP8, FP8_PER_CHANNEL_PER_TOKEN, FP8_PB_WO, NVFP4, W4A16_NVFP4, MXFP8 and MIXED_PRECISION. Both were tried and both refused it.

The intended path is TensorRT-Edge-LLM v0.10.1:

this checkpoint  ->  tensorrt-edgellm-export  ->  ONNX  ->  tensorrt-edgellm-build  ->  engine (per device)

The exported ONNX is published alongside under onnx/ in this repo. The engine must be built on the target device, because TensorRT compiles for the exact GPU it will run on — an engine built elsewhere will not load on Orin. Build with the same Edge-LLM tag used here (v0.10.1, commit e8b2952); a tag mismatch across the export/build boundary is the first thing to rule out when a build fails.

What this is

Quantisation	W4A16_AWQ, group size 128, pre_quant_scale: true
Producer	NVIDIA ModelOpt 0.45.0 via TensorRT-Edge-LLM v0.10.1 (commit e8b2952)
KV cache	not quantised
Visual encoder	FP16 — excluded (model.visual*)
Size	2.24 GiB (2.41 GB decimal), from ~4.9 GB FP16
Reasoner parameters	2.44 B (the source repo's 4.56 B figure includes a bundled diffusion stack)

Why FP16 visual and an unquantised KV cache

Jetson Orin supports FP16, INT8 and INT4 at runtime — not FP8. On Edge-LLM v0.10.1 the --visual_quantization and --kv_cache_quantization flags accept only fp8, so both were omitted rather than set. The emitted hf_quant_config.json records the result: kv_cache_quant_algo: null and model.visual* excluded. Do not re-enable them for a device that cannot run FP8.

Calibration

AWQ is activation-aware, so the calibration data shapes the result. This checkpoint was calibrated on 198 real rover frames using the same posture prompt the robot runs, not a generic corpus:

• The frames are segment-disjoint from the evaluation slice — not merely frame-disjoint. A different frame of the same continuous run is the same place, lighting and people, and calibrating on it would flatter the model on the scenes used to judge it. Slice definitions: ubrobotics-ai/nvidia-codefest-2026 → data/slices/.
• Worth knowing: the quantiser runs image calibration for a multimodal model even when only --text_dataset is given, defaulting to lmms-lab/MMMU. Override with EDGELLM_QUANT_DATASET_MMMU=<local .jsonl>.
• The calibration images themselves are not published — they show identifiable people.

Validation status — read this

No accuracy number is published here, because the checkpoint could not be executed on the machine that made it.

• transformers cannot load it: W4A16_AWQ weights are GPTQ-packed and their shapes do not match the FP16 originals.
• vLLM does not support W4A16_AWQ under its ModelOpt backend (only FP8/NVFP4/MXFP8 variants).
• TensorRT was not available on the host.

Update, 2026-09-10: it now runs on a Jetson Orin Nano 8 GB. Measured on device — decode 18.92 ms/token, 52.8 tok/s; prefill 221.98 ms / 512 tokens; visual encoder 259.8 ms at 896x448; engine build 2m21s. Engines total 2.24 GiB on disk (llm.engine 838.9 MiB + visual.engine 942 MiB + embedding.safetensors 512 MiB), against the 2.3 GiB predicted from the ONNX. Peak memory 3,377 MB for the AWQ/TensorRT arm; a separate system-wide figure of ~4,830 MB is also reported for the full VLM and the two are being reconciled, so quote the scope alongside whichever number you use.

Update, 2026-09-11: accuracy is now measured, on two benchmarks.

Spatial reasoning, 1,442 items, synthetic warehouse scenes. Against the bf16 source model on identical items, this checkpoint's INT4 kernels are statistically indistinguishable once the distance calibration constant is re-fitted per arm: paired exact McNemar p = 0.931 on distance, p = 0.724 on left/right, p = 0.451 on region-choice. 4-bit shifts the distance scale, not the spatial judgement — the constant moves from 1.1546 at bf16 to 1.2148 here, and a run that carries the bf16 constant onto this checkpoint throws away 2.88 points for no reason. Fit k on your own engine's outputs: it is a property of the quantised artefact and of the engine build, and it drifts with range (the far bucket wants ~1.34 where a global fit gives 1.21).

Grounded detection, 24 real frames, 12 target instances. On a rescue-target benchmark it finds 9–10 of 12 targets at IoU 0.5 against 0–1 of 12 for the incumbent VLM it was compared with, with 0 of 14 false triggers on hard negatives against that model's 5. Reasoning must be on for this (request.enableThinking); it takes detections from 8 to 10. Note the strict scorer reads 0.0 for every arm here because it gates on schema validity — see the structured-output warning at the top of this card; the detections are real and are recovered by a tolerant parser.

At the time of the original export, the quantisation delta and the smoke test could not be run. What was verified is structural: the ONNX export passes onnx.checker on all components, with 169 Int4GroupwiseGemmPluginV2 nodes and 169 INT8 weight tensors (838.9 M elements) in the LLM graph, and none in the visual graph — the quantisation is where the config says it is. The exported fragments were also checked numerically: recomputing the cuteDSL fragment layout from the checkpoint's packed uint8 weights reproduces the ONNX tensors byte for byte.

The INT4 delta has since been measured by simulation (mtq.quantize() fake-quant in PyTorch — correct for weight precision, not for runtime). On 1,442 PhysicalAI Spatial-Intelligence-Warehouse validation items, paired McNemar against a same-environment bf16 control:

task	bf16	INT4 (simulated)	p	verdict
distance, raw	29.63%	21.81%	0.00016	real
distance, each arm with its own re-fitted k	38.75%	35.27%	0.163	not distinguishable
left / right	64.80%	67.20%	0.175	not distinguishable
mcq	17.76%	16.67%	0.551	not distinguishable

Once the distance-calibration constant is re-fitted per precision, no task shows a loss this benchmark can detect. k moves from 1.155 at bf16 to 1.211 at INT4 — carrying the bf16 value over invents a 7.8-point collapse that is not there. Re-fit it again on whatever your engine actually executes.

Deploying to Orin

The ONNX export yields six components, split across two pipelines. The VLM reasoner path is llm (1.4 GiB, INT4) · visual (943 MiB, FP16) · text_tokenizer (17 MiB) ≈ 2.3 GiB. und_prefill, gen and vae_encoder are the Cosmos3 policy path (COSMOS3_COMPONENTS = ("und_prefill", "gen", "vae_encoder") in models/cosmos3/export.py) and are not needed for perception. onnx/ here ships the three reasoner components plus und_prefill; gen and vae_encoder are not included.

Fit: 2.3 GiB is comfortable on an 8 GB Orin Nano and plausible, if tight, on a 4 GB one. The largest remaining item after the INT4 decoder is the FP16 vision encoder at 943 MiB; on tag v0.10.1 --visual_quantization accepts only fp8, which Orin cannot run, so it stays FP16 here.

Correction, 2026-09-10. An earlier upload of this repo carried a broken onnx/ export: a stale model.safetensors.index.json staged into the quantised directory made the loader read the diffusion and vision towers instead of the quantised weights, so the decoder shipped as dense FP16 with zero INT8 tensors while still carrying INT4 plugin nodes — 3.7 GB pretending to be INT4, unbuildable on a Jetson. It also quantised the multimodal projector, which breaks the visual encoder export. A third defect followed: the remap renames only the attention projections, so all 56 decoder MLP weights (mlp.fc1/mlp.fc2 vs mlp.up_proj/mlp.down_proj) never bound and shipped as all-zero INT8 tensors — every MLP a zero matrix, passing both onnx.checker and a dtype check. All three are fixed and each tensor is now verified byte-for-byte against the checkpoint, one per module class. If you pulled onnx/ before 2026-09-10 evening, re-pull. JETSON_ORIN_NANO_NOTES.md documents all three traps.

Other notes: llm consumes inputs_embeds, not input_ids, so the runtime performs embedding lookup itself; Int4GroupwiseGemmPluginV2 is a custom plugin, so the graph needs --plugin-path and will not run under stock ONNX Runtime. JETSON_ORIN_NANO_NOTES.md in this repo records seven staging gaps that each cost an hour.

Intended use

Perception and spatial reasoning on a Jetson Orin Nano for an offline-first search-and-rescue ground robot: given a camera frame and a question, answer in text about what is present and where. Vision-language in, text out.

Prototyping should use the unquantised nvidia/Cosmos3-Edge, which loads with transformers today. This repo is the deployment artefact.

Limitations

• Measured on two benchmarks, neither of them a live mission. The quantisation delta and the detection numbers above are bench evidence on 1,442 synthetic items and 24 real frames. Nothing here has run in a robot's control loop, and 24 frames is a small set — treat single-target differences as noise.
• The accuracy numbers were taken on two different GPUs. The spatial-reasoning result was measured on a DGX B300 (SM103) under patches not yet merged upstream; the detection result on a Jetson Orin Nano (SM87). Engines are per-SM: the same ONNX export produces different plans on the two parts, so an engine-equivalence check is owed before an accuracy claim is transferred between them. It has not been run.
• Requires the Edge-LLM ONNX path plus an engine built on the target device. It is not loadable by transformers, vLLM, or stock ONNX Runtime — the LLM graph uses a custom Int4GroupwiseGemmPluginV2 plugin.
• Calibrated on one narrow domain. The 198 calibration frames are indoor urban-search-and-rescue footage from a single robot, with one posture prompt. AWQ fits activation ranges to what it is shown, so behaviour on outdoor scenes, other camera geometries, other lighting, or different prompt styles is unmeasured. The in-domain calibration is deliberate and is also a real narrowing.
• Modest capability even before quantisation. The bf16 source model scores 64.8% on a two-way left/right question (chance 50%) and 17.8% on a region-choice question whose true chance level is 14.2% — that task offers between 3 and 13 regions per item, so chance is the mean of 1/n_regions, not 1/4 or 1/8. The mcq margin over chance is therefore about 3 points, not 5. It is better than the alternative we benchmarked, and it is not a reliable spatial reasoner.
• It also carries a directional prior. It answers left on 72% of a perfectly balanced 250/250 left-right set. It still clears chance by 15 points, so the bias is costing it accuracy it already has.
• und_prefill is shipped but is not part of the reasoner. It is a Cosmos3 policy component, grouped with gen and vae_encoder in the exporter. It is unquantised FP16 at 3.7 GiB and does not load on the VLM path. Earlier versions of this card counted it in the reasoner payload; that was wrong.

Licence

OpenMDW 1.1, inherited from nvidia/Cosmos3-Edge. This is a derivative work; the upstream terms govern.