feat: run:ai gpu_direct model loading with quantization-aware device maps

[khaiwang]

Summary

• Add run:ai SafetensorsStreamer as the default loading backend for LanguageModel, with silent fallback to from_pretrained when runai-model-streamer is not installed — zero behavior c
  hange for existing users
• gpu_direct: copies tensors directly from Run:AI's I/O buffer to the target GPU during streaming, making HF's subsequent .to(device) a no-op
• Quantization-aware _resolve_device_map: creates HfQuantizer and runs preprocess_model on the meta model before computing the device map, so compute_module_sizes sees correct quantized parameter shapes. Without this, MXFP4/GPT
  Q models get 4× memory overestimates, causing unnecessary CPU offloading on tight GPU configs
• Includes benchmark suite (benchmark_loading.py) and documentation (README.md)

Benchmark Results (cold cache, 8× A100-80GB PCIe)

Model	Type	Size	HF (s)	gpu_direct (s)	Speedup
Qwen/Qwen3-8B	Dense (2 GPU)	15.3 GB	10.3	5.0	2.0×
Qwen/Qwen3-30B-A3B	MoE (8 GPU)	56.9 GB	49.3	23.8	2.1×
openai/gpt-oss-120b	MXFP4 quantized (8 GPU)	60.8 GB	49.7	17.7	2.8×

How it works

HF baseline:       disk → mmap page faults → CPU → HF .to(cuda) → GPU
Run:AI gpu_direct:  disk → run:ai buffer → .to(cuda) → GPU → HF .to() [no-op]

The Run:AI streamer uses N concurrent C++ pthreads for disk I/O (no GIL contention). The loader resolves device_map upfront, then copies each tensor to its target GPU as it arrives from the streamer. HF's _materialize_copy().to() becom
es a device/dtype no-op. The streamer's background reads overlap naturally with the blocking .to(cuda) calls.

Usage

from nnsight import LanguageModel

# Default: uses run:ai gpu_direct (silently falls back to from_pretrained)
model = LanguageModel("openai/gpt-oss-120b", device_map="auto", dispatch=True)

# Force HF from_pretrained
model = LanguageModel("openai/gpt-oss-120b", load_format="from_pretrained",
                      device_map="auto", dispatch=True)

# Tune Run:AI I/O concurrency (default 16)
model = LanguageModel("openai/gpt-oss-120b", concurrency=4,
                      device_map="auto", dispatch=True)

Files changed

File	Description
src/nnsight/modeling/loader.py	New — RunAIShardCache, LazyRunAITensor, build_lazy_state_dict
src/nnsight/modeling/transformers.py	_load dispatch, _resolve_device_map with quantizer, _load_streamed
tests/performance/loading/benchmark_loading.py	Benchmark suite (hf vs gpu_direct, page cache control, correctness verification)
tests/performance/loading/README.md	Architecture docs, usage, benchmark results

Test plan

• pytest tests/test_tiny.py --device cpu — 18/18 pass
• Dense model (Qwen3-8B): correctness + performance verified on 2 GPUs
• MoE model (Qwen3-30B-A3B): correctness + performance verified on 8 GPUs
• MXFP4 quantized (gpt-oss-120b): correctness + performance verified on 4 and 8 GPUs
• MXFP4 quantized (gpt-oss-20b): correctness verified on 1 GPU — validates quantization-aware device map fix
• Fallback: loads correctly without runai-model-streamer installed (uses from_pretrained)

[greptile-apps]

Your free trial has ended. If you'd like to continue receiving code reviews, you can add a payment method here.