# SPDX-FileCopyrightText: 2024-2025 ModelCloud.ai # SPDX-FileCopyrightText: 2024-2025 qubitium@modelcloud.ai # SPDX-License-Identifier: Apache-2.0 # Contact: qubitium@modelcloud.ai, x.com/qubitium from typing import Any, Dict, Optional import torch import transformers from torch import Tensor, nn from torch.nn import Parameter from torch.nn.modules.conv import _ConvNd from ..utils.logger import setup_logger from ..utils.module_locks import get_parent_lock from ..utils.stream import stream_sync as stream_sync_events from ..utils.stream import stream_tensor_dict_to_cpu log = setup_logger() class NamedModule(torch.nn.Module): """Thread-safe wrapper that adds stable names and scratch state to a module.""" def __init__(self, module: torch.nn.Module, name: str, full_name:str, layer_index: int) -> None: """Wraps a submodule with naming metadata used by the looper pipeline.""" super().__init__() self._parent_lock = get_parent_lock(full_name) self.module = module # wrapped module self.module_dtype = next(module.parameters()).dtype self.name = name # module name self.full_name = full_name # full dotted path inside model self.layer_index = layer_index # layer index for repeated blocks if hasattr(module, "module_name") and module.module_name is None: module.module_name = full_name # persistent work state for named module (used by some LoopProcessors) # store all `processed()` work state/data/result here self.state = {} # print(f"NamedModule init: name: `{name}, full-name: `{full_name}`") # store original in/out features since weight.data will changed later on if isinstance(module, nn.Linear): in_features = module.in_features out_features = module.out_features elif isinstance(module, _ConvNd): in_features = module.in_channels out_features = module.out_channels # elif isinstance(module, nn.Conv2d): # in_features = module.in_channels # out_features = module.out_channels # elif isinstance(module, nn.Conv2d): # in_features = module.in_channels # out_features = module.out_channels elif isinstance(module, transformers.pytorch_utils.Conv1D): in_features = module.weight.shape[0] out_features = module.weight.shape[1] else: in_features = None out_features = None if in_features and out_features: self.state.update({ "in_features": in_features, "out_features": out_features, }) def parameters(self, recurse: bool = True): """Delegates parameter iteration to the wrapped module.""" return self.module.parameters(recurse=recurse) def named_parameters(self, prefix: str = "", recurse: bool = True): """Delegates named parameter iteration to the wrapped module.""" return self.module.named_parameters(prefix=prefix, recurse=recurse) def buffers(self, recurse: bool = True): """Delegates buffer iteration to the wrapped module.""" return self.module.buffers(recurse=recurse) def named_buffers(self, prefix: str = "", recurse: bool = True): """Delegates named buffer iteration to the wrapped module.""" return self.module.named_buffers(prefix=prefix, recurse=recurse) def register_forward_hook( self, *args, **kwargs ): """Registers a forward hook while holding the parent module lock.""" with self._parent_lock: return self.module.register_forward_hook(*args, **kwargs) def register_buffer( self, name: str, tensor: Optional[Tensor], persistent: bool = True ) -> None: """Registers a buffer on the wrapped module under the parent lock.""" with self._parent_lock: return self.module.register_buffer(name, tensor, persistent) def unregister_buffer(self, name: str): """Removes a buffer from the wrapped module if it exists.""" with self._parent_lock: if name in self.module._buffers: del self.module._buffers[name] if hasattr(self.module, name): delattr(self.module, name) def register_parameter(self, name: str, param: Optional[Parameter]) -> None: """Registers a parameter on the wrapped module under the parent lock.""" with self._parent_lock: return self.module.register_parameter(name, param) def unregister_parameter(self, name: str) -> None: """Removes a parameter from the wrapped module if it exists.""" with self._parent_lock: if name in self.module._parameters: del self.module._parameters[name] if hasattr(self.module, name): delattr(self.module, name) # return stats for mo # def stats(self) -> Dict[str, float]: # # -1 means no stats have yet to gathered for the stat property # return { # STAT_GPTQ_DURATION: self.state.get(STAT_GPTQ_DURATION, -1), # STAT_GPTQ_AVG_LOSS: self.state.get(STAT_GPTQ_AVG_LOSS, -1), # STAT_GPTQ_DAMP_PERCENT: self.state.get(STAT_GPTQ_DAMP_PERCENT, -1), # STAT_GPTQ_FWD_TIME: self.state.get(STAT_GPTQ_FWD_TIME, -1), # } # getattr is only called if python cannot find attr for `self` def __getattr__(self, name: str): """Falls back to the wrapped module while preserving lock discipline.""" lock = object.__getattribute__(self, "_parent_lock") module = object.__getattribute__(self, "module") with lock: return getattr(module, name) # setattr is always called by python even if attr exists in `self` def __setattr__(self, name: str, value: Any) -> None: """Routes non-wrapper attributes to the wrapped module under lock.""" if name in [ "module", "module_dtype", "name", "full_name", "layer_index", "state", "_parent_lock", "target_device", "register_buffer", "unregister_buffer", "register_parameter", "unregister_parameter", ]: object.__setattr__(self, name, value) return lock = object.__getattribute__(self, "_parent_lock") module = object.__getattribute__(self, "module") with lock: setattr(module, name, value) def stream_state_payload_to_cpu( self, tensors: Dict[str, torch.Tensor], ) -> Dict[str, torch.Tensor]: """Streams arbitrary state tensors to CPU and records them in `state`.""" state_lock = self._parent_lock return stream_tensor_dict_to_cpu( tensors, store_callback=lambda host_map: self.state.update(host_map), state=self.state, state_lock=state_lock, ) def stream_parameters_to_cpu(self) -> Dict[str, torch.Tensor]: """Streams this module's direct parameters to CPU-backed state storage.""" state_lock = self._parent_lock tensor_map = dict(self.module.named_parameters(recurse=False)) return stream_tensor_dict_to_cpu( tensor_map, store_callback=lambda host_map: self.state.setdefault("parameters_cpu", {}).update(host_map), state=self.state, state_lock=state_lock, ) def stream_buffers_to_cpu(self) -> Dict[str, torch.Tensor]: """Streams this module's direct buffers to CPU-backed state storage.""" state_lock = self._parent_lock tensor_map = dict(self.module.named_buffers(recurse=False)) return stream_tensor_dict_to_cpu( tensor_map, store_callback=lambda host_map: self.state.setdefault("buffers_cpu", {}).update(host_map), state=self.state, state_lock=state_lock, ) def stream_all_to_cpu(self) -> Dict[str, Dict[str, torch.Tensor]]: """Streams direct parameters and buffers to CPU in one call.""" params = self.stream_parameters_to_cpu() buffers = self.stream_buffers_to_cpu() return {"parameters": params, "buffers": buffers} def stream_sync(self) -> None: """Waits for any outstanding asynchronous stream-to-CPU transfers.""" stream_sync_events(self.state, self._parent_lock)