Skip to content

ramtorch + quanto + PEFT: fixes for moving between devices #2884

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bghira merged 1 commit into from
Jul 23, 2026
Merged

ramtorch + quanto + PEFT: fixes for moving between devices #2884

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions simpletuner/helpers/training/quantisation/peft_workarounds.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"""

Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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]
Expand Down
151 changes: 72 additions & 79 deletions simpletuner/helpers/utils/ramtorch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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.
Expand Down
Loading