Metadata-Version: 2.4
Name: gb10-load-llm
Version: 0.1.2
Summary: Small GB10/GX10 helpers for loading LLMs without slow cold-page CUDA H2D transfers.
Project-URL: Homepage, https://github.com/liusida/gb10-load-llm
Project-URL: Repository, https://github.com/liusida/gb10-load-llm
Project-URL: Issues, https://github.com/liusida/gb10-load-llm/issues
Author: Sida Liu
License-Expression: MIT
License-File: LICENSE
Keywords: cuda,gb10,llm,safetensors,transformers
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Requires-Python: >=3.10
Requires-Dist: torch
Requires-Dist: transformers
Provides-Extra: dev
Requires-Dist: pytest; extra == 'dev'
Description-Content-Type: text/markdown

# gb10-load-llm

Small helpers for loading LLMs on NVIDIA GB10 / GX10 systems without hitting
very slow CUDA Host-to-Device transfers from cold CPU pages.

The core workaround is:

1. Load the model on CPU.
2. Read one byte per CPU page from model parameters and buffers.
3. Move the model to CUDA.

This package keeps that process explicit and reusable.

## Install


```bash
pip install gb10-load-llm
```

## Transformers Usage

```python
from transformers import AutoModelForCausalLM

from gb10_load_llm import load_model_to_cuda

model = load_model_to_cuda(
    AutoModelForCausalLM,
    "Qwen/Qwen2.5-Coder-3B",
    dtype="auto",
    touch="auto",
)
```

`load_model_to_cuda` passes extra keyword arguments to
`AutoModelForCausalLM.from_pretrained`. By default, Transformers loads models on
CPU before this helper applies the touch policy and moves the model to CUDA.

For single-GPU models that fit in VRAM, this CPU-load, touch, then CUDA-move
route is the recommended path on GB10.

## Avoid Accelerate / device_map

For single-GPU models that fit in VRAM, this package does not recommend using
Accelerate / `device_map` as the loading path on GB10.

On the tested GB10 system with cached `Qwen/Qwen2.5-Coder-3B` in bfloat16:

- CPU load + touch + `model.to("cuda")`: about 2-4 seconds to get the model onto
  GPU.
- Unpatched `device_map="cuda"` through Accelerate: about 52 seconds.
- Experimental patched `device_map="cuda"` that touched pages inside
  Transformers' private loading path: about 5-7 seconds.

Even with the experimental touch fix, the direct CPU-touch-then-CUDA route was
faster and simpler. Use Accelerate when you need its features, such as
multi-GPU placement, CPU/disk offload, or models that do not fit on one GPU.
For fastest single-GPU loading on GB10, prefer this package's CPU-first helper.

## Lower-Level Usage

```python
import torch
from transformers import AutoModelForCausalLM

from gb10_load_llm import touch_model_cpu_pages

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-Coder-3B",
    dtype=torch.bfloat16,
    local_files_only=True,
)

touched = touch_model_cpu_pages(model)
model = model.to("cuda")
torch.cuda.synchronize()
```

## Touch Policy

`move_model_to_cuda` accepts:

- `touch=True`: always touch CPU pages before moving.
- `touch=False`: never touch.
- `touch="auto"`: touch only when the CUDA device name looks like GB10.

The Transformers wrapper defaults to `touch=True`, because this package is
GB10-focused and explicit.

You can override all callers with:

```bash
GB10_LOAD_LLM_TOUCH=0  # disable
GB10_LOAD_LLM_TOUCH=1  # enable
```

## Why This Exists

On the tested ASUS Ascent GX10 / NVIDIA GB10 system, moving cold
safetensors-backed CPU weights to CUDA can be tens of seconds slower than
expected. Reading one byte per CPU page before `model.to("cuda")` avoids most of
that cost.

This appears hardware / driver dependent. It is not intended as a universal
PyTorch rule.

## Development

```bash
uv sync --extra dev
uv run pytest
```

Build:

```bash
uv build
```

Publish later with:

```bash
uv publish
```

## Blog post

https://liusida.com/post/cuda-h2d-slowdown-from-cold-mmap-backed-safetensors-pages-on-gb10
