patdev/NitroGen-RTX2060-ONNX

NitroGen RTX 2060 ONNX

Turing-oriented ONNX acceleration package for NVIDIA NitroGen. Target machine: RTX 2060 6 GB / Windows 11. The repeated 122M-parameter DiT denoise block is exported natively to ONNX FP16 (opset 20) while NitroGen's control loop, tokenizer, vision tower and small action heads remain upstream PyTorch. The gameplay loader skips the PyTorch DiT weights entirely on GPU to reduce peak VRAM.

This is intentionally hybrid rather than a monolithic export: NitroGen calls the DiT 4-16 times per action chunk, so moving this repeated block to ONNX Runtime gives the useful acceleration while preserving upstream behavior.

Artifact

• onnx/dit_fp16.onnx — 233 MiB, native FP16 export
• input hidden_states: [B,18,1024] FP16
• input encoder_hidden_states: [B,256,1024] FP16
• input timestep: [B] INT64
• output pred: [B,18,1024] FP16

The source checkpoint is nvidia/NitroGen/ng.pt and is downloaded from the official repository at runtime. Its NVIDIA non-commercial license continues to apply.

Why this profile for RTX 2060

RTX 2060 (Turing, SM 7.5) has fast FP16 Tensor Cores but no native FP8/NVFP4 path. The existing FP8/NVFP4 NitroGen exports target later NVIDIA hardware. This package therefore uses FP16 and defaults to:

1. TensorRT Execution Provider when already available
2. CUDA Execution Provider otherwise
3. CPU only for synthetic validation/benchmark

The gameplay session uses FP16 autocast instead of the upstream BF16 autocast, which is a better match for Turing.

Windows setup

powershell -ExecutionPolicy Bypass -File scripts/setup_windows.ps1

Then benchmark the ONNX DiT:

.\.venv\Scripts\python.exe scripts\benchmark_dit.py

For the first gameplay test, start with 4 denoise steps:

.\scripts\test_game.ps1 -Process "YourGame.exe" -Steps 4

Or run the server manually:

.\.venv\Scripts\python.exe scripts\serve_onnx.py --steps 4
.\.venv\Scripts\python.exe .vendor\NitroGen\scripts\play.py --process "YourGame.exe"

Try --steps 8 if action quality is noticeably worse; 16 reproduces the checkpoint's original denoise-step count but is slower. The setup pins upstream NitroGen commit 32608444660950ffda95e1e57c79632ad65bea10 so future upstream changes do not silently break the wrapper.

Validation performed during export

The graph passes onnx.checker. A synthetic DiT comparison against the source PyTorch weights produced cosine similarity 0.99994. On the CPU build machine, one synthetic DiT step measured roughly 404 ms PyTorch FP16 vs 20 ms ONNX Runtime; that CPU number is only an export sanity check, not an RTX 2060 performance claim. Run benchmark_dit.py locally for the real GPU number.

Notes

• NitroGen itself is a research model and is strongest on gamepad-oriented games.
• TensorRT engine caches are hardware/runtime specific and are generated locally under .ort-cache; they are not uploaded as portable .plan files.
• onnxruntime-gpu is capped below 1.27 in the setup script because current ORT 1.27+ PyPI GPU packages default to CUDA 13, while the 1.21-1.26 line uses CUDA 12.8 and is typically simpler on Turing/Windows installations.