diff --git a/README.md b/README.md index 42382e3a..9032b11b 100644 --- a/README.md +++ b/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: diff --git a/air_llm/airllm/airllm_base.py b/air_llm/airllm/airllm_base.py index d8e28351..080d031f 100644 --- a/air_llm/airllm/airllm_base.py +++ b/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,6 +446,25 @@ 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: @@ -438,7 +472,7 @@ def forward( 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: diff --git a/air_llm/airllm/device_utils.py b/air_llm/airllm/device_utils.py new file mode 100644 index 00000000..ec06b593 --- /dev/null +++ b/air_llm/airllm/device_utils.py @@ -0,0 +1,122 @@ +""" +Device abstraction utilities for AirLLM. + +Adds support for: + - NVIDIA CUDA (device="cuda:0") + - Intel / AMD integrated GPU via DirectML on Windows (device="privateuseone:0") + - Apple Silicon via MPS (device="mps") [MLX path handles this separately] + - CPU fallback (device="cpu") +""" + +import torch + + +# --------------------------------------------------------------------------- +# DirectML detection +# --------------------------------------------------------------------------- + +try: + import torch_directml # pip install torch-directml + + _directml_available = True +except ImportError: + _directml_available = False + + +def is_directml_available() -> bool: + return _directml_available + + +def get_directml_device(index: int = 0): + """Return a torch-directml device handle, or None if not available.""" + if _directml_available: + import torch_directml + return torch_directml.device(index) + return None + + +# --------------------------------------------------------------------------- +# Device type helpers +# --------------------------------------------------------------------------- + +def get_device_type(device: str) -> str: + """ + Normalise a device string to one of: + "cuda" | "directml" | "mps" | "cpu" + """ + d = str(device).lower() + if d.startswith("cuda"): + return "cuda" + # torch-directml registers as "privateuseone" internally + if d.startswith("privateuseone") or d.startswith("dml") or d.startswith("directml"): + return "directml" + if d.startswith("mps"): + return "mps" + return "cpu" + + +def is_cuda_device(device: str) -> bool: + return get_device_type(device) == "cuda" + + +def is_directml_device(device: str) -> bool: + return get_device_type(device) == "directml" + + +def can_pin_memory(device: str) -> bool: + """ + pin_memory() is only meaningful when copying to CUDA. + It's a no-op (and sometimes errors) for DirectML / MPS / CPU targets. + """ + return is_cuda_device(device) + + +# --------------------------------------------------------------------------- +# Device-agnostic cache clearing +# --------------------------------------------------------------------------- + +def empty_cache(device: str) -> None: + """Free unused memory on the given device.""" + dtype = get_device_type(device) + if dtype == "cuda": + torch.cuda.empty_cache() + elif dtype == "mps": + # torch.mps.empty_cache() is available in PyTorch >= 2.0 + # but calling it on a non-Mac machine raises RuntimeError + if hasattr(torch.mps, "empty_cache"): + try: + torch.mps.empty_cache() + except RuntimeError: + pass + # DirectML and CPU: nothing to do + + +# --------------------------------------------------------------------------- +# Device-agnostic free-memory query +# --------------------------------------------------------------------------- + +def get_free_memory_bytes(device: str) -> int: + """ + Return free device memory in bytes, or -1 if unavailable. + """ + dtype = get_device_type(device) + if dtype == "cuda": + try: + free, _ = torch.cuda.mem_get_info() + return free + except Exception: + return -1 + # MPS / DirectML / CPU: no reliable API yet + return -1 + + +# --------------------------------------------------------------------------- +# Compression support check +# --------------------------------------------------------------------------- + +def supports_bitsandbytes(device: str) -> bool: + """ + bitsandbytes only works on NVIDIA CUDA devices. + Integrated GPUs (DirectML / MPS) must skip compression. + """ + return is_cuda_device(device) diff --git a/air_llm/airllm/profiler.py b/air_llm/airllm/profiler.py index d457605d..44017f2e 100644 --- a/air_llm/airllm/profiler.py +++ b/air_llm/airllm/profiler.py @@ -1,25 +1,29 @@ import torch +from .device_utils import get_free_memory_bytes class LayeredProfiler: - def __init__(self, print_memory=False): + def __init__(self, print_memory=False, device: str = "cuda"): self.profiling_time_dict = {} self.print_memory = print_memory - self.min_free_mem = 1024*1024*1024*1024 - + self.device = device + self.min_free_mem = 1024 * 1024 * 1024 * 1024 def add_profiling_time(self, item, time): - if not item in self.profiling_time_dict: + if item not in self.profiling_time_dict: self.profiling_time_dict[item] = [] self.profiling_time_dict[item].append(time) if self.print_memory: - free_mem = torch.cuda.mem_get_info()[0] - self.min_free_mem = min(self.min_free_mem, free_mem) - print(f"free vmem @{item}: {free_mem/1024/1024/1024:.02f}GB, min free: {self.min_free_mem/1024/1024/1024:.02f}GB") + free_mem = get_free_memory_bytes(self.device) + if free_mem >= 0: + self.min_free_mem = min(self.min_free_mem, free_mem) + print(f"free vmem @{item}: {free_mem/1024/1024/1024:.02f}GB, min free: {self.min_free_mem/1024/1024/1024:.02f}GB") + else: + print(f"free vmem @{item}: n/a (device '{self.device}' does not expose memory info)") def clear_profiling_time(self): for item in self.profiling_time_dict.keys(): diff --git a/air_llm/airllm/utils.py b/air_llm/airllm/utils.py index e2536e99..1aa9a512 100644 --- a/air_llm/airllm/utils.py +++ b/air_llm/airllm/utils.py @@ -20,9 +20,11 @@ import torch import torch.nn as nn +from safetensors import safe_open from safetensors.torch import load_file, save_file from .persist import ModelPersister +from .device_utils import empty_cache, supports_bitsandbytes try: @@ -72,17 +74,25 @@ class NotEnoughSpaceException(Exception): pass # Function to clean RAM & vRAM -def clean_memory(): +def clean_memory(device: str = "cuda"): gc.collect() try: ctypes.CDLL("libc.so.6").malloc_trim(0) - except Exception as ex: - # maybe platform + except Exception: + # not available on Windows / macOS pass - torch.cuda.empty_cache() + empty_cache(device) -def uncompress_layer_state_dict(layer_state_dict): +def uncompress_layer_state_dict(layer_state_dict, device: str = "cuda"): + """Decompress a quantized layer state dict back to float16. + + ``device`` is the *target* compute device. bitsandbytes requires CUDA + for dequantization, so we always dequantize on CUDA and then move the + result to the requested device. This means compression can only be used + when at least one CUDA device is present (even if the inference device is + DirectML/MPS/CPU). + """ uncompressed_layer_state_dict = None if any(['4bit' in k for k in layer_state_dict.keys()]): uncompressed_layer_state_dict = {} @@ -90,38 +100,32 @@ def uncompress_layer_state_dict(layer_state_dict): if '4bit' not in k: quant_state_dict = {kk[len(k):]: kv for kk, kv in layer_state_dict.items() if kk.startswith(k) and k != kk} quant_state = bnb.functional.QuantState.from_dict(qs_dict=quant_state_dict, device="cuda") - dqv = bnb.functional.dequantize_nf4(v.cuda(), quant_state) - uncompressed_layer_state_dict[k] = dqv + uncompressed_layer_state_dict[k] = dqv.to(device) del layer_state_dict elif any(['8bit' in k for k in layer_state_dict.keys()]): uncompressed_layer_state_dict = {} for k, v in layer_state_dict.items(): if '8bit' not in k: - absmax = layer_state_dict[k + ".8bit.absmax"] code = layer_state_dict[k + ".8bit.code"] - dqv = bnb.functional.dequantize_blockwise(v.cuda(), bnb.functional.QuantState(absmax=absmax.cuda(), code=code.cuda(), blocksize=2048, dtype=torch.float16)) - uncompressed_layer_state_dict[k] = dqv + uncompressed_layer_state_dict[k] = dqv.to(device) del layer_state_dict return layer_state_dict if uncompressed_layer_state_dict is None else uncompressed_layer_state_dict -def load_layer(local_path, layer_name, profiling=False): - #layer_state_dict = load_file(Path(local_path) / (layer_name + ".safetensors"), device="cpu") +def load_layer(local_path, layer_name, profiling=False, device: str = "cuda"): layer_state_dict = ModelPersister.get_model_persister().load_model(layer_name, local_path) if profiling: t = time.process_time() - to_return = uncompress_layer_state_dict(layer_state_dict) - - #clean_memory() + to_return = uncompress_layer_state_dict(layer_state_dict, device=device) if profiling: elapsed_time = time.process_time() - t @@ -155,21 +159,27 @@ def check_space(checkpoint_path, layer_shards_saving_path=None, compression=None ) def compress_layer_state_dict(layer_state_dict, compression=None): + """Quantize layer weights using bitsandbytes (CUDA-only). + + Compression always happens on CUDA regardless of the inference device + because bitsandbytes only supports NVIDIA GPUs. The compressed tensors + are stored on CPU / disk and later dequantized at load time. + """ compressed_layer_state_dict = None if compression == '4bit': compressed_layer_state_dict = {} for k, v in layer_state_dict.items(): v_quant, quant_state = bnb.functional.quantize_nf4(v.cuda(), blocksize=64) - compressed_layer_state_dict[k] = v_quant + compressed_layer_state_dict[k] = v_quant.cpu() for quant_state_k, quant_state_v in save_quant_state_to_dict(quant_state).items(): - compressed_layer_state_dict[k + ".4bit." + quant_state_k] = quant_state_v + compressed_layer_state_dict[k + ".4bit." + quant_state_k] = quant_state_v.cpu() if isinstance(quant_state_v, torch.Tensor) else quant_state_v elif compression == '8bit': compressed_layer_state_dict = {} for k, v in layer_state_dict.items(): v_quant, quant_state = bnb.functional.quantize_blockwise(v.cuda(), blocksize=2048) - absmax = quant_state.absmax.clone().contiguous() - code = quant_state.code.clone().contiguous() - compressed_layer_state_dict[k] = v_quant + absmax = quant_state.absmax.clone().contiguous().cpu() + code = quant_state.code.clone().contiguous().cpu() + compressed_layer_state_dict[k] = v_quant.cpu() compressed_layer_state_dict[k + ".8bit.absmax"] = absmax compressed_layer_state_dict[k + ".8bit.code"] = code @@ -208,11 +218,26 @@ def split_and_save_layers(checkpoint_path, layer_shards_saving_path=None, splitt if os.path.exists(checkpoint_path / 'pytorch_model.bin.index.json'): with open(checkpoint_path / 'pytorch_model.bin.index.json', 'rb') as f: index = json.load(f)['weight_map'] - else: + elif os.path.exists(checkpoint_path / 'model.safetensors.index.json'): safetensors_format = True - assert os.path.exists(checkpoint_path / 'model.safetensors.index.json'), f'model.safetensors.index.json should exist.' with open(checkpoint_path / 'model.safetensors.index.json', 'rb') as f: index = json.load(f)['weight_map'] + elif os.path.exists(checkpoint_path / 'model.safetensors'): + # Single-file safetensors (no shard index) — common for small models <= ~7B + safetensors_format = True + with safe_open(checkpoint_path / 'model.safetensors', framework='pt', device='cpu') as f: + index = {k: 'model.safetensors' for k in f.keys()} + elif os.path.exists(checkpoint_path / 'pytorch_model.bin'): + # Single-file PyTorch bin (no shard index) + _state = torch.load(checkpoint_path / 'pytorch_model.bin', map_location='cpu') + index = {k: 'pytorch_model.bin' for k in _state.keys()} + del _state + else: + raise FileNotFoundError( + f"No model weights found in {checkpoint_path}. " + "Expected one of: pytorch_model.bin.index.json, model.safetensors.index.json, " + "model.safetensors, or pytorch_model.bin." + ) if layer_names is None: n_layers = len(set([int(k.split('.')[2]) for k in index.keys() if 'model.layers' in k])) diff --git a/air_llm/setup.py b/air_llm/setup.py index fcfead43..7af657f8 100644 --- a/air_llm/setup.py +++ b/air_llm/setup.py @@ -35,8 +35,15 @@ def run(self): 'optimum', 'huggingface-hub', 'scipy', - #'bitsandbytes' set it to optional to support fallback when not installable + # 'bitsandbytes' -- optional; required only for 4-bit/8-bit compression on NVIDIA GPUs + # 'torch-directml' -- optional; required for Intel/AMD integrated GPU on Windows + # install with: pip install torch-directml + # then use: device="privateuseone:0" ], + extras_require={ + 'compression': ['bitsandbytes'], + 'directml': ['torch-directml'], + }, cmdclass={ 'install': PostInstallCommand, }, diff --git a/air_llm/tests/test_compression.py b/air_llm/tests/test_compression.py index 7456b014..18834629 100644 --- a/air_llm/tests/test_compression.py +++ b/air_llm/tests/test_compression.py @@ -2,13 +2,11 @@ import unittest import torch -sys.path.insert(0, '../airllm') - -from airllm import compress_layer_state_dict, uncompress_layer_state_dict - +from airllm.utils import compress_layer_state_dict, uncompress_layer_state_dict +@unittest.skipUnless(torch.cuda.is_available(), "compression requires a CUDA (NVIDIA) GPU — skipping on this machine") class TestCompression(unittest.TestCase): def setUp(self): pass @@ -16,11 +14,10 @@ def tearDown(self): pass def test_should_compress_uncompress(self): - #torch.manual_seed(0) a0 = torch.normal(0, 1, (32, 128), dtype=torch.float16).cuda() a1 = torch.normal(0, 1, (32, 128), dtype=torch.float16).cuda() - a_state_dict = {'a0':a0, 'a1':a1} + a_state_dict = {'a0': a0, 'a1': a1} loss_fn = torch.nn.MSELoss() @@ -29,15 +26,14 @@ def test_should_compress_uncompress(self): b = compress_layer_state_dict(a_state_dict, compression) if iloop < 2: - print(f"for compression {compression}, compressed to: { {k:v.shape for k,v in b.items()} }") + print(f"for compression {compression}, compressed to: { {k: v.shape for k, v in b.items()} }") aa = uncompress_layer_state_dict(b) for k in aa.keys(): - if compression is None: self.assertTrue(torch.equal(aa[k], a_state_dict[k])) else: RMSE_loss = torch.sqrt(loss_fn(aa[k], a_state_dict[k])).detach().cpu().item() print(f"compression {compression} loss: {RMSE_loss}") - self.assertLess(RMSE_loss, 0.1) \ No newline at end of file + self.assertLess(RMSE_loss, 0.1) diff --git a/air_llm/tests/test_device_utils.py b/air_llm/tests/test_device_utils.py new file mode 100644 index 00000000..360d70e9 --- /dev/null +++ b/air_llm/tests/test_device_utils.py @@ -0,0 +1,121 @@ +import unittest +import torch + +from airllm.device_utils import ( + get_device_type, + is_cuda_device, + is_directml_device, + can_pin_memory, + supports_bitsandbytes, + is_directml_available, + get_directml_device, + empty_cache, + get_free_memory_bytes, +) + + +class TestGetDeviceType(unittest.TestCase): + + def test_cuda_variants(self): + self.assertEqual(get_device_type("cuda"), "cuda") + self.assertEqual(get_device_type("cuda:0"), "cuda") + self.assertEqual(get_device_type("cuda:1"), "cuda") + self.assertEqual(get_device_type("CUDA:0"), "cuda") + + def test_directml_variants(self): + self.assertEqual(get_device_type("privateuseone:0"), "directml") + self.assertEqual(get_device_type("privateuseone:1"), "directml") + self.assertEqual(get_device_type("dml:0"), "directml") + self.assertEqual(get_device_type("directml:0"), "directml") + + def test_mps(self): + self.assertEqual(get_device_type("mps"), "mps") + self.assertEqual(get_device_type("mps:0"), "mps") + + def test_cpu(self): + self.assertEqual(get_device_type("cpu"), "cpu") + + +class TestDeviceBoolHelpers(unittest.TestCase): + + def test_is_cuda_device(self): + self.assertTrue(is_cuda_device("cuda:0")) + self.assertFalse(is_cuda_device("privateuseone:0")) + self.assertFalse(is_cuda_device("cpu")) + self.assertFalse(is_cuda_device("mps")) + + def test_is_directml_device(self): + self.assertTrue(is_directml_device("privateuseone:0")) + self.assertTrue(is_directml_device("dml:0")) + self.assertFalse(is_directml_device("cuda:0")) + self.assertFalse(is_directml_device("cpu")) + + def test_can_pin_memory(self): + # pin_memory is only useful when copying to CUDA + self.assertTrue(can_pin_memory("cuda:0")) + self.assertFalse(can_pin_memory("privateuseone:0")) + self.assertFalse(can_pin_memory("mps")) + self.assertFalse(can_pin_memory("cpu")) + + def test_supports_bitsandbytes(self): + # bitsandbytes is NVIDIA-only + self.assertTrue(supports_bitsandbytes("cuda:0")) + self.assertFalse(supports_bitsandbytes("privateuseone:0")) + self.assertFalse(supports_bitsandbytes("mps")) + self.assertFalse(supports_bitsandbytes("cpu")) + + +class TestEmptyCache(unittest.TestCase): + + def test_empty_cache_cpu_does_not_crash(self): + empty_cache("cpu") + + def test_empty_cache_directml_does_not_crash(self): + empty_cache("privateuseone:0") + + @unittest.skipUnless( + hasattr(torch.backends, "mps") and torch.backends.mps.is_available(), + "MPS not available on this machine" + ) + def test_empty_cache_mps_does_not_crash(self): + empty_cache("mps") + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA not available") + def test_empty_cache_cuda(self): + empty_cache("cuda:0") + + +class TestGetFreeMemoryBytes(unittest.TestCase): + + def test_returns_minus_one_for_cpu(self): + self.assertEqual(get_free_memory_bytes("cpu"), -1) + + def test_returns_minus_one_for_directml(self): + self.assertEqual(get_free_memory_bytes("privateuseone:0"), -1) + + def test_returns_minus_one_for_mps(self): + self.assertEqual(get_free_memory_bytes("mps"), -1) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA not available") + def test_returns_positive_for_cuda(self): + free = get_free_memory_bytes("cuda:0") + self.assertGreater(free, 0) + + +class TestDirectMLDetection(unittest.TestCase): + + def test_is_directml_available_returns_bool(self): + result = is_directml_available() + self.assertIsInstance(result, bool) + + @unittest.skipUnless(is_directml_available(), "torch-directml not installed") + def test_get_directml_device_returns_valid_device(self): + dev = get_directml_device(0) + self.assertIsNotNone(dev) + self.assertIn("privateuseone", str(dev)) + + def test_get_directml_device_returns_none_when_unavailable(self): + if is_directml_available(): + self.skipTest("torch-directml is installed on this machine") + result = get_directml_device(0) + self.assertIsNone(result) diff --git a/air_llm/tests/test_directml.py b/air_llm/tests/test_directml.py new file mode 100644 index 00000000..6b38d8d4 --- /dev/null +++ b/air_llm/tests/test_directml.py @@ -0,0 +1,99 @@ +import unittest +import torch + +from airllm.device_utils import is_directml_available, get_directml_device + + +@unittest.skipUnless(is_directml_available(), "torch-directml not installed — skipping DirectML tests") +class TestDirectMLTensorOps(unittest.TestCase): + """ + Validates that basic PyTorch tensor operations work correctly on an + Intel / AMD integrated GPU via torch-directml. + + These tests are automatically skipped when torch-directml is not installed. + Install with: pip install torch-directml + """ + + def setUp(self): + self.device = get_directml_device(0) + + def test_device_string(self): + self.assertIn("privateuseone", str(self.device)) + + def test_tensor_creation_on_device(self): + t = torch.randn(64, 64).to(self.device) + self.assertEqual(str(t.device), "privateuseone:0") + + def test_matmul(self): + a = torch.randn(512, 512).to(self.device) + b = torch.randn(512, 512).to(self.device) + c = torch.matmul(a, b) + self.assertEqual(c.shape, torch.Size([512, 512])) + self.assertEqual(str(c.device), "privateuseone:0") + + def test_float16_matmul(self): + """AirLLM uses float16 by default — verify it works on iGPU.""" + a = torch.randn(256, 256, dtype=torch.float16).to(self.device) + b = torch.randn(256, 256, dtype=torch.float16).to(self.device) + c = torch.matmul(a, b) + self.assertEqual(c.dtype, torch.float16) + + def test_layer_move_to_device(self): + """Simulate AirLLM moving a transformer layer's weights to the iGPU.""" + layer_weights = { + 'self_attn.q_proj.weight': torch.randn(512, 512), + 'self_attn.k_proj.weight': torch.randn(512, 512), + 'mlp.gate_proj.weight': torch.randn(1024, 512), + } + moved = {k: v.to(self.device) for k, v in layer_weights.items()} + for name, tensor in moved.items(): + self.assertEqual(str(tensor.device), "privateuseone:0", + f"{name} not on iGPU") + + def test_layer_unload_to_cpu(self): + """Simulate AirLLM unloading a layer back to CPU after forward pass.""" + t = torch.randn(512, 512).to(self.device) + t_cpu = t.cpu() + self.assertEqual(t_cpu.device.type, "cpu") + self.assertEqual(t_cpu.shape, torch.Size([512, 512])) + + def test_airllm_layer_cycle(self): + """ + Full cycle: CPU → iGPU (load) → compute → CPU (unload). + Mirrors what AirLLM does for every transformer layer. + """ + # Weights on CPU (as loaded from disk) + weight = torch.randn(256, 256, dtype=torch.float16) + x = torch.randn(1, 256, dtype=torch.float16) + + # Move to iGPU + weight_gpu = weight.to(self.device) + x_gpu = x.to(self.device) + + # Forward pass (linear layer equivalent) + out = x_gpu @ weight_gpu.T + + self.assertEqual(str(out.device), "privateuseone:0") + self.assertEqual(out.shape, torch.Size([1, 256])) + + # Unload weight (free iGPU memory) + del weight_gpu + result = out.cpu() + self.assertEqual(result.device.type, "cpu") + + def test_multiple_sequential_layers(self): + """ + Verify sequential layer processing works — each layer loaded, used, + then freed, mimicking AirLLM's sharded inference loop. + """ + x = torch.randn(1, 128, dtype=torch.float16) + + for i in range(5): + layer_weight = torch.randn(128, 128, dtype=torch.float16) + weight_gpu = layer_weight.to(self.device) + x_gpu = x.to(self.device) + x = (x_gpu @ weight_gpu.T).cpu() + del weight_gpu, x_gpu + + self.assertEqual(x.shape, torch.Size([1, 128])) + self.assertEqual(x.device.type, "cpu") diff --git a/air_llm/tests/test_single_file_model.py b/air_llm/tests/test_single_file_model.py new file mode 100644 index 00000000..e383b7d8 --- /dev/null +++ b/air_llm/tests/test_single_file_model.py @@ -0,0 +1,112 @@ +""" +Tests for single-file model support in split_and_save_layers. + +Covers the case where a model ships as a single model.safetensors or +pytorch_model.bin file (no shard index), which is common for models <= ~7B. +""" +import json +import os +import tempfile +import shutil +import unittest + +import torch +from safetensors.torch import save_file + + +class TestSingleFileModelSplit(unittest.TestCase): + + def setUp(self): + self.tmpdir = tempfile.mkdtemp(prefix="airllm_single_file_test_") + + def tearDown(self): + shutil.rmtree(self.tmpdir, ignore_errors=True) + + def _make_fake_model_state(self): + """Minimal Llama-style state dict with 1 decoder layer.""" + hidden = 64 + inter = 128 + vocab = 100 + heads = 4 + return { + "model.embed_tokens.weight": torch.randn(vocab, hidden), + "model.layers.0.input_layernorm.weight": torch.randn(hidden), + "model.layers.0.self_attn.q_proj.weight": torch.randn(hidden, hidden), + "model.layers.0.self_attn.k_proj.weight": torch.randn(hidden // heads, hidden), + "model.layers.0.self_attn.v_proj.weight": torch.randn(hidden // heads, hidden), + "model.layers.0.self_attn.o_proj.weight": torch.randn(hidden, hidden), + "model.layers.0.mlp.gate_proj.weight": torch.randn(inter, hidden), + "model.layers.0.mlp.up_proj.weight": torch.randn(inter, hidden), + "model.layers.0.mlp.down_proj.weight": torch.randn(hidden, inter), + "model.layers.0.post_attention_layernorm.weight": torch.randn(hidden), + "model.norm.weight": torch.randn(hidden), + "lm_head.weight": torch.randn(vocab, hidden), + } + + # ------------------------------------------------------------------ + # single model.safetensors (no index) + # ------------------------------------------------------------------ + def test_split_single_safetensors_file(self): + state = self._make_fake_model_state() + save_file(state, os.path.join(self.tmpdir, "model.safetensors")) + + from airllm.utils import split_and_save_layers + split_path = split_and_save_layers(self.tmpdir) + + self.assertTrue(os.path.isdir(split_path)) + expected_files = [ + "model.embed_tokens.safetensors", + "model.layers.0.safetensors", + "model.norm.safetensors", + "lm_head.safetensors", + ] + for fname in expected_files: + self.assertTrue( + os.path.exists(os.path.join(split_path, fname)), + f"Expected shard file missing: {fname}", + ) + + # ------------------------------------------------------------------ + # single pytorch_model.bin (no index) + # ------------------------------------------------------------------ + def test_split_single_pytorch_bin_file(self): + state = self._make_fake_model_state() + torch.save(state, os.path.join(self.tmpdir, "pytorch_model.bin")) + + from airllm.utils import split_and_save_layers + split_path = split_and_save_layers(self.tmpdir) + + self.assertTrue(os.path.isdir(split_path)) + self.assertTrue( + os.path.exists(os.path.join(split_path, "model.embed_tokens.safetensors")) + ) + + # ------------------------------------------------------------------ + # sharded model.safetensors.index.json still works (regression) + # ------------------------------------------------------------------ + def test_split_sharded_safetensors_still_works(self): + state = self._make_fake_model_state() + shard_file = "model-00001-of-00001.safetensors" + save_file(state, os.path.join(self.tmpdir, shard_file)) + index = {"metadata": {}, "weight_map": {k: shard_file for k in state}} + with open(os.path.join(self.tmpdir, "model.safetensors.index.json"), "w") as f: + json.dump(index, f) + + from airllm.utils import split_and_save_layers + split_path = split_and_save_layers(self.tmpdir) + + self.assertTrue( + os.path.exists(os.path.join(split_path, "model.embed_tokens.safetensors")) + ) + + # ------------------------------------------------------------------ + # no weights at all → FileNotFoundError + # ------------------------------------------------------------------ + def test_raises_when_no_weights(self): + from airllm.utils import split_and_save_layers + with self.assertRaises(FileNotFoundError): + split_and_save_layers(self.tmpdir) + + +if __name__ == "__main__": + unittest.main()