diff --git a/simpletuner/helpers/training/quantisation/peft_workarounds.py b/simpletuner/helpers/training/quantisation/peft_workarounds.py index 4c6eb0788..c19c061c3 100644 --- a/simpletuner/helpers/training/quantisation/peft_workarounds.py +++ b/simpletuner/helpers/training/quantisation/peft_workarounds.py @@ -32,6 +32,16 @@ QConv2d, QLinear = None, None +def _ensure_lora_adapter_on_device(layer: LoraLayer, adapter_name: str, device: torch.device) -> None: + for adapters in (layer.lora_A, layer.lora_B): + module = adapters.get(adapter_name) + if module is None: + continue + weight = getattr(module, "weight", None) + if weight is not None and weight.device != device: + module.to(device) + + class QuantoLoraLinear(torch.nn.Module, LoraLayer): """LoRA layer implementation for quanto QLinear""" @@ -77,6 +87,7 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: for active_adapter in self.active_adapters: if active_adapter not in self.lora_A.keys(): continue + _ensure_lora_adapter_on_device(self, active_adapter, x.device) lora_A = self.lora_A[active_adapter] lora_B = self.lora_B[active_adapter] dropout = self.lora_dropout[active_adapter] @@ -245,6 +256,7 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: for active_adapter in self.active_adapters: if active_adapter not in self.lora_A.keys(): continue + _ensure_lora_adapter_on_device(self, active_adapter, x.device) lora_A = self.lora_A[active_adapter] lora_B = self.lora_B[active_adapter] dropout = self.lora_dropout[active_adapter] diff --git a/simpletuner/helpers/utils/ramtorch.py b/simpletuner/helpers/utils/ramtorch.py index 386200138..1405b4ded 100644 --- a/simpletuner/helpers/utils/ramtorch.py +++ b/simpletuner/helpers/utils/ramtorch.py @@ -11,6 +11,7 @@ _peft_lora_config_patched = False _peft_lora_model_patched = False _peft_inject_patched = False +_peft_lora_forward_patched = False @lru_cache(maxsize=1) @@ -47,6 +48,7 @@ def ensure_available() -> dict: imports = _ramtorch_imports() _maybe_patch_peft_lora_config() _maybe_patch_peft_lora_model() + _maybe_patch_peft_lora_forward() _maybe_patch_peft_inject() return imports except ImportError as exc: @@ -106,6 +108,42 @@ def _is_replaceable_linear(module: nn.Module) -> bool: return _is_quanto_linear(module) and hasattr(module, "in_features") and hasattr(module, "out_features") +def _move_peft_lora_adapters_to_device(peft_layer: Any, device: torch.device | None) -> None: + if device is None or getattr(peft_layer, "disable_adapters", False): + return + + try: + active = peft_layer.active_adapters + except Exception: + try: + active = peft_layer.active_adapter + except Exception: + active = None + + if active is None: + active_adapters = list(getattr(peft_layer, "lora_A", {}).keys()) + elif isinstance(active, (list, tuple, set)): + active_adapters = list(active) + else: + active_adapters = [active] + + def _move_module(module: nn.Module | None) -> None: + if module is None: + return + try: + weight = getattr(module, "weight", None) + if weight is not None and weight.device != device: + module.to(device) + except Exception: + pass + + for adapter_name in active_adapters: + if adapter_name in getattr(peft_layer, "lora_A", {}): + _move_module(peft_layer.lora_A[adapter_name]) + if adapter_name in getattr(peft_layer, "lora_B", {}): + _move_module(peft_layer.lora_B[adapter_name]) + + def _tensor_uses_subclass_storage(tensor: Any) -> bool: tensor_type = type(tensor) return tensor_type is not torch.Tensor and tensor_type is not nn.Parameter @@ -354,52 +392,7 @@ class RamTorchPeftLinear(PeftLinear): # type: ignore[misc] """ def _ensure_lora_on_device(self, device): - if device is None: - return - try: - active = self.active_adapter - except Exception: - active = None - - def _active_adapters(candidate): - if candidate is None: - return [] - if isinstance(candidate, (list, tuple, set)): - return list(candidate) - return [candidate] - - def _move_layer(layer): - try: - if hasattr(layer, "to") and layer.weight.device != device: - layer.to(device) - except Exception: - pass - - actives = _active_adapters(active) - keys_to_move = actives or list(self.lora_A.keys()) - - for name in keys_to_move: - try: - key = name - if key in self.lora_A: - _move_layer(self.lora_A[key]) - if key in self.lora_B: - _move_layer(self.lora_B[key]) - if key in getattr(self.lora_embedding_A, "_parameters", {}): - emb_a = self.lora_embedding_A[key] - emb_b = self.lora_embedding_B.get(key) - try: - if hasattr(emb_a, "to") and emb_a.device != device: - emb_a.data = emb_a.data.to(device) - except Exception: - pass - try: - if emb_b is not None and hasattr(emb_b, "to") and emb_b.device != device: - emb_b.data = emb_b.data.to(device) - except Exception: - pass - except Exception: - continue + _move_peft_lora_adapters_to_device(self, device) def forward(self, x: torch.Tensor, *args, **kwargs): # type: ignore[override] self._ensure_lora_on_device(x.device if torch.is_tensor(x) else None) @@ -525,39 +518,7 @@ def _maybe_patch_peft_lora_model() -> bool: class RamTorchPeftLinear(PeftLinear): # type: ignore[misc] def _ensure_lora_on_device(self, device): - if device is None: - return - try: - active = self.active_adapter - except Exception: - active = None - - def _active_adapters(candidate): - if candidate is None: - return [] - if isinstance(candidate, (list, tuple, set)): - return list(candidate) - return [candidate] - - def _move_layer(layer): - try: - if hasattr(layer, "to") and layer.weight.device != device: - layer.to(device) - except Exception: - pass - - actives = _active_adapters(active) - keys_to_move = actives or list(self.lora_A.keys()) - - for name in keys_to_move: - try: - key = name - if key in self.lora_A: - _move_layer(self.lora_A[key]) - if key in self.lora_B: - _move_layer(self.lora_B[key]) - except Exception: - continue + _move_peft_lora_adapters_to_device(self, device) def forward(self, x: torch.Tensor, *args, **kwargs): # type: ignore[override] self._ensure_lora_on_device(x.device if torch.is_tensor(x) else None) @@ -600,6 +561,38 @@ def _wrapped_create_new_module(lora_config, adapter_name, target, **kwargs): return True +def _maybe_patch_peft_lora_forward() -> bool: + """ + Patch PEFT's stock LoRA Linear forward when RamTorch is active. Some models + can still dispatch to the normal PEFT wrapper after quantization/offload, and + those adapters need the same activation-device relocation as RamTorchLinear. + """ + + global _peft_lora_forward_patched + if _peft_lora_forward_patched: + return False + + try: + import torch + from peft.tuners.lora.layer import Linear as PeftLinear + except Exception: + return False + + orig_forward = PeftLinear.forward + if getattr(orig_forward, "_ramtorch_wrapped", False): + _peft_lora_forward_patched = True + return False + + def _wrapped_forward(self, x: torch.Tensor, *args, **kwargs): # type: ignore[override] + _move_peft_lora_adapters_to_device(self, x.device if torch.is_tensor(x) else None) + return orig_forward(self, x, *args, **kwargs) + + _wrapped_forward._ramtorch_wrapped = True + PeftLinear.forward = _wrapped_forward + _peft_lora_forward_patched = True + return True + + def _maybe_patch_peft_inject() -> bool: """ Patch PEFT inject_adapter_in_model to duplicate LoRA state_dict keys with _orig_mod.