fix: onnx export, consolidate the model onto one device.

unidrive_vla_nusc_base_quantize_all.py

@@ -429,11 +429,11 @@ def quantize_nvfp4(
     onnx_path = os.path.join(output_dir, "unidrive_vla_nvfp4.onnx")
     try:
         # Build a single sample input from the first calibration item.
-        # export_to_onnx expects processed tensors, not raw messages metadata.
+        # export_to_onnx expects a list of processed input dicts and uses [0].
         sample_input = _process_calibration_batch(
             calibration_items[:1], processor, device
         )
-        onnx_path = export_to_onnx(quantized_model, sample_input, onnx_path)
+        onnx_path = export_to_onnx(quantized_model, [sample_input], onnx_path)
     except Exception as e:
         print(f"  [WARN] ONNX export failed: {e}")
         print(f"  The quantized PyTorch model is still saved at: {output_path}.")
@@ -473,24 +473,22 @@ def export_to_onnx(
     # Use the first calibration sample as dummy input
     sample_input = calibration_inputs[0]
 
-    # --- Consolidate the model onto CPU before tracing. ---
-    # torch.onnx.export(dynamo=False) traces via torch.jit. Tracing on GPU works
-    # for the forward, but the bf16 weights + trace activations don't fit in
-    # small-VRAM cards, and jit-trace of a model carrying accelerate hooks is
-    # unreliable. Consolidating onto one device avoids both. CPU is used because
-    # the NVFP4 ONNX export path is device-agnostic —
-    # configure_linear_module_onnx_quantizers only flags quantizers to emit graph
-    # nodes (no cuda kernels), and modelopt's own get_onnx_bytes_and_metadata
-    # disables cuda autocast for fp4 models — so tracing on CPU yields the same
-    # ONNX graph. (The model is loaded single-device now, but we still strip any
-    # residual accelerate hooks defensively and move to CPU explicitly.)
+    # --- Consolidate the model onto one device before tracing. ---
+    # NVFP4 dynamic block-quantization requires amax to be a CUDA tensor, so
+    # tracing must run on GPU. On cards with limited VRAM this may OOM; however,
+    # the UniDriveVLA base model (~2.1B params) fits comfortably on an H20,
+    # so we default to cuda and only fall back to CPU if GPU tracing fails.
     try:
         from accelerate.hooks import remove_hook_from_module
         if hasattr(model, "hf_device_map"):
             remove_hook_from_module(model, recurse=True)
     except ImportError:
         pass  # no accelerate -> model is already single-device
-    export_device = torch.device("cpu")
+
+    if torch.cuda.is_available():
+        export_device = torch.device("cuda")
+    else:
+        export_device = torch.device("cpu")
     model.to(export_device)
     sample_input = {
         k: (v.to(export_device) if torch.is_tensor(v) else v)