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 project before this is published:

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

After publishing:

pip install "gb10-load-llm[transformers]"

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.0.tar.gz (78.3 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.0-py3-none-any.whl (6.0 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: gb10_load_llm-0.1.0.tar.gz
  • Upload date:
  • Size: 78.3 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.0.tar.gz
Algorithm Hash digest
SHA256 493c6a009faa07770921aa36e51e34aee5d81d18e15b9000c219707b97c1b375
MD5 3a4ac62d4ae676c0c88f2d0288c7742e
BLAKE2b-256 68402a430774bf60799e7baceb0b4e05813a7b1aa4b4adae6e2c81f8722ad0d6

See more details on using hashes here.

File details

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

File metadata

  • Download URL: gb10_load_llm-0.1.0-py3-none-any.whl
  • Upload date:
  • Size: 6.0 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.0-py3-none-any.whl
Algorithm Hash digest
SHA256 efd933831686029e99a6926e1ed7a309b8acd1d3ac20a246809e1ce487db5fcf
MD5 73369ce0f3f9266e0f97443d65b3d988
BLAKE2b-256 5c45dc55f488bf64dfe2da5e94459a8c411f8338a1c05161a855596b0feac572

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