| --- |
| license: other |
| license_name: nvidia-oneway-noncommercial |
| license_link: https://huggingface.co/nvidia/NitroGen/blob/main/LICENSE |
| tags: |
| - onnx |
| - nitrogen |
| - gaming |
| - fp16 |
| - cuda |
| - tensorrt |
| - rtx-2060 |
| - turing |
| base_model: nvidia/NitroGen |
| --- |
| |
| # 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 |
| powershell -ExecutionPolicy Bypass -File scripts/setup_windows.ps1 |
| ``` |
|
|
| Then benchmark the ONNX DiT: |
|
|
| ```powershell |
| .\.venv\Scripts\python.exe scripts\benchmark_dit.py |
| ``` |
|
|
| For the first gameplay test, start with 4 denoise steps: |
|
|
| ```powershell |
| .\scripts\test_game.ps1 -Process "YourGame.exe" -Steps 4 |
| ``` |
|
|
| Or run the server manually: |
|
|
| ```powershell |
| .\.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. |
|
|
