Skip to content

feat: add Windows integrated GPU support via torch-directml (Intel/AMD iGPU) - #271

Open
RajeshKumar11 wants to merge 6 commits into
lyogavin:mainfrom
RajeshKumar11:feat/integrated-gpu-directml
Open

feat: add Windows integrated GPU support via torch-directml (Intel/AMD iGPU)#271
RajeshKumar11 wants to merge 6 commits into
lyogavin:mainfrom
RajeshKumar11:feat/integrated-gpu-directml

Conversation

@RajeshKumar11

@RajeshKumar11 RajeshKumar11 commented Mar 19, 2026 •
edited
Loading

Copy link
Copy Markdown

Closes #270

What this PR does

Adds support for Intel and AMD integrated GPUs on Windows via torch-directml, enabling AirLLM to run on laptops and desktops with no discrete NVIDIA GPU.

Also fixes three transformers 5.x compatibility bugs that broke inference on all devices (CUDA, MPS, CPU) — discovered while testing on iGPU.

Tested on: Intel Iris Xe Graphics — TinyLlama-1.1B correctly answers prompts end-to-end. 28 tests pass, 5 correctly skipped (CUDA/MPS guards).

Changes

New file: air_llm/airllm/device_utils.py

Central device abstraction module:

  • get_device_type(device) — normalises device strings to "cuda" / "directml" / "mps" / "cpu"
  • is_cuda_device / is_directml_device / can_pin_memory / supports_bitsandbytes
  • empty_cache(device) — device-agnostic cache flush
  • get_free_memory_bytes(device) — safe memory query (returns -1 for non-CUDA)
  • get_directml_device(index) — convenience wrapper

air_llm/airllm/airllm_base.py

  • Import device_utils; guard BetterTransformer import with try/except (optimum ≥ 2.0 removed it)
  • Compression guard: raise ValueError with clear message when compression= is used on a non-CUDA device
  • torch.inference_mode() → torch.no_grad() in the layer loop — set_module_tensor_to_device() cannot modify tensors created inside inference_mode (fixes all devices)
  • Pre-compute rotary position_embeddings before the layer loop; pass via get_pos_emb_args() — transformers ≥ 5.x LlamaDecoderLayer no longer computes cos/sin internally (fixes all devices)
  • Handle tensor-vs-tuple decoder layer return: transformers 5.x returns a bare tensor; layer(...)[0] was stripping the batch dim (fixes all devices)
  • Prefetch CUDA stream gated on is_cuda_device(device) instead of bare device.startswith("cuda")
  • pin_memory gated on can_pin_memory(device) (CUDA only)
  • All clean_memory() calls pass device argument

air_llm/airllm/utils.py

  • clean_memory(device) uses empty_cache(device) instead of torch.cuda.empty_cache()
  • uncompress_layer_state_dict / compress_layer_state_dict / load_layer pass device through
  • Single-file model support: handles model.safetensors and pytorch_model.bin without a shard index (small models ≤ ~7B)

air_llm/airllm/profiler.py

  • LayeredProfiler accepts device parameter; uses get_free_memory_bytes(device) instead of torch.cuda.mem_get_info()

air_llm/setup.py

  • extras_require = {'compression': ['bitsandbytes'], 'directml': ['torch-directml']}

README.md

  • New Windows Integrated GPU (Intel/AMD) section with install steps and code example

New tests (air_llm/tests/)

File Tests Notes
test_device_utils.py 19 All device helper functions; CUDA/MPS tests skip gracefully
test_directml.py 8 DirectML tensor ops, layer move/unload, full AirLLM layer cycle on Intel Iris Xe
test_single_file_model.py 4 Single-file safetensors/bin splitting; no GPU required
test_compression.py 1 Added CUDA skip guard

Usage (Intel/AMD iGPU)

pip install torch-directml
from airllm import AutoModel

model = AutoModel.from_pretrained(
    "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    device="privateuseone:0",   # DirectML device string
    dtype=torch.float16,
)

Note: compression='4bit'/'8bit' requires bitsandbytes (CUDA only) and will raise a clear ValueError on iGPU.

Related

PR #272 (fix/single-file-model-support) splits out the single-file model fix as a standalone patch that can be merged independently.

RajeshKumar11 and others added 5 commits March 19, 2026 13:19
@RajeshKumar11 @claude
feat: add device_utils module for device-agnostic GPU abstraction
909cab4 
Introduces device_utils.py with helpers to detect and work with
any compute device uniformly:
- get_device_type()      — normalises device string to cuda/directml/mps/cpu
- is_cuda_device()       — True only for NVIDIA CUDA devices
- is_directml_device()   — True for Intel/AMD iGPU via torch-directml
- can_pin_memory()       — True only when target is a CUDA device
- supports_bitsandbytes()— True only for CUDA (guards compression)
- empty_cache()          — device-agnostic VRAM/cache clear
- get_free_memory_bytes()— device-agnostic memory query (-1 if unavailable)
- is_directml_available()— detects torch-directml install at runtime

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@RajeshKumar11
feat: add Intel/AMD integrated GPU support via torch-directml
6dd0642 
Replaces all CUDA-hardcoded calls with device-agnostic equivalents.

- utils.py: clean_memory(device) uses empty_cache(device); compress/
  uncompress accept device param; load_layer forwards device param
- airllm_base.py: guards compression on non-CUDA with clear ValueError;
  prefetch stream only for CUDA; pin_memory gated by can_pin_memory();
  BetterTransformer import wrapped in try/except for optimum >= 2.0
- profiler.py: LayeredProfiler accepts device param, uses
  get_free_memory_bytes() safe for DirectML/MPS/CPU
@RajeshKumar11
test: add test_device_utils and test_directml; fix test_compression skip
9fc9215 
- test_device_utils.py: 19 unit tests for all device_utils helpers
  covering CUDA/DirectML/MPS/CPU paths; CUDA/MPS tests auto-skip when
  hardware is unavailable
- test_directml.py: 8 integration tests verifying tensor ops on Intel/AMD
  iGPU via DirectML including the full AirLLM layer load/compute/unload
  cycle; entire suite skips when torch-directml is not installed
- test_compression.py: wrap class with @unittest.skipUnless so tests
  skip gracefully on machines without NVIDIA CUDA; fix import path to
  airllm.utils (compress_layer_state_dict was never in airllm __init__)
@RajeshKumar11
docs: add Windows integrated GPU section to README and setup.py extras
5064931 
- README.md: new "Windows Integrated GPU (Intel/AMD)" section with
  install steps, inference example, and notes on compression limitation;
  add iGPU link to nav header and table of contents; add v2.12.0 entry
  to Updates log
- setup.py: add extras_require with 'directml' (torch-directml) and
  'compression' (bitsandbytes) optional groups
@RajeshKumar11
fix: support transformers 5.x API changes for Llama-family inference
ab53f3a 
Three breaking changes in transformers 5.x broke AirLLM on any device:

1. torch.inference_mode() → torch.no_grad() in the layer loop:
   set_module_tensor_to_device() cannot modify inference tensors created
   inside inference_mode; no_grad disables gradients without that restriction.

2. Pre-compute rotary position_embeddings before the layer loop:
   transformers >= 5.x LlamaDecoderLayer no longer computes cos/sin
   internally — they must be passed explicitly as position_embeddings=(cos, sin).
   A fresh LlamaRotaryEmbedding is instantiated from config on CPU and the
   (cos, sin) tensors are moved to the target device once, then sliced per
   layer in get_pos_emb_args().

3. Decoder layer return type: tensor not tuple:
   transformers 5.x LlamaDecoderLayer.forward() returns a bare tensor while
   4.x returned a tuple. The layer(…)[0] indexing was stripping the batch
   dimension, causing wrong shapes in every subsequent layer. Now handles
   both tuple (4.x) and tensor (5.x) returns.

Verified: TinyLlama-1.1B-Chat-v1.0 runs end-to-end on Intel Iris Xe via
DirectML and correctly answers "What is 2 + 2?" → "2 + 2 = 4".

Also adds run_real_model_test.py: a real-hardware test script that downloads
TinyLlama from HuggingFace and runs it layer-by-layer via AirLLM AutoModel.
@RajeshKumar11
RajeshKumar11 force-pushed the feat/integrated-gpu-directml branch from 75430b8 to ab53f3a Compare March 19, 2026 09:28
@RajeshKumar11
fix: support single-file safetensors and pytorch_model.bin without sh…
d206d76 
…ard index

Models <= ~7B (e.g. TinyLlama, Phi, Gemma-2B) are distributed as a single
model.safetensors or pytorch_model.bin file with no shard-index JSON.
AirLLM previously hard-asserted that model.safetensors.index.json must exist,
making these models fail on first use.

Changes in split_and_save_layers() (utils.py):
- model.safetensors.index.json  → handled (existing behaviour, now elif)
- model.safetensors (no index)  → NEW: reads tensor keys via safe_open header
  (no data loaded) and builds weight_map in-memory
- pytorch_model.bin (no index)  → NEW: loads state dict to extract key list
- none of the above             → raises FileNotFoundError with a clear message

Also adds tests/test_single_file_model.py with 4 cases:
  - single model.safetensors splits correctly
  - single pytorch_model.bin splits correctly
  - sharded index path still works (regression guard)
  - missing weights raises FileNotFoundError
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Feature: Add integrated GPU support via torch-directml (Intel/AMD iGPU on Windows)

1 participant

@RajeshKumar11