diff --git a/src/nnsight/modeling/loader.py b/src/nnsight/modeling/loader.py new file mode 100644 index 00000000..f0949815 --- /dev/null +++ b/src/nnsight/modeling/loader.py @@ -0,0 +1,375 @@ +"""Streaming model weight loading via run:ai SafetensorsStreamer.""" + +import json +import os +import struct +import time +import threading +from pathlib import Path + +import torch + +# -- Safetensors dtype string → torch.dtype ----------------------------------- + +_SAFETENSORS_DTYPE = { + "F64": torch.float64, + "F32": torch.float32, + "F16": torch.float16, + "BF16": torch.bfloat16, + "I64": torch.int64, + "I32": torch.int32, + "I16": torch.int16, + "I8": torch.int8, + "U8": torch.uint8, + "BOOL": torch.bool, +} + +_FLOATING_DTYPES = frozenset( + {torch.float16, torch.bfloat16, torch.float32, torch.float64} +) + +def resolve_shard_paths(repo_id: str, revision: str = "main") -> list[str]: + """Resolve local .safetensors shard paths from HF cache. + + Raises ValueError if no .safetensors files found (no .bin fallback — + SafetensorsStreamer only handles safetensors format). + """ + from huggingface_hub import snapshot_download + + model_dir = snapshot_download(repo_id, revision=revision, local_files_only=True) + paths = sorted(Path(model_dir).glob("*.safetensors")) + if not paths: + raise ValueError( + f"No .safetensors files found for {repo_id} (rev={revision}). " + f"SafetensorsStreamer requires safetensors format." + ) + return [str(p) for p in paths] + + +# -- Shard-by-shard streaming ------------------------------------------------- + + +def _parse_safetensors_keys( + shard_paths: list[str], +) -> tuple[dict[str, tuple[str, list[int], str]], dict[str, int]]: + """Parse safetensors JSON headers to build key→metadata mapping. + + Reads only the 8-byte length prefix + JSON header per file — no tensor + data is touched. ~0.2 ms per file vs ~130 ms for ``safe_open`` on + networked FS. + + Returns: + key_map: ``{key: (shard_path, shape, dtype_str)}`` + shard_key_counts: ``{shard_path: num_keys}`` + """ + key_map: dict[str, tuple[str, list[int], str]] = {} + shard_key_counts: dict[str, int] = {} + + for path in shard_paths: + with open(path, "rb") as f: + (header_size,) = struct.unpack(" None: + self._concurrency = concurrency + self._device_map = device_map + self._torch_dtype = torch_dtype + self._gpu_direct = ( + device_map is not None and torch.cuda.is_available() + ) + self._lock = threading.Lock() + self._condition = threading.Condition(self._lock) + + # Per-tensor storage (flat, not grouped by shard) + self._tensors: dict[str, torch.Tensor] = {} + + # Shard-level tracking + self._refcounts: dict[str, int] = {} + self._shard_loading: set[str] = set() + self._errors: dict[str, Exception] = {} + + # Profiling stats (thread-safe via _lock/_condition) + self.stats_shard_wall_s = 0.0 # wall time of _stream_shard_incremental + self.stats_io_wait_s = 0.0 # time waiting for streamer (uncovered I/O) + self.stats_clone_s = 0.0 # time cloning tensors (CPU-target path) + self.stats_gpu_copy_s = 0.0 # time in buffer→GPU .to() copies + self.stats_notify_s = 0.0 # time holding lock for store+notify + self.stats_consumer_wait_s = 0.0 # cumulative consumer wait on Condition + self.stats_pop_count = 0 # tensors served + + def register(self, shard_path: str, num_keys: int) -> None: + """Set the initial refcount for a shard (= number of keys it holds).""" + self._refcounts[shard_path] = num_keys + + def get(self, shard_path: str, key: str) -> torch.Tensor: + """Return the tensor for *key*, streaming the shard if needed.""" + should_load = False + t_wait_start = time.perf_counter() + + with self._condition: + while True: + if shard_path in self._errors: + raise self._errors[shard_path] + + if key in self._tensors: + self.stats_consumer_wait_s += time.perf_counter() - t_wait_start + return self._pop_tensor(shard_path, key) + + if shard_path not in self._shard_loading: + self._shard_loading.add(shard_path) + should_load = True + break + + # Wait for notification (spurious wakeups handled by loop) + self._condition.wait() + + if should_load: + try: + self._stream_shard_incremental(shard_path) + except Exception as e: + with self._condition: + self._errors[shard_path] = e + self._condition.notify_all() + raise + + # Loader's own tensor is now available + with self._condition: + if shard_path in self._errors: + raise self._errors[shard_path] + return self._pop_tensor(shard_path, key) + + # -- internals ------------------------------------------------------------- + + def _pop_tensor(self, shard_path: str, key: str) -> torch.Tensor: + """Pop *key* from the cache and clean up shard tracking if exhausted.""" + tensor = self._tensors.pop(key) + self._refcounts[shard_path] -= 1 + self.stats_pop_count += 1 + if self._refcounts[shard_path] == 0: + del self._refcounts[shard_path] + self._shard_loading.discard(shard_path) + return tensor + + def _resolve_device(self, checkpoint_key: str) -> torch.device: + """Map a checkpoint key to its target device via the device_map.""" + from transformers.integrations.accelerate import expand_device_map + + expanded = expand_device_map(self._device_map, [checkpoint_key]) + device = expanded[checkpoint_key] + if isinstance(device, int): + return torch.device("cuda", device) + if device == "disk": + return torch.device("cpu") + return torch.device(device) + + def _resolve_dtype(self, tensor: torch.Tensor) -> torch.dtype | None: + """Return target dtype for a tensor, or None to keep original. + + Only floating-point tensors are cast; integer/bool buffers are + left as-is. + """ + if self._torch_dtype is None: + return None + if tensor.dtype in _FLOATING_DTYPES: + return self._torch_dtype + return None + + @staticmethod + def _enable_expandable_segments() -> None: + """Enable CUDA expandable segments to avoid fragmentation. + + Streaming hundreds of individually-sized tensors to GPU causes + fragmentation in the default CUDA caching allocator. Expandable + segments let the allocator grow incrementally instead of reserving + large fixed blocks, with no measurable performance penalty. + """ + key = "PYTORCH_CUDA_ALLOC_CONF" + conf = os.environ.get(key, "") + if "expandable_segments" not in conf: + entry = "expandable_segments:True" + os.environ[key] = f"{conf},{entry}" if conf else entry + + def _stream_shard_incremental(self, shard_path: str) -> None: + """Stream a shard via run:ai, storing and notifying per tensor.""" + from runai_model_streamer import SafetensorsStreamer + + os.environ["RUNAI_STREAMER_CONCURRENCY"] = str(self._concurrency) + + # Avoid CUDA memory fragmentation from streaming many varied-size tensors + if self._gpu_direct: + self._enable_expandable_segments() + + t_shard_start = time.perf_counter() + t_iter_start = t_shard_start + with SafetensorsStreamer() as streamer: + streamer.stream_files([shard_path], device="cpu") + for name, tensor in streamer.get_tensors(): + t_yield = time.perf_counter() + io_wait = t_yield - t_iter_start + + if self._gpu_direct: + target_device = self._resolve_device(name) + target_dtype = self._resolve_dtype(tensor) + if target_device.type == "cuda": + # Direct buffer → GPU (one copy, no intermediate clone) + result = tensor.to( + device=target_device, + dtype=target_dtype or tensor.dtype, + ) + t_copied = time.perf_counter() + with self._condition: + self.stats_io_wait_s += io_wait + self.stats_gpu_copy_s += t_copied - t_yield + self._tensors[name] = result + self._condition.notify_all() + else: + # CPU/disk target — must clone since buffer is transient + cloned = tensor.clone() + t_cloned = time.perf_counter() + if target_dtype is not None: + cloned = cloned.to(dtype=target_dtype) + with self._condition: + self.stats_io_wait_s += io_wait + self.stats_clone_s += t_cloned - t_yield + self._tensors[name] = cloned + self._condition.notify_all() + else: + # CPU-only path (no device_map or no CUDA) + cloned = tensor.clone() + t_cloned = time.perf_counter() + with self._condition: + self.stats_io_wait_s += io_wait + self.stats_clone_s += t_cloned - t_yield + self._tensors[name] = cloned + self._condition.notify_all() + + t_iter_start = time.perf_counter() + shard_wall = time.perf_counter() - t_shard_start + + with self._condition: + self.stats_shard_wall_s += shard_wall + + +class LazyRunAITensor: + """Drop-in replacement for ``SafetensorSlice``. + + Implements the minimal interface that transformers' + ``convert_and_load_state_dict_in_model`` expects from state-dict values: + ``__getitem__``, ``.dtype``, ``.shape``, ``.get_shape()``, and + ``.is_floating_point()``. + + The actual tensor is streamed on the first (and only) ``__getitem__`` + call via the shared :class:`RunAIShardCache`. + """ + + def __init__( + self, + shard_path: str, + key: str, + cache: RunAIShardCache, + shape: list[int], + dtype_str: str, + ) -> None: + self._shard_path = shard_path + self._key = key + self._cache = cache + self._shape_list = shape + self._dtype = _SAFETENSORS_DTYPE[dtype_str] + + def __getitem__(self, idx): + tensor = self._cache.get(self._shard_path, self._key) + return tensor[idx] + + def to(self, *args, **kwargs): + """Materialize the tensor and move it to the target device/dtype.""" + tensor = self._cache.get(self._shard_path, self._key) + return tensor.to(*args, **kwargs) + + @property + def dtype(self) -> torch.dtype: + return self._dtype + + @property + def shape(self) -> torch.Size: + return torch.Size(self._shape_list) + + def get_shape(self) -> list[int]: + return list(self._shape_list) + + def is_floating_point(self) -> bool: + return self._dtype in _FLOATING_DTYPES + + +def build_lazy_state_dict( + shard_paths: list[str], + concurrency: int = 16, + device_map: dict | None = None, + torch_dtype: torch.dtype | None = None, +) -> dict[str, LazyRunAITensor]: + """Build a lazy state dict backed by run:ai streaming. + + Values are :class:`LazyRunAITensor` instances that stream the + underlying safetensors shard on first ``__getitem__`` access. Shards + are evicted from memory once all their keys have been consumed, + keeping peak memory to ~1 tensor instead of the full model. + + Args: + shard_paths: Paths to ``.safetensors`` shard files. + concurrency: Passed to ``RUNAI_STREAMER_CONCURRENCY``. + device_map: Resolved device map (param name → device). When + provided with GPU targets, the streaming cache copies tensors + directly from the Run:AI buffer to GPU, so that HF's + ``_materialize_copy().to()`` is a no-op. + torch_dtype: Target dtype for floating-point tensors. Applied + during GPU transfer so HF's dtype cast is also a no-op. + + Returns: + Dict mapping checkpoint key names to :class:`LazyRunAITensor`. + """ + # Fail fast if run:ai is not installed. + from runai_model_streamer import SafetensorsStreamer # noqa: F401 + + key_map, shard_key_counts = _parse_safetensors_keys(shard_paths) + + cache = RunAIShardCache( + concurrency, device_map=device_map, torch_dtype=torch_dtype, + ) + for shard_path, count in shard_key_counts.items(): + cache.register(shard_path, count) + + return { + key: LazyRunAITensor(shard_path, key, cache, shape, dtype_str) + for key, (shard_path, shape, dtype_str) in key_map.items() + } diff --git a/src/nnsight/modeling/transformers.py b/src/nnsight/modeling/transformers.py index d6385a12..85dcaa04 100755 --- a/src/nnsight/modeling/transformers.py +++ b/src/nnsight/modeling/transformers.py @@ -75,10 +75,177 @@ def _load( **kwargs, ) -> PreTrainedModel: + load_format = kwargs.pop("load_format", None) + gpu_direct = kwargs.pop("gpu_direct", True) + concurrency = kwargs.pop("concurrency", 16) self._load_config(repo_id, revision=revision, **kwargs) + # Tensor parallelism: skip the run:ai streamer entirely. + # With single-node TP each rank mmap's the same files and the OS + # page cache shares the physical pages (1× read for N ranks). + # RunAI's per-rank independent streaming reads N× the data with no + # page sharing, measured 2× slower than mmap for Qwen3-8B on 2 GPUs. + # gpu_direct is also incompatible: TP needs torch.narrow() on CPU + # before GPU placement, so full-tensor GPU copies would require an + # extra GPU→CPU round-trip. + if kwargs.get("tp_plan") is not None: + load_format = "from_pretrained" + + # Default: try run:ai streamer, fall back to from_pretrained if not installed + if load_format != "from_pretrained": + try: + model = self._load_streamed( + repo_id, revision=revision, + gpu_direct=gpu_direct, concurrency=concurrency, + **kwargs, + ) + self.config = model.config + return model + except ImportError: + if load_format == "runai_streamer": + raise # explicit request — don't swallow the error + # else load_format is None (default) — fall through silently + except ValueError: + if load_format == "runai_streamer": + raise + # No .safetensors files (e.g. old repos with only .bin) — fall through + model = self.automodel.from_pretrained(repo_id, revision=revision, **kwargs) self.config = model.config return model + + def _resolve_device_map(self, device_map_str: str, max_memory=None, + torch_dtype=None) -> dict: + """Expand a string device_map ('auto', 'balanced', etc.) to a dict. + + Creates a throwaway meta-device model to compute the layer→device + mapping, then discards it. This lets us know each tensor's target + GPU *before* streaming begins. + + ``torch_dtype`` sets the default dtype during meta-model creation + so that memory estimates match the actual loading dtype (e.g. + bfloat16 vs float32). + + When the model config contains a ``quantization_config``, the + corresponding ``HfQuantizer`` is created and its + ``preprocess_model`` is called to replace modules with quantized + variants (e.g. ``Mxfp4GptOssExperts``). This ensures + ``compute_module_sizes`` sees the smaller quantized parameter + shapes rather than overestimating at full-precision sizes. + """ + from transformers.integrations.accelerate import _get_device_map + import torch + + # Build quantizer so device-map sizing accounts for quantized weights + hf_quantizer = None + if getattr(self.config, "quantization_config", None) is not None: + try: + from transformers.quantizers.auto import AutoHfQuantizer + hf_quantizer = AutoHfQuantizer.from_config( + self.config.quantization_config, pre_quantized=True, + ) + except Exception: + pass # Missing deps — fall back to unquantized sizing + + model_class = self.automodel._model_mapping[type(self.config)] + old_dtype = torch.get_default_dtype() + effective_dtype = torch_dtype or getattr(self.config, "torch_dtype", None) + if effective_dtype is not None: + torch.set_default_dtype(effective_dtype) + try: + with torch.device("meta"): + meta_model = model_class(self.config) + # Run quantizer's module replacement (e.g. GptOssExperts → + # Mxfp4GptOssExperts) so that compute_module_sizes sees the + # smaller quantized parameter shapes. This mirrors what + # from_pretrained does before calling _get_device_map. + if hf_quantizer is not None: + try: + hf_quantizer.preprocess_model( + meta_model, dtype=effective_dtype, + ) + except Exception: + pass # Non-fatal — sizes will be overestimated + resolved = _get_device_map( + meta_model, device_map_str, max_memory, hf_quantizer=hf_quantizer, + ) + finally: + torch.set_default_dtype(old_dtype) + del meta_model + return resolved + + def _load_streamed( + self, + repo_id: str, + revision: Optional[str] = None, + concurrency: int = 16, + gpu_direct: bool = True, + **kwargs, + ) -> PreTrainedModel: + """Load model using run:ai SafetensorsStreamer for fast disk I/O. + + Builds a lazy state dict whose values stream incrementally on + first ``__getitem__`` access — peak CPU memory is ~1 tensor + instead of the full model. + + When *gpu_direct* is True (default) and ``device_map`` is a + string like ``"auto"``, it is resolved to a concrete dict + *before* building the lazy state dict so that the streaming + cache can copy tensors directly to their target GPU, making + HF's ``_materialize_copy().to()`` a no-op. + + When *gpu_direct* is False, tensors are cloned to CPU and HF + workers handle the GPU transfer (the pre-GPU-direct path). + + ``from_pretrained(None, state_dict=...)`` handles weight + renaming, conversion, dtype casting, device placement, and + tied weight resolution. + """ + from .loader import ( + resolve_shard_paths, + build_lazy_state_dict, + ) + + shard_paths = resolve_shard_paths(repo_id, revision=revision) + + # tp_plan and device_map are mutually exclusive in from_pretrained. + tp_plan = kwargs.get("tp_plan") + + device_map = kwargs.pop("device_map", None) + resolved_device_map = None + + if gpu_direct: + # Resolve device_map early so the cache can place tensors on GPU + if isinstance(device_map, str) and device_map in ( + "auto", "balanced", "balanced_low_0", "sequential", + ): + resolved_device_map = self._resolve_device_map( + device_map, max_memory=kwargs.get("max_memory"), + torch_dtype=kwargs.get("torch_dtype"), + ) + elif isinstance(device_map, dict): + resolved_device_map = device_map + + state_dict = build_lazy_state_dict( + shard_paths, concurrency=concurrency, + device_map=resolved_device_map, + torch_dtype=kwargs.get("torch_dtype") if gpu_direct else None, + ) + + # Resolve concrete model class — Auto classes reject None as path + model_class = self.automodel._model_mapping[type(self.config)] + + # Only pass device_map when not using TP (they're mutually exclusive) + if tp_plan is None: + kwargs["device_map"] = resolved_device_map or device_map + + model = model_class.from_pretrained( + None, + config=self.config, + state_dict=state_dict, + revision=revision, + **kwargs, + ) + return model diff --git a/tests/performance/loading/README.md b/tests/performance/loading/README.md new file mode 100644 index 00000000..2e5098e3 --- /dev/null +++ b/tests/performance/loading/README.md @@ -0,0 +1,104 @@ +# Run:AI Model Loading with GPU-Direct Tensor Placement + +Fast model loading for nnsight's `LanguageModel` via [run:ai model streamer](https://github.com/run-ai/runai-model-streamer) with direct buffer-to-GPU tensor placement, bypassing CPU intermediaries. + +## Loading Paths + +### 1. HF `from_pretrained` (baseline) + +Standard HuggingFace loading: safetensors files are memory-mapped, tensors are materialized via 4 KB page faults, then HF worker threads call `.to(device)` for GPU placement. Bottlenecked by mmap read granularity and synchronous `cudaMemcpyAsync` blocking on each worker thread. + +### 2. Run:AI GPU-Direct (default) + +The loader thread resolves `device_map` before streaming begins, then copies tensors directly from the Run:AI buffer to the target GPU via `.to(device=cuda:N, dtype=dtype)`. When HF workers later call `.to()`, it's a no-op since the tensor is already on the correct device and dtype. The Run:AI streamer's background I/O threads read the next shard from disk concurrently with the blocking `.to()`, providing natural overlap between disk reads and GPU transfers. + +``` +HF baseline: disk → mmap page faults → CPU → HF .to(cuda) → GPU +Run:AI GPU-direct: disk → run:ai buffer → .to(cuda) → GPU cache → HF .to() [no-op] +``` + +## Architecture + +``` +┌─────────────────────────────────────────────────────────────┐ +│ TransformersModel._load_streamed() │ +│ │ +│ 1. resolve_shard_paths() → list of .safetensors paths │ +│ 2. _resolve_device_map() → {"model.layers.0": 0, ...} │ +│ 3. build_lazy_state_dict() → {key: LazyRunAITensor, ...} │ +│ 4. from_pretrained(None, state_dict=...) │ +└─────────────────────────────────────────────────────────────┘ + +┌─────────────────────────────────────────────────────────────┐ +│ RunAIShardCache (loader thread) │ +│ │ +│ SafetensorsStreamer.get_tensors() yields (name, buffer) │ +│ ├─ GPU target? → tensor.to(device=cuda:N, dtype=dtype) │ +│ └─ CPU target? → tensor.clone() │ +│ Store in _tensors[name], notify_all() waiting workers │ +└─────────────────────────────────────────────────────────────┘ + +┌─────────────────────────────────────────────────────────────┐ +│ HF GLOBAL_WORKERS threads │ +│ │ +│ _materialize_copy(lazy_tensor, device, dtype) │ +│ → lazy_tensor[...] → cache.get() → returns GPU tensor │ +│ → tensor.to(device, dtype) → no-op (already there) │ +└─────────────────────────────────────────────────────────────┘ +``` + +### Key Implementation Details + +- **`_resolve_device_map` dtype fix:** The meta model used for device map resolution is created with the model's native dtype (from config) instead of float32 default. Without this, large models (e.g., Qwen3-32B at 61 GB in BF16 vs 122 GB in FP32) get incorrect memory estimates and trigger disk offloading. + +- **Quantization-aware device map:** When the model config contains `quantization_config`, the loader creates an `HfQuantizer` and runs `preprocess_model` on the meta model before computing the device map. This ensures `compute_module_sizes` sees the correct quantized parameter shapes (e.g., MXFP4 expert weights) instead of overestimating at full-precision sizes. + +- **`expandable_segments`:** GPU-direct streaming allocates hundreds of individually-sized tensors on GPU, which can cause CUDA memory fragmentation. The loader automatically enables `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` on first GPU-direct use. + +### Key Files + +| File | Role | +|---|---| +| `src/nnsight/modeling/loader.py` | `RunAIShardCache`, `LazyRunAITensor`, `build_lazy_state_dict` | +| `src/nnsight/modeling/transformers.py` | `_load_streamed`, `_resolve_device_map`, `gpu_direct` flag | + +### User-Facing API + +```python +from nnsight import LanguageModel + +# GPU-direct (default when run:ai is installed) +model = LanguageModel("Qwen/Qwen3-8B", device_map="auto", dispatch=True) + +# Force HF from_pretrained +model = LanguageModel("Qwen/Qwen3-8B", device_map="auto", dispatch=True, + load_format="from_pretrained") + +# Tune Run:AI I/O concurrency (default 16) +model = LanguageModel("Qwen/Qwen3-8B", device_map="auto", dispatch=True, + concurrency=32) +``` + +If `runai-model-streamer` is not installed, loading falls back to `from_pretrained` silently. + +## Benchmark Results + +### 8× A100-80GB PCIe, RAID /dev/md0, cold cache, best concurrency config + +| Model | Type | Size | HF (s) | gpu_direct (s) | Speedup | +|---|---|---|---|---|---| +| Qwen/Qwen3-8B | Dense (2 GPU) | 15.3 GB | 10.31 | **5.04** | **2.0×** | +| Qwen/Qwen3-30B-A3B | MoE (8 GPU) | 56.9 GB | 49.26 | **23.76** | **2.1×** | +| openai/gpt-oss-120b | MXFP4 quantized (8 GPU) | 60.8 GB | 49.70 | **17.73** | **2.8×** | + +## Running Benchmarks + +```bash +# Quick comparison (hf vs gpu_direct) +python benchmark_loading.py \ + --model Qwen/Qwen3-8B --gpus 0 \ + --experiments hf gpu_direct --no-verify + +# Full sweep with multiple concurrency/worker configs +python benchmark_loading.py --model Qwen/Qwen3-30B-A3B --gpus 0,1,2,3 --repeats 3 +``` diff --git a/tests/performance/loading/benchmark_loading.py b/tests/performance/loading/benchmark_loading.py new file mode 100644 index 00000000..c4793664 --- /dev/null +++ b/tests/performance/loading/benchmark_loading.py @@ -0,0 +1,774 @@ +"""Benchmark: nnsight LanguageModel loading — HF vs Run:AI streaming. + +Compares wall-clock time, peak GPU/CPU memory, and output correctness for +loading HuggingFace models through nnsight's LanguageModel interface: + + hf — standard from_pretrained (safetensors mmap → threaded loading) + runai_stream — run:ai streaming, CPU clone, HF workers do .to(cuda) + runai_gpu_direct — run:ai + direct .to(cuda) from Run:AI buffer to target GPU + +Page cache invalidation between experiments: + By default, uses posix_fadvise(FADV_DONTNEED) on model shard files to evict + their pages from the kernel page cache. This is reliable and requires no sudo. + With --sudo-drop-caches, uses 'echo 3 > /proc/sys/vm/drop_caches' instead. + With --no-drop-caches, skips invalidation entirely (warm-cache runs). + +Usage: + # Cold cache (default): fadvise evicts model pages before each run + python benchmark_loading.py --model meta-llama/Llama-3.1-8B --gpus 0,1 + + # Warm cache: warmup reads shards into page cache, then no eviction + python benchmark_loading.py --model meta-llama/Llama-3.1-8B --warmup --no-drop-caches + + # Specific experiments + python benchmark_loading.py --model Qwen/Qwen2.5-7B-Instruct --experiments hf runai_gpu_direct + + # Both methods, 3 repeats, JSON output + python benchmark_loading.py --model meta-llama/Llama-3.1-8B --repeats 3 --output results.json +""" + +import argparse +import gc +import json +import os +import subprocess +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + +import torch + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +@dataclass +class TimingResult: + experiment: str + config: dict + wall_time_s: float + peak_gpu_mem_mb: float = 0.0 + peak_gpu_alloc_mb: float = 0.0 + peak_gpu_reserved_mb: float = 0.0 + peak_rss_mb: float = 0.0 + peak_private_mb: float = 0.0 + disk_read_gib: Optional[float] = None + net_rx_gib: Optional[float] = None + net_tx_gib: Optional[float] = None + error: Optional[str] = None + + +def _read_diskstats(device: str = "nvme0n1") -> Optional[float]: + """Read cumulative bytes read for *device* from sysfs. Returns GiB.""" + stat_path = f"/sys/class/block/{device}/stat" + try: + with open(stat_path) as f: + fields = f.read().split() + sectors_read = int(fields[2]) # 3rd field + return sectors_read * 512 / 1024**3 + except (FileNotFoundError, IndexError, ValueError): + return None + + +def _read_iface_bytes(iface: str = "hsn0") -> Optional[tuple[float, float]]: + """Read (rx_gib, tx_gib) for *iface* from /proc/net/dev.""" + try: + with open("/proc/net/dev") as f: + for line in f.readlines()[2:]: + name, data = line.split(":", 1) + if name.strip() == iface: + fields = data.split() + rx = int(fields[0]) / 1024**3 + tx = int(fields[8]) / 1024**3 + return rx, tx + except (FileNotFoundError, IndexError, ValueError): + pass + return None + + +def get_process_rss_mb() -> float: + """Current process RSS in MB via /proc/self/status.""" + try: + with open("/proc/self/status") as f: + for line in f: + if line.startswith("VmRSS:"): + return int(line.split()[1]) / 1024 + except Exception: + pass + return 0.0 + + +def get_private_dirty_mb() -> float: + """Private dirty memory in MB via /proc/self/smaps_rollup. + + Returns only Private_Dirty — pages the process allocated and wrote to. + Excludes Private_Clean (which includes MAP_PRIVATE mmap read-only pages + like safetensors mmap that don't actually consume extra physical RAM + beyond the page cache). + """ + private_kb = 0 + try: + with open("/proc/self/smaps_rollup") as f: + for line in f: + if line.startswith("Private_Dirty:"): + private_kb += int(line.split()[1]) + except Exception: + return get_process_rss_mb() # fallback + return private_kb / 1024 + + +class PeakMemMonitor: + """Context manager that polls process memory in a background thread. + + Tracks both RSS (total resident, including shared mmap pages) and + private memory (excluding shared pages like safetensors mmap). + Records the peak of each observed during the context. + """ + + def __init__(self, interval: float = 0.1): + self.interval = interval + self.peak_rss_mb: float = 0.0 + self.peak_private_mb: float = 0.0 + self._stop = False + self._thread = None + + def _poll(self): + while not self._stop: + rss = get_process_rss_mb() + private = get_private_dirty_mb() + if rss > self.peak_rss_mb: + self.peak_rss_mb = rss + if private > self.peak_private_mb: + self.peak_private_mb = private + time.sleep(self.interval) + + def __enter__(self): + import threading + self.peak_rss_mb = get_process_rss_mb() + self.peak_private_mb = get_private_dirty_mb() + self._stop = False + self._thread = threading.Thread(target=self._poll, daemon=True) + self._thread.start() + return self + + def __exit__(self, *exc): + self._stop = True + self._thread.join(timeout=2) + # One final sample + rss = get_process_rss_mb() + private = get_private_dirty_mb() + if rss > self.peak_rss_mb: + self.peak_rss_mb = rss + if private > self.peak_private_mb: + self.peak_private_mb = private + + +def get_gpu_mem_allocated_mb(gpu_ids: list[int]) -> float: + return sum(torch.cuda.memory_allocated(i) / (1024 ** 2) for i in gpu_ids) + + +def build_max_memory(gpu_ids: list[int], model_size_bytes: int = 0) -> dict: + """Build a max_memory dict that forces an even split across *gpu_ids*. + + ``device_map="auto"`` fills GPUs sequentially, so setting each GPU to + 90% of its capacity lets the whole model land on the first GPU when it + fits. Instead, we cap each GPU at ``(model_size / n_gpus) * 1.15`` so + accelerate is forced to distribute layers evenly. Falls back to 90% + of physical memory when *model_size_bytes* is unknown or when only one + GPU is used. + """ + max_memory = {} + n_target = len(gpu_ids) + for i in range(torch.cuda.device_count()): + if i in gpu_ids: + physical = torch.cuda.get_device_properties(i).total_memory + if model_size_bytes > 0 and n_target > 1: + # Per-GPU budget: even share + 15 % headroom, capped at 90 % physical + per_gpu = int(model_size_bytes / n_target * 1.15) + max_memory[i] = min(per_gpu, int(physical * 0.9)) + else: + max_memory[i] = int(physical * 0.9) + else: + max_memory[i] = 0 + return max_memory + + +def unload_model(model): + """Fully unload a model and free GPU memory.""" + if model is not None: + del model + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.synchronize() + + +def resolve_model_size_gb(model_id: str, revision: str = "main") -> float: + """Total size of model shard files in GB.""" + from huggingface_hub import snapshot_download + + try: + model_dir = snapshot_download(model_id, revision=revision, local_files_only=True) + except Exception: + return 0.0 + paths = list(Path(model_dir).glob("*.safetensors")) + if not paths: + paths = list(Path(model_dir).glob("*.bin")) + return sum(p.stat().st_size for p in paths) / (1024 ** 3) + + +# --------------------------------------------------------------------------- +# Page cache invalidation +# --------------------------------------------------------------------------- + +def evict_model_pages(model_id: str, revision: str = "main"): + """Drop page cache for model shard files using posix_fadvise(FADV_DONTNEED).""" + from huggingface_hub import snapshot_download + + try: + model_dir = snapshot_download(model_id, revision=revision, local_files_only=True) + except Exception as e: + print(f" [cache] Could not locate model files: {e}") + return + + shard_paths = sorted(Path(model_dir).glob("*.safetensors")) + if not shard_paths: + shard_paths = sorted(Path(model_dir).glob("*.bin")) + if not shard_paths: + print(" [cache] No shard files found to evict") + return + + total_bytes = 0 + for path in shard_paths: + fd = os.open(str(path), os.O_RDONLY) + try: + size = os.fstat(fd).st_size + os.posix_fadvise(fd, 0, size, os.POSIX_FADV_DONTNEED) + total_bytes += size + finally: + os.close(fd) + print(f" [cache] Evicted {len(shard_paths)} shard files " + f"({total_bytes / 1024**3:.1f} GB) via FADV_DONTNEED") + + +def warm_model_pages(model_id: str, revision: str = "main", + num_threads: int = 16): + """Read all model shard files into page cache (warmup for warm-cache runs).""" + from huggingface_hub import snapshot_download + + try: + model_dir = snapshot_download(model_id, revision=revision, + local_files_only=True) + except Exception as e: + print(f" [warmup] Could not locate model files: {e}") + return + + shard_paths = sorted(Path(model_dir).glob("*.safetensors")) + if not shard_paths: + shard_paths = sorted(Path(model_dir).glob("*.bin")) + if not shard_paths: + print(" [warmup] No shard files found") + return + + total_bytes = sum(p.stat().st_size for p in shard_paths) + + def _read_file(path): + with open(path, "rb") as f: + while f.read(16 * 1024 * 1024): + pass + + t0 = time.perf_counter() + print(f" [warmup] Reading {len(shard_paths)} shard files " + f"({total_bytes / 1024**3:.1f} GB) into page cache...") + with ThreadPoolExecutor(max_workers=num_threads) as pool: + list(pool.map(_read_file, shard_paths)) + elapsed = time.perf_counter() - t0 + bw = total_bytes / 1024**3 / elapsed if elapsed > 0 else 0 + print(f" [warmup] Done in {elapsed:.1f}s ({bw:.1f} GB/s)") + + +def sudo_drop_caches(): + print(" [cache] Dropping page caches (sudo)...") + subprocess.run( + ["sudo", "sh", "-c", "sync; echo 3 > /proc/sys/vm/drop_caches"], + check=True, timeout=30, + ) + + +# --------------------------------------------------------------------------- +# Experiments +# --------------------------------------------------------------------------- + +def _patch_hf_workers(n_workers: int): + """Set the number of worker threads for from_pretrained shard loading. + + Supports both transformers v5 (GLOBAL_WORKERS) and v4 (env var). + Returns a restore function. + """ + try: + import transformers.core_model_loading as cml + original = cml.GLOBAL_WORKERS + cml.GLOBAL_WORKERS = n_workers + + def restore(): + cml.GLOBAL_WORKERS = original + return restore + except (ImportError, AttributeError): + pass + + # Fallback: transformers v4 env var API + old_enable = os.environ.get("HF_ENABLE_PARALLEL_LOADING") + old_workers = os.environ.get("HF_PARALLEL_LOADING_WORKERS") + os.environ["HF_ENABLE_PARALLEL_LOADING"] = "1" + os.environ["HF_PARALLEL_LOADING_WORKERS"] = str(n_workers) + + def restore(): + if old_enable is None: + os.environ.pop("HF_ENABLE_PARALLEL_LOADING", None) + else: + os.environ["HF_ENABLE_PARALLEL_LOADING"] = old_enable + if old_workers is None: + os.environ.pop("HF_PARALLEL_LOADING_WORKERS", None) + else: + os.environ["HF_PARALLEL_LOADING_WORKERS"] = old_workers + return restore + + +def run_hf(model_id: str, gpu_ids: list[int], + workers: int = 4, revision: str = "main", + model_size_bytes: int = 0) -> TimingResult: + """Load via LanguageModel with load_format='from_pretrained'.""" + from nnsight import LanguageModel + + max_memory = build_max_memory(gpu_ids, model_size_bytes=model_size_bytes) + restore = _patch_hf_workers(workers) + + try: + with PeakMemMonitor() as mem: + for i in gpu_ids: + torch.cuda.reset_peak_memory_stats(i) + torch.cuda.synchronize() + t0 = time.perf_counter() + + model = LanguageModel( + model_id, + load_format="from_pretrained", + device_map="auto", + max_memory=max_memory, + revision=revision, + dispatch=True, + ) + + torch.cuda.synchronize() + wall = time.perf_counter() - t0 + + peak_gpu = get_gpu_mem_allocated_mb(gpu_ids) + peak_alloc = sum(torch.cuda.max_memory_allocated(i) / (1024 ** 2) for i in gpu_ids) + peak_reserved = sum(torch.cuda.max_memory_reserved(i) / (1024 ** 2) for i in gpu_ids) + unload_model(model) + finally: + restore() + + return TimingResult( + experiment="hf", + config={"workers": workers}, + wall_time_s=wall, + peak_gpu_mem_mb=peak_gpu, + peak_gpu_alloc_mb=peak_alloc, + peak_gpu_reserved_mb=peak_reserved, + peak_rss_mb=mem.peak_rss_mb, + peak_private_mb=mem.peak_private_mb, + ) + + +def run_runai(model_id: str, gpu_ids: list[int], + experiment: str = "gpu_direct", + concurrency: int = 16, + workers: int = 4, + gpu_direct: bool = True, + revision: str = "main", + model_size_bytes: int = 0) -> TimingResult: + """Load via run:ai streaming. + + *concurrency* controls Run:AI I/O threads (disk bandwidth). + *workers* controls HF GLOBAL_WORKERS (GPU transfer parallelism). + + When *gpu_direct* is True, the cache copies tensors directly from the + Run:AI buffer to GPU — HF workers just consume pre-placed tensors. + When False, tensors are cloned to CPU and HF workers handle GPU transfer. + """ + from nnsight import LanguageModel + + max_memory = build_max_memory(gpu_ids, model_size_bytes=model_size_bytes) + restore = _patch_hf_workers(workers) + + try: + with PeakMemMonitor() as mem: + for i in gpu_ids: + torch.cuda.reset_peak_memory_stats(i) + torch.cuda.synchronize() + t0 = time.perf_counter() + + model = LanguageModel( + model_id, + device_map="auto", + max_memory=max_memory, + concurrency=concurrency, + gpu_direct=gpu_direct, + revision=revision, + dispatch=True, + ) + + torch.cuda.synchronize() + wall = time.perf_counter() - t0 + + peak_gpu = get_gpu_mem_allocated_mb(gpu_ids) + peak_alloc = sum(torch.cuda.max_memory_allocated(i) / (1024 ** 2) for i in gpu_ids) + peak_reserved = sum(torch.cuda.max_memory_reserved(i) / (1024 ** 2) for i in gpu_ids) + unload_model(model) + finally: + restore() + + return TimingResult( + experiment=experiment, + config={"concurrency": concurrency, "workers": workers}, + wall_time_s=wall, + peak_gpu_mem_mb=peak_gpu, + peak_gpu_alloc_mb=peak_alloc, + peak_gpu_reserved_mb=peak_reserved, + peak_rss_mb=mem.peak_rss_mb, + peak_private_mb=mem.peak_private_mb, + ) + + +# --------------------------------------------------------------------------- +# Correctness verification +# --------------------------------------------------------------------------- + +def _load_and_get_logits(model_id, gpu_ids, revision, prompt, + model_size_bytes=0, **extra_kwargs): + """Load model, run one forward pass, return (logits, decoded_token, token_id).""" + from nnsight import LanguageModel + + max_memory = build_max_memory(gpu_ids, model_size_bytes=model_size_bytes) + model = LanguageModel( + model_id, device_map="auto", max_memory=max_memory, + revision=revision, dispatch=True, **extra_kwargs, + ) + with model.trace(prompt): + logits = model.lm_head.output.save() + token_id = logits[0, -1].argmax(dim=-1).item() + decoded = model.tokenizer.decode(token_id) + logits_cpu = logits.cpu().float() + unload_model(model) + return logits_cpu, decoded, token_id + + +def verify_outputs(model_id: str, gpu_ids: list[int], + experiments: list[str], + revision: str = "main", + model_size_bytes: int = 0) -> bool: + """Load with each selected path, run same prompt, compare logits.""" + if len(experiments) < 2: + print(" Skipping verification (need at least 2 experiments)") + return True + + prompt = "The Eiffel Tower is in the city of" + print(f"\n Verifying output correctness (prompt: {prompt!r})...") + + load_kwargs = { + "hf": {"load_format": "from_pretrained"}, + "gpu_direct": {}, + } + + # Decide baseline order: prefer hf first + ordered = sorted(experiments, key=lambda e: (e != "hf", e)) + results = {} + for exp in ordered: + print(f" Loading {exp}...") + results[exp] = _load_and_get_logits( + model_id, gpu_ids, revision, prompt, + model_size_bytes=model_size_bytes, **load_kwargs[exp], + ) + + baseline_name = ordered[0] + baseline_logits, baseline_decoded, baseline_token = results[baseline_name] + print(f" Baseline ({baseline_name}): {baseline_decoded!r} (id={baseline_token})") + + all_ok = True + for exp in ordered[1:]: + logits, decoded, token_id = results[exp] + match = torch.allclose(baseline_logits, logits, atol=1e-4) + max_diff = (baseline_logits - logits).abs().max().item() + token_ok = token_id == baseline_token + print(f" {exp}: {decoded!r} (id={token_id}) | " + f"logits {'MATCH' if match else 'MISMATCH'} " + f"(max diff: {max_diff:.2e}) | " + f"token {'MATCH' if token_ok else 'MISMATCH'}") + if not match or not token_ok: + all_ok = False + + return all_ok + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +ALL_EXPERIMENTS = ["hf", "gpu_direct"] + +# HF workers and Run:AI concurrency are independent knobs: +# workers — HF GLOBAL_WORKERS: threads doing GPU transfers (.to(device)) +# concurrency — RUNAI_STREAMER_CONCURRENCY: I/O threads (disk bandwidth) +DEFAULT_RUNAI_WORKERS = 4 + +EXPERIMENT_CONFIGS = { + "hf": [("hf", {"workers": w}) for w in [1, 4, 8]], + "gpu_direct": [("gpu_direct", {"concurrency": c}) for c in [1, 4, 8]], +} + +_RUNNERS = { + "hf": lambda mid, gids, cfg, rev, msb: run_hf( + mid, gids, workers=cfg.get("workers", 4), revision=rev, + model_size_bytes=msb), + "gpu_direct": lambda mid, gids, cfg, rev, msb: run_runai( + mid, gids, experiment="gpu_direct", + concurrency=cfg.get("concurrency", 16), + workers=cfg.get("workers", DEFAULT_RUNAI_WORKERS), + gpu_direct=True, revision=rev, + model_size_bytes=msb), +} + + +def run_single_config(exp_name, config, model_id, gpu_ids, revision, + model_size_bytes=0): + runner = _RUNNERS.get(exp_name) + if runner is None: + raise ValueError(f"Unknown experiment: {exp_name}") + return runner(model_id, gpu_ids, config, revision, model_size_bytes) + + +def main(): + parser = argparse.ArgumentParser( + description="Benchmark nnsight model loading: hf vs runai_stream vs runai_gpu_direct" + ) + parser.add_argument("--model", default="openai-community/gpt2", + help="HuggingFace model ID") + parser.add_argument("--revision", default="main") + parser.add_argument("--gpus", default=None, + help="Comma-separated GPU IDs (default: all)") + parser.add_argument("--experiments", nargs="+", + default=ALL_EXPERIMENTS, + choices=ALL_EXPERIMENTS) + parser.add_argument("--repeats", type=int, default=1) + parser.add_argument("--sudo-drop-caches", action="store_true") + parser.add_argument("--no-drop-caches", action="store_true") + parser.add_argument("--warmup", action="store_true", + help="Read model shards into page cache before timing") + parser.add_argument("--disk-device", default="nvme0n1", + help="Block device for disk read stats (default: nvme0n1)") + parser.add_argument("--net-iface", default="hsn0", + help="Network interface for RX/TX stats (default: hsn0)") + parser.add_argument("--no-verify", action="store_true", default=False, + help="Skip output correctness verification") + parser.add_argument("--output", default=None, help="JSON output file") + args = parser.parse_args() + + if args.gpus: + gpu_ids = [int(x) for x in args.gpus.split(",")] + else: + gpu_ids = list(range(torch.cuda.device_count())) + + # Check run:ai availability + try: + from runai_model_streamer import SafetensorsStreamer # noqa: F401 + has_runai = True + except ImportError: + has_runai = False + + runai_exps = {"gpu_direct"} + if not has_runai: + skipped = [e for e in args.experiments if e in runai_exps] + if skipped: + print(f"WARNING: runai-model-streamer not installed. " + f"Skipping {skipped}.") + args.experiments = [e for e in args.experiments if e not in runai_exps] + if not args.experiments: + print("ERROR: No valid experiments. Exiting.") + sys.exit(1) + + # Ensure model is downloaded before any experiments or cache operations + from huggingface_hub import snapshot_download + print(f"Downloading model (if needed): {args.model}") + snapshot_download(args.model, revision=args.revision) + + # Resolve model size + model_size_gb = resolve_model_size_gb(args.model, args.revision) + + print(f"Model: {args.model}") + print(f"Model size: {model_size_gb:.1f} GB") + print(f"GPUs: {gpu_ids}") + print(f"Repeats: {args.repeats}") + print(f"Experiments: {args.experiments}") + print(f"run:ai: {'available' if has_runai else 'NOT installed'}") + + # Cache invalidation setup + if args.sudo_drop_caches: + cache_mode = "sudo_drop_caches" + print("Cache mode: sudo drop_caches") + try: + subprocess.run(["sudo", "-n", "true"], check=True, timeout=5, + capture_output=True) + except Exception: + print("ERROR: --sudo-drop-caches requires passwordless sudo.") + sys.exit(1) + elif args.no_drop_caches: + cache_mode = "warm" if args.warmup else "disabled" + print(f"Cache mode: {cache_mode}") + else: + cache_mode = "cold" + print("Cache mode: fadvise (FADV_DONTNEED on model shards)") + + def drop_caches(): + if args.no_drop_caches: + return + if args.sudo_drop_caches: + sudo_drop_caches() + else: + evict_model_pages(args.model, args.revision) + + # Correctness verification + if not args.no_verify and len(args.experiments) > 1: + ok = verify_outputs(args.model, gpu_ids, args.experiments, args.revision, + model_size_bytes=int(model_size_gb * 1024**3)) + if not ok: + print("\n WARNING: Output mismatch between loading paths!") + else: + print(" PASS: All outputs match.") + + # Warmup: populate page cache before timed runs + if args.warmup: + warm_model_pages(args.model, args.revision) + + # Run experiments + all_results = [] + for exp_name in args.experiments: + gc.collect() + drop_caches() + if args.warmup: + warm_model_pages(args.model, args.revision) + + print(f"\n{'=' * 60}") + print(f"Experiment: {exp_name}") + print(f"{'=' * 60}") + + for rep in range(args.repeats): + if args.repeats > 1: + print(f"\n--- Repeat {rep + 1}/{args.repeats} ---") + + for _, config in EXPERIMENT_CONFIGS[exp_name]: + drop_caches() + gc.collect() + torch.cuda.empty_cache() + torch.cuda.synchronize() + + try: + disk_before = _read_diskstats(args.disk_device) + net_before = _read_iface_bytes(args.net_iface) + + result = run_single_config( + exp_name, config, args.model, gpu_ids, + args.revision, + model_size_bytes=int(model_size_gb * 1024**3), + ) + + disk_after = _read_diskstats(args.disk_device) + net_after = _read_iface_bytes(args.net_iface) + + result.config["repeat"] = rep + + if disk_before is not None and disk_after is not None: + result.disk_read_gib = disk_after - disk_before + if net_before is not None and net_after is not None: + result.net_rx_gib = net_after[0] - net_before[0] + result.net_tx_gib = net_after[1] - net_before[1] + + bw = model_size_gb / result.wall_time_s if result.wall_time_s > 0 else 0 + print(f"\n [OK] {result.experiment} | config={result.config}") + print(f" wall_time: {result.wall_time_s:.2f}s") + print(f" bw: {bw:.2f} GB/s") + print(f" peak_gpu_mem: {result.peak_gpu_mem_mb:.0f} MB") + print(f" peak_gpu_alloc: {result.peak_gpu_alloc_mb:.0f} MB") + print(f" peak_gpu_resrv: {result.peak_gpu_reserved_mb:.0f} MB") + print(f" peak_rss: {result.peak_rss_mb:.0f} MB") + print(f" peak_private: {result.peak_private_mb:.0f} MB") + if result.disk_read_gib is not None: + print(f" disk_read: {result.disk_read_gib:.3f} GiB") + if result.net_rx_gib is not None: + print(f" net_rx: {result.net_rx_gib:.3f} GiB") + + all_results.append(result) + except Exception as e: + print(f"\n [ERROR] {exp_name} config={config}: {e}") + import traceback + traceback.print_exc() + all_results.append(TimingResult( + experiment=exp_name, config={**config, "repeat": rep}, + wall_time_s=0, error=str(e), + )) + + gc.collect() + torch.cuda.empty_cache() + torch.cuda.synchronize() + + # Summary + print(f"\n{'=' * 60}") + print("SUMMARY") + print(f"{'=' * 60}") + hdr = (f"{'Experiment':<12} {'Config':<28} {'Wall (s)':<10} " + f"{'BW (GB/s)':<10} {'GPU (MB)':<10} {'PkAlloc(MB)':<12} " + f"{'PkResrv(MB)':<12} {'RSS (MB)':<10} {'Priv (MB)':<10}") + print(hdr) + print("-" * len(hdr)) + for r in all_results: + if r.error: + continue + config_str = json.dumps(r.config, default=str) + if len(config_str) > 26: + config_str = config_str[:23] + "..." + bw = model_size_gb / r.wall_time_s if r.wall_time_s > 0 else 0 + print(f"{r.experiment:<12} {config_str:<28} {r.wall_time_s:<10.2f} " + f"{bw:<10.2f} {r.peak_gpu_mem_mb:<10.0f} {r.peak_gpu_alloc_mb:<12.0f} " + f"{r.peak_gpu_reserved_mb:<12.0f} {r.peak_rss_mb:<10.0f} " + f"{r.peak_private_mb:<10.0f}") + + # JSON output + if args.output: + out = [{ + "experiment": r.experiment, "config": r.config, + "wall_time_s": r.wall_time_s, + "peak_gpu_mem_mb": r.peak_gpu_mem_mb, + "peak_gpu_alloc_mb": r.peak_gpu_alloc_mb, + "peak_gpu_reserved_mb": r.peak_gpu_reserved_mb, + "peak_rss_mb": r.peak_rss_mb, + "peak_private_mb": r.peak_private_mb, + "disk_read_gib": r.disk_read_gib, + "net_rx_gib": r.net_rx_gib, + "net_tx_gib": r.net_tx_gib, + "error": r.error, + } for r in all_results] + with open(args.output, "w") as f: + json.dump({ + "model": args.model, "model_size_gb": model_size_gb, + "gpus": gpu_ids, "cache_mode": cache_mode, + "results": out, + }, f, indent=2) + print(f"\nResults saved to {args.output}") + + +if __name__ == "__main__": + main()