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

README.md

@@ -1,9 +1,10 @@
 ![airllm_logo](https://github.com/lyogavin/airllm/blob/main/assets/airllm_logo_sm.png?v=3&raw=true)
 
-[**Quickstart**](#quickstart) | 
-[**Configurations**](#configurations) | 
-[**MacOS**](#macos) | 
-[**Example notebooks**](#example-python-notebook) | 
+[**Quickstart**](#quickstart) |
+[**Configurations**](#configurations) |
+[**MacOS**](#macos) |
+[**Windows iGPU**](#windows-integrated-gpu-intelamd) |
+[**Example notebooks**](#example-python-notebook) |
 [**FAQ**](#faq)
 
 **AirLLM** optimizes inference memory usage, allowing 70B large language models to run inference on a single 4GB GPU card without quantization, distillation and pruning. And you can run **405B Llama3.1** on **8GB vram** now.
@@ -29,6 +30,8 @@
 * [Best AI Facial Expression Editor](https://crazyfaceai.com)
 
 ## Updates
+[2025/03/19] v2.12.0: Add **Windows Integrated GPU** support (Intel/AMD iGPU via DirectML). Run models on Intel Iris Xe and AMD Radeon integrated graphics without an NVIDIA GPU.
+
 [2024/08/20] v2.11.0: Support Qwen2.5
 
 [2024/08/18] v2.10.1 Support CPU inference. Support non sharded models. Thanks @NavodPeiris for the great work! 
@@ -63,6 +66,7 @@
 * [Model Compression](#model-compression---3x-inference-speed-up)
 * [Configurations](#configurations)
 * [Run on MacOS](#macos)
+* [Windows Integrated GPU (Intel/AMD)](#windows-integrated-gpu-intelamd)
 * [Example notebooks](#example-python-notebook)
 * [Supported Models](#supported-models)
 * [Acknowledgement](#acknowledgement)
@@ -168,6 +172,58 @@ Just install airllm and run the code the same as on linux. See more in [Quick St
 Example [python notebook] (https://github.com/lyogavin/airllm/blob/main/air_llm/examples/run_on_macos.ipynb)
 
 
+## Windows Integrated GPU (Intel/AMD)
+
+AirLLM supports Intel and AMD **integrated GPUs** on Windows via [torch-directml](https://pypi.org/project/torch-directml/). This enables running large language models on laptops and desktops with no discrete NVIDIA GPU.
+
+**Supported hardware:**
+* Intel UHD Graphics, Iris Xe, Arc (integrated)
+* AMD Radeon integrated graphics
+
+### 1. Install
+
+```bash
+pip install torch-directml
+pip install airllm[directml]
+```
+
+### 2. Inference
+
+Pass `device="privateuseone:0"` to target the integrated GPU (index `0` is usually the iGPU on laptops):
+
+```python
+from airllm import AutoModel
+
+MAX_LENGTH = 128
+model = AutoModel.from_pretrained(
+    "garage-bAInd/Platypus2-70B-instruct",
+    device="privateuseone:0",   # Intel / AMD iGPU via DirectML
+)
+
+input_text = ['What is the capital of United States?']
+
+input_tokens = model.tokenizer(input_text,
+    return_tensors="pt",
+    return_attention_mask=False,
+    truncation=True,
+    max_length=MAX_LENGTH,
+    padding=False)
+
+generation_output = model.generate(
+    input_tokens['input_ids'].to("privateuseone:0"),
+    max_new_tokens=20,
+    use_cache=True,
+    return_dict_in_generate=True)
+
+print(model.tokenizer.decode(generation_output.sequences[0]))
+```
+
+### Notes
+
+* **Compression not supported on iGPU** — `bitsandbytes` (used for `4bit`/`8bit` compression) requires an NVIDIA GPU. Do not pass `compression=` when using DirectML. AirLLM will raise a clear error if you try.
+* Prefetching (background disk loading) still works on DirectML via `ThreadPoolExecutor`.
+* Memory info (`profiling_mode=True`) reports `n/a` for iGPU — DirectML does not expose a free-memory API.
+
 ## Example Python Notebook
 
 Example colabs here:

air_llm/airllm/airllm_base.py

@@ -15,10 +15,17 @@
 
 from .profiler import LayeredProfiler
 
-from optimum.bettertransformer import BetterTransformer
+try:
+    from optimum.bettertransformer import BetterTransformer
+    bettertransformer_available = True
+except (ImportError, ModuleNotFoundError):
+    bettertransformer_available = False
 
 from .utils import clean_memory, load_layer, \
     find_or_create_local_splitted_path
+from .device_utils import (is_cuda_device, is_directml_available,
+                            can_pin_memory, supports_bitsandbytes,
+                            get_device_type)
 
 try:
     import bitsandbytes as bnb
@@ -84,7 +91,7 @@ def __init__(self, model_local_path_or_repo_id, device="cuda:0", dtype=torch.flo
 
 
         self.profiling_mode = profiling_mode
-        self.profiler = LayeredProfiler()
+        self.profiler = LayeredProfiler(device=device)
 
         self.total_disk_loading_time = None
         self.total_gpu_loading_time = None
@@ -95,6 +102,12 @@ def __init__(self, model_local_path_or_repo_id, device="cuda:0", dtype=torch.flo
         if compression is not None:
             if not bitsandbytes_installed:
                 raise ImportError('WARNING: bitsandbytes not found. Compression needs bitsandbytes. To use compression, please install bitsandbytes: `pip install bitsandbytes`')
+            if not supports_bitsandbytes(device):
+                raise ValueError(
+                    f"Compression ('4bit'/'8bit') requires a CUDA (NVIDIA) device but got device='{device}'. "
+                    f"Integrated GPU / DirectML / MPS devices do not support bitsandbytes. "
+                    f"Run without compression=... on this device."
+                )
 
 
         self.compression = compression
@@ -153,8 +166,10 @@ def __init__(self, model_local_path_or_repo_id, device="cuda:0", dtype=torch.flo
             self.prefetching = False
             print(f"not support prefetching for compression for now. loading with no prepetching mode.")
 
-        # this operation should run only if gpu is available
-        if prefetching and device.startswith("cuda"):
+        # CUDA streams are only available on NVIDIA GPUs.
+        # For DirectML / MPS / CPU we still prefetch using ThreadPoolExecutor
+        # but without a CUDA stream (stream = None).
+        if prefetching and is_cuda_device(device):
             self.stream = torch.cuda.Stream()
         else:
             self.stream = None
@@ -184,14 +199,14 @@ def init_model(self):
         # Load meta model (no memory used)
         self.model = None
 
-        if self.get_use_better_transformer():
+        if self.get_use_better_transformer() and bettertransformer_available:
             try:
                 with init_empty_weights():
                     self.model = AutoModelForCausalLM.from_config(self.config, trust_remote_code=True)
                     self.model = BetterTransformer.transform(self.model)  # enable flash attention
             except ValueError as ve:
                 del self.model
-                clean_memory()
+                clean_memory(self.running_device)
                 self.model = None
 
             if self.model is None:
@@ -207,7 +222,7 @@ def init_model(self):
 
                 except TypeError as ve:
                     del self.model
-                    clean_memory()
+                    clean_memory(self.running_device)
                     self.model = None
 
         # fallback to original way
@@ -270,7 +285,7 @@ def load_layer_to_cpu(self, layer_name):
 
         t = time.time()
 
-        load_layer_output = load_layer(self.checkpoint_path, layer_name, self.profiling_mode)
+        load_layer_output = load_layer(self.checkpoint_path, layer_name, self.profiling_mode, device=self.running_device)
         elapsed_time = time.time() - t
 
         if self.profiling_mode:
@@ -283,15 +298,12 @@ def load_layer_to_cpu(self, layer_name):
         else:
             state_dict = load_layer_output
 
-        # pin memory:
+        # pin memory (only beneficial when copying to a CUDA device):
         if self.prefetching:
             t = time.time()
-            if torch.cuda.is_available():  # Check if CUDA is available
+            if can_pin_memory(self.running_device):
                 for k in state_dict.keys():
-                    state_dict[k].pin_memory()
-            else:
-                # For CPU, no action is needed, but you could optionally add a log or message
-                print("Prefetching is enabled, but no pin_memory operation is needed for CPU.")
+                    state_dict[k] = state_dict[k].pin_memory()
 
             elapsed_time = time.time() - t
             if self.profiling_mode:
@@ -374,6 +386,9 @@ def get_sequence_len(self, seq):
         return seq.shape[1]
 
     def get_pos_emb_args(self, len_p, len_s):
+        if getattr(self, '_cached_position_embeddings', None) is not None:
+            cos, sin = self._cached_position_embeddings
+            return {'position_embeddings': (cos[:, len_p:len_p + len_s], sin[:, len_p:len_p + len_s])}
         return {}
 
     def get_past_key_value_args(self, k_cache, v_cache):
@@ -419,7 +434,7 @@ def forward(
 
         # Reboot the model to make sure buffers are loaded and memory is clean
         del self.model
-        clean_memory()
+        clean_memory(self.running_device)
         self.init_model()
 
         batch = [input_ids_unit.to(self.running_device).unsqueeze(0) for input_ids_unit in input_ids]
@@ -431,14 +446,33 @@ def forward(
         attention_mask = attention_mask.to(self.running_device)
         position_ids = torch.arange(self.max_seq_len, dtype=torch.long, device=self.running_device)[None, :]
 
+        # Pre-compute rotary position embeddings for transformers 5.x models.
+        # In transformers >= 5.x, decoder layers no longer compute cos/sin internally;
+        # they must be passed explicitly as position_embeddings=(cos, sin).
+        self._cached_position_embeddings = None
+        _rotary_emb_mod = getattr(getattr(self.model, 'model', None), 'rotary_emb', None)
+        if _rotary_emb_mod is not None:
+            try:
+                # Re-instantiate on CPU so we can compute without meta weights
+                _re = type(_rotary_emb_mod)(self.config)
+                _cpu_pos_ids = position_ids.cpu()
+                _dummy = torch.zeros(1, dtype=torch.float32)  # dtype determines output dtype
+                _cos, _sin = _re(_dummy, _cpu_pos_ids)
+                self._cached_position_embeddings = (
+                    _cos.to(device=self.running_device, dtype=self.running_dtype),
+                    _sin.to(device=self.running_device, dtype=self.running_dtype),
+                )
+            except Exception:
+                pass
+
         kv_cache_list = [] if use_cache else None
         if use_cache:
             for x in self.layers:
                 kv_cache_list.append(([], []))
         all_hidden_states = [] * len(self.layers) if output_hidden_states else None
         all_self_attns = [] * len(self.layers) if output_attentions else None
 
-        with torch.inference_mode(), ThreadPoolExecutor() as executor:
+        with torch.no_grad(), ThreadPoolExecutor() as executor:
 
             # Load first layer
             if self.prefetching:
@@ -535,7 +569,12 @@ def forward(
                             layer_outputs = layer(seq,
                                                   **kwargs
                                                   )
-                            new_seq = layer_outputs[0]
+                            # transformers 5.x decoder layers return a bare tensor; 4.x returned a tuple
+                            if isinstance(layer_outputs, tuple):
+                                new_seq = layer_outputs[0]
+                            else:
+                                new_seq = layer_outputs
+                                layer_outputs = (new_seq,)  # normalise for attentions/cache access below
 
                             if output_attentions:
                                 all_self_attns[i].append(layer_outputs[1])
@@ -566,7 +605,8 @@ def forward(
                                 kwargs = {**kwargs, **pos_embed_args, **attention_mask_args, **position_ids_args}
 
 
-                                new_seq = layer(seq, **kwargs)[0]
+                                _out = layer(seq, **kwargs)
+                                new_seq = _out[0] if isinstance(_out, tuple) else _out
                             else:
 
                                 kwargs = {'use_cache': True,
@@ -597,7 +637,7 @@ def forward(
                     layer.to("meta")
 
                 layer.to("meta")
-                clean_memory()  # proposed by CPMP
+                clean_memory(self.running_device)  # proposed by CPMP
 
         logits = torch.cat(batch, 0)
         if use_cache: