diff --git a/toolkit/memory_management/manager.py b/toolkit/memory_management/manager.py index 6fd195b..5211499 100644 --- a/toolkit/memory_management/manager.py +++ b/toolkit/memory_management/manager.py @@ -8,6 +8,8 @@ from .manager_modules import ( OstrisLinearLayerMemoryManager, EmbeddingLayerMemoryManager, _DEVICE_STATE, + _move_own_param, + _storage_device, ) import random @@ -290,9 +292,12 @@ class MemoryManager: for sub in module.modules(): if hasattr(sub, "_layer_memory_manager"): continue - for p in sub.parameters(recurse=False): - if p is not None and p.device != device: - p.data = p.data.to(device) + # quantized tensor subclasses (torchao) ignore `p.data = ...` -- an unmanaged + # quantized linear (offload_percent < 1) would stay on cpu; _move_own_param + # swaps the Parameter in that case + for name, p in list(sub.named_parameters(recurse=False)): + if p is not None and _storage_device(p) != device: + _move_own_param(sub, name, p, device) for name, b in sub._buffers.items(): if b is not None and b.device != device: sub._buffers[name] = b.to(device) diff --git a/toolkit/memory_management/manager_modules.py b/toolkit/memory_management/manager_modules.py index f0720e8..59abfd7 100644 --- a/toolkit/memory_management/manager_modules.py +++ b/toolkit/memory_management/manager_modules.py @@ -249,6 +249,34 @@ def _ensure_cpu_pinned(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]: return t +def _storage_device(t: torch.Tensor) -> torch.device: + """Device of a tensor's actual storage. A wrapper subclass (torchao) can report one + device on the outside while its inner tensors sit elsewhere (a `.data =` move updates + only the outer metadata), so walk __tensor_flatten__ down to the leaves.""" + try: + names, _ = t.__tensor_flatten__() + except Exception: + return t.device + for name in names: + inner = getattr(t, name, None) + if isinstance(inner, torch.Tensor): + return _storage_device(inner) + return t.device + + +def _move_own_param(module: nn.Module, name: str, param: nn.Parameter, device) -> None: + """Move one of ``module``'s direct parameters to ``device``. Plain tensors go through + `.data =`; quantized tensor subclasses (torchao AffineQuantizedTensor) get the whole + Parameter replaced -- `.data =` on them moves only the outer metadata, leaving the + inner storage behind (same reason _move_params_to_cpu_and_pin replaces them).""" + with torch.no_grad(): + moved = param.data.to(device) + if _is_quantized_tensor(param.data): + setattr(module, name, nn.Parameter(moved, requires_grad=param.requires_grad)) + else: + param.data = moved + + def _move_params_to_cpu_and_pin(module: nn.Module): """Force parameters to CPU (+pinned) so we can 'bounce' them per forward/backward.""" with torch.no_grad():