Skip to main content

Small GB10/GX10 helpers for loading LLMs without slow cold-page CUDA H2D transfers.

Project description

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

For local development:

uv sync --extra dev

From another local project:

uv pip install -e /path/to/gb10-load-llm

From PyPI:

pip install gb10-load-llm

Transformers Usage

import torch
from transformers import AutoModelForCausalLM

from gb10_load_llm import load_model_to_cuda

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

print(f"touched {touched / 1024**3:.2f} GiB")

load_model_to_cuda passes extra keyword arguments to AutoModelForCausalLM.from_pretrained. By default, Transformers loads models on CPU before this helper touches pages 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 device_map="cuda" for load speed unless you have measured it in your workload or need Accelerate features such as multi-GPU placement or offload.

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

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:

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

uv sync --extra dev
uv run pytest

Build:

uv build

Publish later with:

uv publish

Blog post

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

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

gb10_load_llm-0.1.1.tar.gz (78.2 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

gb10_load_llm-0.1.1-py3-none-any.whl (5.9 kB view details)

Uploaded Python 3

File details

Details for the file gb10_load_llm-0.1.1.tar.gz.

File metadata

  • Download URL: gb10_load_llm-0.1.1.tar.gz
  • Upload date:
  • Size: 78.2 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.10.9 {"installer":{"name":"uv","version":"0.10.9","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}

File hashes

Hashes for gb10_load_llm-0.1.1.tar.gz
Algorithm Hash digest
SHA256 48dfed11af4c15c49cedcc3d28d1c83d7ad86101c9b77ef003027e13d1c532c2
MD5 325676bd1eacf35f616d4c2af826e785
BLAKE2b-256 c72063a4b8524dd782e3d7d39b64d6487a718ae894fcdb595805aada88af6778

See more details on using hashes here.

File details

Details for the file gb10_load_llm-0.1.1-py3-none-any.whl.

File metadata

  • Download URL: gb10_load_llm-0.1.1-py3-none-any.whl
  • Upload date:
  • Size: 5.9 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.10.9 {"installer":{"name":"uv","version":"0.10.9","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}

File hashes

Hashes for gb10_load_llm-0.1.1-py3-none-any.whl
Algorithm Hash digest
SHA256 bc82cdcc517fca0811bbe7346124d5bbb329650f21bcca506d36749b2af15447
MD5 614238aa3055d83654e98d3bb5254c30
BLAKE2b-256 0e5ca030875d1344ce006dd16217cd1de6463365f306c03b6c4721f351d29894

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page