From 3043becc2063cf9b5a19694dcb516d90a8db9d29 Mon Sep 17 00:00:00 2001 From: xlycae Date: Wed, 1 Jul 2026 16:46:31 +0800 Subject: [PATCH 1/2] Optimize offload weight pinning with mmap host registration --- .../common/ops/embedding/embedding_weight.py | 24 +- lightx2v/common/ops/mm/mm_weight.py | 65 ++--- lightx2v/common/ops/norm/layer_norm_weight.py | 9 +- lightx2v/common/ops/norm/rms_norm_weight.py | 5 +- lightx2v/common/ops/tensor/tensor.py | 24 +- lightx2v/common/ops/utils.py | 247 +++++++++++++++++- lightx2v_platform/ops/mm/template.py | 101 ++----- lightx2v_platform/ops/norm/norm_template.py | 53 ++-- test_cases/test_mmap_pinned_weights.py | 233 +++++++++++++++++ 9 files changed, 569 insertions(+), 192 deletions(-) create mode 100644 test_cases/test_mmap_pinned_weights.py diff --git a/lightx2v/common/ops/embedding/embedding_weight.py b/lightx2v/common/ops/embedding/embedding_weight.py index 57bdf99a2..02e19e5e0 100755 --- a/lightx2v/common/ops/embedding/embedding_weight.py +++ b/lightx2v/common/ops/embedding/embedding_weight.py @@ -2,10 +2,10 @@ from abc import ABCMeta from pathlib import Path -import torch import torch.nn.functional as F from safetensors import safe_open +from lightx2v.common.ops.utils import create_pin_tensor, move_tensor_back_to_cpu from lightx2v.utils.envs import * from lightx2v.utils.registry_factory import EMBEDDING_WEIGHT_REGISTER from lightx2v_platform.base.global_var import AI_DEVICE @@ -48,32 +48,29 @@ def _get_weight_tensor(self, weight_dict=None, use_infer_dtype=False): lazy_load_file_path = os.path.join(self.lazy_load_file, f"block_{self.weight_name.split('.')[1]}.safetensors") with safe_open(lazy_load_file_path, framework="pt", device="cpu") as lazy_load_file: tensor = lazy_load_file.get_tensor(self.weight_name) - if use_infer_dtype: - tensor = tensor.to(self.infer_dtype) else: tensor = weight_dict[self.weight_name] + if use_infer_dtype: + tensor = tensor.to(self.infer_dtype) return tensor - def _create_cpu_pin_weight(self, tensor): - pin_tensor = torch.empty(tensor.shape, pin_memory=True, dtype=tensor.dtype) - pin_tensor.copy_(tensor) - del tensor - return pin_tensor + def _create_cpu_pin_weight(self, tensor, dtype=None): + return create_pin_tensor(tensor, dtype=dtype) def _load_cuda_buffer(self, weight_dict): weight_tensor = self._get_weight_tensor(weight_dict, use_infer_dtype=self.lazy_load) self.weight_cuda_buffer = weight_tensor.to(AI_DEVICE) def _load_cpu_pin_buffer(self): - weight_tensor = self._get_weight_tensor(use_infer_dtype=True) - self.pin_weight = self._create_cpu_pin_weight(weight_tensor) + weight_tensor = self._get_weight_tensor() + self.pin_weight = self._create_cpu_pin_weight(weight_tensor, dtype=self.infer_dtype) def to_cuda(self, non_blocking=False): self.weight = self.pin_weight.to(AI_DEVICE, non_blocking=non_blocking) def to_cpu(self, non_blocking=False): if hasattr(self, "pin_weight"): - self.weight = self.pin_weight.copy_(self.weight, non_blocking=non_blocking).cpu() + move_tensor_back_to_cpu(self, "weight", non_blocking=non_blocking) else: self.weight = self.weight.to("cpu", non_blocking=non_blocking) @@ -105,9 +102,8 @@ def load_state_dict_from_disk(self, block_index, adapter_block_index=None): else: lazy_load_file_path = os.path.join(self.lazy_load_file, f"block_{block_index}.safetensors") with safe_open(lazy_load_file_path, framework="pt", device="cpu") as lazy_load_file: - weight_tensor = lazy_load_file.get_tensor(self.weight_name).to(self.infer_dtype) - self.pin_weight = self.pin_weight.copy_(weight_tensor) - del weight_tensor + weight_tensor = lazy_load_file.get_tensor(self.weight_name) + self.pin_weight = create_pin_tensor(weight_tensor, dtype=self.infer_dtype) @EMBEDDING_WEIGHT_REGISTER("Default") diff --git a/lightx2v/common/ops/mm/mm_weight.py b/lightx2v/common/ops/mm/mm_weight.py index be61890e9..101c45cdf 100755 --- a/lightx2v/common/ops/mm/mm_weight.py +++ b/lightx2v/common/ops/mm/mm_weight.py @@ -365,14 +365,12 @@ def load_state_dict_from_disk(self, block_index, adapter_block_index=None): lazy_load_file_path = get_lazy_load_file_path(self.lazy_load_file, self.weight_name) with safe_open(lazy_load_file_path, framework="pt", device="cpu") as lazy_load_file: - weight_tensor = lazy_load_file.get_tensor(self.weight_name).t() - self.pin_weight = self.pin_weight.copy_(weight_tensor) - del weight_tensor + weight_tensor = lazy_load_file.get_tensor(self.weight_name) + self.pin_weight = create_pin_tensor(weight_tensor, transpose=True) if self.bias_name is not None: bias_tensor = lazy_load_file.get_tensor(self.bias_name) - self.pin_bias.copy_(bias_tensor) - del bias_tensor + self.pin_bias = create_pin_tensor(bias_tensor) @MM_WEIGHT_REGISTER("Default-ForceFp32") @@ -579,22 +577,17 @@ def load_state_dict_from_disk(self, block_index, adapter_block_index=None): lazy_load_file_path = get_lazy_load_file_path(self.lazy_load_file, self.weight_name) with safe_open(lazy_load_file_path, framework="pt", device="cpu") as lazy_load_file: - if self.weight_need_transpose: - weight_tensor = lazy_load_file.get_tensor(self.weight_name).t() - else: - weight_tensor = lazy_load_file.get_tensor(self.weight_name) - - self.pin_weight = self.pin_weight.copy_(weight_tensor) - del weight_tensor + weight_tensor = lazy_load_file.get_tensor(self.weight_name) + self.pin_weight = create_pin_tensor(weight_tensor, transpose=self.weight_need_transpose) weight_scale_tensor = lazy_load_file.get_tensor(self.weight_scale_name) - self.pin_weight_scale = self.pin_weight_scale.copy_(weight_scale_tensor) - del weight_scale_tensor + scale_dtype = torch.float32 if self.scale_force_fp32 else None + self.pin_weight_scale = create_pin_tensor(weight_scale_tensor, dtype=scale_dtype) if self.bias_name is not None: bias_tensor = lazy_load_file.get_tensor(self.bias_name) - self.pin_bias.copy_(bias_tensor) - del bias_tensor + bias_dtype = torch.float32 if self.bias_force_fp32 else self.infer_dtype + self.pin_bias = create_pin_tensor(bias_tensor, dtype=bias_dtype) def per_block_cast_to_fp8(self, x): assert x.dim() == 2 @@ -1148,23 +1141,15 @@ def _get_cpu_pin_bias_tensor(self, source, is_lazy): ) with safe_open(lazy_load_file_path, framework="pt", device="cpu") as source: bias_tensor = source.get_tensor(self.bias_name) - if not self.bias_force_fp32: - bias_tensor = bias_tensor.to(self.infer_dtype) - if self.bias_force_fp32: - bias_tensor = bias_tensor.to(torch.float32) - return self._create_pin_tensor(bias_tensor) + bias_dtype = torch.float32 if self.bias_force_fp32 else self.infer_dtype + return self._create_pin_tensor(bias_tensor, dtype=bias_dtype) else: bias_tensor = source[self.bias_name] - if self.bias_force_fp32: - bias_tensor = bias_tensor.to(torch.float32) - return self._create_pin_tensor(bias_tensor) + bias_dtype = torch.float32 if self.bias_force_fp32 else None + return self._create_pin_tensor(bias_tensor, dtype=bias_dtype) def _create_pin_tensor(self, tensor, dtype=None): - dtype = dtype or tensor.dtype - pin_tensor = torch.empty(tensor.shape, pin_memory=True, dtype=dtype) - pin_tensor.copy_(tensor) - del tensor - return pin_tensor + return create_pin_tensor(tensor, dtype=dtype) def _load_default_tensors(self, weight_dict): if not self.lazy_load: @@ -1261,13 +1246,13 @@ def to_cuda(self, non_blocking=False): def to_cpu(self, non_blocking=False): if hasattr(self, "pin_weight"): - self.weight = self.pin_weight.copy_(self.weight, non_blocking=non_blocking).cpu() + move_tensor_back_to_cpu(self, "weight", non_blocking=non_blocking) if hasattr(self, "weight_scale_name"): - self.weight_scale = self.pin_weight_scale.copy_(self.weight_scale, non_blocking=non_blocking).cpu() - self.input_global_scale = self.pin_input_global_scale.copy_(self.input_global_scale, non_blocking=non_blocking).cpu() - self.alpha = self.pin_alpha.copy_(self.alpha, non_blocking=non_blocking).cpu() + move_tensor_back_to_cpu(self, "weight_scale", non_blocking=non_blocking) + move_tensor_back_to_cpu(self, "input_global_scale", non_blocking=non_blocking) + move_tensor_back_to_cpu(self, "alpha", non_blocking=non_blocking) if self.bias is not None: - self.bias = self.pin_bias.copy_(self.bias, non_blocking=non_blocking).cpu() + move_tensor_back_to_cpu(self, "bias", non_blocking=non_blocking) else: self.weight = self.weight.to("cpu", non_blocking=non_blocking) if hasattr(self, "weight_scale"): @@ -1321,17 +1306,11 @@ def load_state_dict_from_disk(self, block_index, adapter_block_index=None): lazy_load_file_path = get_lazy_load_file_path(self.lazy_load_file, self.weight_name) with safe_open(lazy_load_file_path, framework="pt", device="cpu") as lazy_load_file: - if self.weight_need_transpose: - weight_tensor = lazy_load_file.get_tensor(self.weight_name).t() - else: - weight_tensor = lazy_load_file.get_tensor(self.weight_name) - - self.pin_weight = self.pin_weight.copy_(weight_tensor) - del weight_tensor + weight_tensor = lazy_load_file.get_tensor(self.weight_name) + self.pin_weight = create_pin_tensor(weight_tensor, transpose=self.weight_need_transpose) weight_scale_tensor = lazy_load_file.get_tensor(self.weight_scale_name) - self.pin_weight_scale = self.pin_weight_scale.copy_(weight_scale_tensor) - del weight_scale_tensor + self.pin_weight_scale = create_pin_tensor(weight_scale_tensor) @MM_WEIGHT_REGISTER("CalibMax") diff --git a/lightx2v/common/ops/norm/layer_norm_weight.py b/lightx2v/common/ops/norm/layer_norm_weight.py index 732d9f273..3eb2407c8 100755 --- a/lightx2v/common/ops/norm/layer_norm_weight.py +++ b/lightx2v/common/ops/norm/layer_norm_weight.py @@ -168,14 +168,13 @@ def load_state_dict_from_disk(self, block_index, adapter_block_index=None): self.is_post_adapter, ) with safe_open(lazy_load_file_path, framework="pt", device="cpu") as lazy_load_file: - weight_tensor = lazy_load_file.get_tensor(self.weight_name).to(self.infer_dtype) - self.pin_weight = self.pin_weight.copy_(weight_tensor) + weight_tensor = lazy_load_file.get_tensor(self.weight_name) + self.pin_weight = create_pin_tensor(weight_tensor, dtype=self.infer_dtype) if self.bias_name is not None: - bias_tensor = lazy_load_file.get_tensor(self.bias_name).to(self.infer_dtype) - self.pin_bias = self.pin_bias.copy_(bias_tensor) + bias_tensor = lazy_load_file.get_tensor(self.bias_name) + self.pin_bias = create_pin_tensor(bias_tensor, dtype=self.infer_dtype) else: self.pin_bias = None - del weight_tensor else: self.weight = None self.bias = None diff --git a/lightx2v/common/ops/norm/rms_norm_weight.py b/lightx2v/common/ops/norm/rms_norm_weight.py index 340b747aa..07b3c0d78 100755 --- a/lightx2v/common/ops/norm/rms_norm_weight.py +++ b/lightx2v/common/ops/norm/rms_norm_weight.py @@ -138,9 +138,8 @@ def load_state_dict_from_disk(self, block_index, adapter_block_index=None): self.weight_name = resolve_block_name(self.weight_name, block_index, adapter_block_index, self.is_post_adapter) lazy_load_file_path = get_lazy_load_file_path(self.lazy_load_file, self.weight_name) with safe_open(lazy_load_file_path, framework="pt", device="cpu") as lazy_load_file: - weight_tensor = lazy_load_file.get_tensor(self.weight_name).to(self.infer_dtype) - self.pin_weight = self.pin_weight.copy_(weight_tensor) - del weight_tensor + weight_tensor = lazy_load_file.get_tensor(self.weight_name) + self.pin_weight = create_pin_tensor(weight_tensor, dtype=self.infer_dtype) @abstractmethod def apply(self, input_tensor): diff --git a/lightx2v/common/ops/tensor/tensor.py b/lightx2v/common/ops/tensor/tensor.py index b1ab7b3a7..3f998b521 100755 --- a/lightx2v/common/ops/tensor/tensor.py +++ b/lightx2v/common/ops/tensor/tensor.py @@ -2,9 +2,9 @@ import re from pathlib import Path -import torch from safetensors import safe_open +from lightx2v.common.ops.utils import create_pin_tensor, move_tensor_back_to_cpu from lightx2v.utils.envs import * from lightx2v.utils.registry_factory import TENSOR_REGISTER from lightx2v_platform.base.global_var import AI_DEVICE @@ -48,25 +48,22 @@ def _get_tensor(self, weight_dict=None, use_infer_dtype=False): lazy_load_file_path = os.path.join(self.lazy_load_file, f"block_{self.tensor_name.split('.')[1]}.safetensors") with safe_open(lazy_load_file_path, framework="pt", device="cpu") as lazy_load_file: tensor = lazy_load_file.get_tensor(self.tensor_name) - if use_infer_dtype: - tensor = tensor.to(self.infer_dtype) else: tensor = weight_dict[self.tensor_name] + if use_infer_dtype: + tensor = tensor.to(self.infer_dtype) return tensor - def _create_cpu_pin_tensor(self, tensor): - pin_tensor = torch.empty(tensor.shape, pin_memory=True, dtype=tensor.dtype) - pin_tensor.copy_(tensor) - del tensor - return pin_tensor + def _create_cpu_pin_tensor(self, tensor, dtype=None): + return create_pin_tensor(tensor, dtype=dtype) def _load_cuda_buffer(self, weight_dict): tensor = self._get_tensor(weight_dict, use_infer_dtype=self.lazy_load) self.tensor_cuda_buffer = tensor.to(AI_DEVICE) def _load_cpu_pin_buffer(self): - tensor = self._get_tensor(use_infer_dtype=True) - self.pin_tensor = self._create_cpu_pin_tensor(tensor) + tensor = self._get_tensor() + self.pin_tensor = self._create_cpu_pin_tensor(tensor, dtype=self.infer_dtype) def to_cuda(self, non_blocking=False): if hasattr(self, "pin_tensor"): @@ -78,7 +75,7 @@ def to_cuda(self, non_blocking=False): def to_cpu(self, non_blocking=False): if hasattr(self, "pin_tensor"): - self.tensor = self.pin_tensor.copy_(self.tensor, non_blocking=non_blocking).cpu() + move_tensor_back_to_cpu(self, "tensor", non_blocking=non_blocking) else: self.tensor = self.tensor.to("cpu", non_blocking=non_blocking) @@ -110,6 +107,5 @@ def load_state_dict_from_disk(self, block_index, adapter_block_index=None): else: lazy_load_file_path = os.path.join(self.lazy_load_file, f"block_{block_index}.safetensors") with safe_open(lazy_load_file_path, framework="pt", device="cpu") as lazy_load_file: - tensor = lazy_load_file.get_tensor(self.tensor_name).to(self.infer_dtype) - self.pin_tensor = self.pin_tensor.copy_(tensor) - del tensor + tensor = lazy_load_file.get_tensor(self.tensor_name) + self.pin_tensor = create_pin_tensor(tensor, dtype=self.infer_dtype) diff --git a/lightx2v/common/ops/utils.py b/lightx2v/common/ops/utils.py index fc6a1a581..266b62af1 100755 --- a/lightx2v/common/ops/utils.py +++ b/lightx2v/common/ops/utils.py @@ -1,5 +1,7 @@ import os import re +import threading +import weakref from pathlib import Path import torch @@ -9,6 +11,151 @@ from lightx2v.utils.envs import * from lightx2v_platform.base.global_var import AI_DEVICE +_CUDA_HOST_REGISTER_DEFAULT = 0 +_CUDA_HOST_REGISTER_FINALIZER_ATTR = "_lightx2v_cuda_host_register_finalizer" +_READONLY_PINNED_SOURCE_ATTR = "_lightx2v_readonly_pinned_source" +_PINNED_WEIGHT_STATS = { + "registered_count": 0, + "registered_bytes": 0, + "fallback_count": 0, + "fallback_bytes": 0, +} +_PINNED_WEIGHT_FALLBACK_WARNED = False +_PINNED_WEIGHT_STATS_LOCK = threading.Lock() + + +def reset_pinned_weight_stats(): + """Reset mmap pinned weight registration counters. + + This helper is used by tests and validation scripts before measuring the + registration/fallback path of weight loading. + """ + global _PINNED_WEIGHT_FALLBACK_WARNED + with _PINNED_WEIGHT_STATS_LOCK: + for key in _PINNED_WEIGHT_STATS: + _PINNED_WEIGHT_STATS[key] = 0 + _PINNED_WEIGHT_FALLBACK_WARNED = False + + +def get_pinned_weight_stats(): + """Get mmap pinned weight registration counters. + + Returns: + dict: Snapshot of registration and fallback counters. + """ + with _PINNED_WEIGHT_STATS_LOCK: + return dict(_PINNED_WEIGHT_STATS) + + +def _record_pinned_weight_fallback(tensor): + """Record one fallback from host registration to copy-based pinning. + + Args: + tensor: Source tensor that could not be registered in place. + """ + with _PINNED_WEIGHT_STATS_LOCK: + _PINNED_WEIGHT_STATS["fallback_count"] += 1 + fallback_bytes = tensor.numel() * tensor.element_size() + try: + fallback_bytes = tensor.untyped_storage().nbytes() + except Exception: + pass + with _PINNED_WEIGHT_STATS_LOCK: + _PINNED_WEIGHT_STATS["fallback_bytes"] += fallback_bytes + + +def _cuda_error_text(error): + """Convert a CUDA runtime error code to readable text. + + Args: + error: CUDA runtime error code returned by cudart. + + Returns: + Error message from CUDA runtime, or the raw error converted to string + if CUDA error text lookup is unavailable. + """ + try: + return torch.cuda.cudart().cudaGetErrorString(error) + except Exception: + return str(error) + + +def cuda_register_host_tensor(tensor): + """Register an existing CPU tensor storage as CUDA pinned host memory. + + Args: + tensor: CPU tensor whose underlying storage should be registered with + ``cudaHostRegister``. + + Returns: + The input tensor after its storage has been registered as pinned host + memory. The returned tensor is marked as a read-only pinned source so + offload code will not write GPU data back into mmap-backed storage. + + Raises: + RuntimeError: If the tensor is not on CPU, CUDA is unavailable, the + current AI device is not CUDA, or CUDA registration fails. + """ + if tensor.device.type != "cpu": + raise RuntimeError(f"cudaHostRegister requires a CPU tensor, got {tensor.device}") + if tensor.numel() == 0: + return tensor + if AI_DEVICE != "cuda": + raise RuntimeError(f"cudaHostRegister is only enabled for cuda AI_DEVICE, got {AI_DEVICE}") + if not torch.cuda.is_available(): + raise RuntimeError("cudaHostRegister requires CUDA to be available") + + storage = tensor.untyped_storage() + finalizer = getattr(storage, _CUDA_HOST_REGISTER_FINALIZER_ATTR, None) + if finalizer is not None and finalizer.alive: + if not tensor.is_pinned(): + raise RuntimeError("storage has a cudaHostRegister marker but tensor.is_pinned() is false") + # Keep a Python reference to the registered storage for async H2D users. + tensor._lightx2v_cuda_host_registered_storage = storage + setattr(tensor, _READONLY_PINNED_SOURCE_ATTR, True) + return tensor + if tensor.is_pinned(): + return tensor + + ptr = storage.data_ptr() + nbytes = storage.nbytes() + cudart = torch.cuda.cudart() + error = cudart.cudaHostRegister(ptr, nbytes, _CUDA_HOST_REGISTER_DEFAULT) + if error != 0: + raise RuntimeError(f"cudaHostRegister failed for {nbytes} bytes at 0x{ptr:x}: {_cuda_error_text(error)}") + + def unregister(address): + if torch.cuda.is_available(): + torch.cuda.cudart().cudaHostUnregister(address) + + # Tie cudaHostUnregister to storage lifetime, not to a temporary tensor view. + finalizer = weakref.finalize(storage, unregister, ptr) + setattr(storage, _CUDA_HOST_REGISTER_FINALIZER_ATTR, finalizer) + # The tensor attribute keeps storage alive after safetensors safe_open closes. + tensor._lightx2v_cuda_host_registered_storage = storage + setattr(tensor, _READONLY_PINNED_SOURCE_ATTR, True) + if not tensor.is_pinned(): + finalizer.detach() + cudart.cudaHostUnregister(ptr) + raise RuntimeError("cudaHostRegister succeeded but tensor.is_pinned() is false") + + with _PINNED_WEIGHT_STATS_LOCK: + _PINNED_WEIGHT_STATS["registered_count"] += 1 + _PINNED_WEIGHT_STATS["registered_bytes"] += nbytes + return tensor + + +def is_readonly_pinned_source(tensor): + """Check whether a tensor is a registered read-only pinned source. + + Args: + tensor: Tensor to inspect. + + Returns: + bool: True if the tensor should only be used as a CPU-to-device source. + """ + return bool(getattr(tensor, _READONLY_PINNED_SOURCE_ATTR, False)) + def resolve_block_name(name, block_index, adapter_block_index=None, is_post_adapter=False): """Resolve the name according to the block index, replacing the block index in the name with the specified block_index. @@ -74,16 +221,59 @@ def get_source_tensor(source_name, weight_dict, lazy_load, lazy_load_file, use_i def create_pin_tensor(tensor, transpose=False, dtype=None): - """Create a tensor with pinned memory for faster data transfer to GPU. + """Create a CPU tensor with pinned memory for faster data transfer to GPU. Args: - tensor: Source tensor to be converted to pinned memory - transpose: Whether to transpose the tensor after creating pinned memory (optional) - dtype: Target data type of the pinned tensor (optional, defaults to source tensor's dtype) + tensor: Source tensor to be converted to pinned memory. + transpose: Whether to transpose the pinned tensor (optional). + dtype: Target data type of the pinned tensor (optional, defaults to + source tensor's dtype). Returns: - Pinned memory tensor (on CPU) with optional transposition applied. - Falls back to regular CPU tensor if pinned memory allocation fails. + Pinned memory tensor on CPU with optional transposition applied. When + the source tensor is a CPU tensor and no dtype conversion is needed, the + source storage is registered in place with ``cudaHostRegister``. If + registration fails or dtype conversion is required, a writable pinned + copy is created instead. The copy path falls back to regular CPU memory + if pinned allocation fails. + """ + global _PINNED_WEIGHT_FALLBACK_WARNED + + dtype = dtype or tensor.dtype + if tensor.device.type == "cpu" and dtype == tensor.dtype: + try: + pin_tensor = cuda_register_host_tensor(tensor) + except Exception as e: + _record_pinned_weight_fallback(tensor) + if not _PINNED_WEIGHT_FALLBACK_WARNED: + logger.warning(f"Failed to register mmap-backed pinned weight; falling back to copy-based pinned memory. First error: {e}") + _PINNED_WEIGHT_FALLBACK_WARNED = True + else: + if transpose: + base_tensor = pin_tensor + pin_tensor = base_tensor.t() + # A transposed view does not own storage; keep the registered + # base tensor alive for async H2D copies. + pin_tensor._lightx2v_cuda_host_registered_base = base_tensor + if is_readonly_pinned_source(base_tensor): + setattr(pin_tensor, _READONLY_PINNED_SOURCE_ATTR, True) + return pin_tensor + + return create_writable_pin_tensor(tensor, transpose=transpose, dtype=dtype) + + +def create_writable_pin_tensor(tensor, transpose=False, dtype=None): + """Create a copy-based CPU buffer that may be written by later copy_ calls. + + Args: + tensor: Source tensor to copy into the new buffer. + transpose: Whether to transpose the copied tensor (optional). + dtype: Target data type of the copied tensor (optional, defaults to + source tensor's dtype). + + Returns: + Writable pinned CPU tensor. Falls back to regular CPU memory if pinned + allocation fails. """ dtype = dtype or tensor.dtype try: @@ -98,6 +288,34 @@ def create_pin_tensor(tensor, transpose=False, dtype=None): return pin_tensor +def move_tensor_back_to_cpu(obj, attr_name, non_blocking=False): + """Move a device tensor attribute back to its CPU-side attribute. + + Args: + obj: Object containing ``attr_name`` and optional ``pin_``. + attr_name: Name of the tensor attribute to move back to CPU. + non_blocking: Whether to perform non-blocking data transfer (optional). + + Notes: + Read-only pinned sources are mmap-backed canonical CPU weights. They + are reused directly on CPU and are not overwritten by GPU results. + Copy-based pinned buffers remain writable and keep the previous + offload behavior of copying the device tensor back into ``pin_*``. + """ + pin_attr_name = f"pin_{attr_name}" + value = getattr(obj, attr_name, None) + if hasattr(obj, pin_attr_name) and getattr(obj, pin_attr_name) is not None: + pin_tensor = getattr(obj, pin_attr_name) + if is_readonly_pinned_source(pin_tensor): + setattr(obj, attr_name, pin_tensor) + elif value is not None: + setattr(obj, attr_name, pin_tensor.copy_(value, non_blocking=non_blocking).cpu()) + else: + setattr(obj, attr_name, pin_tensor) + elif value is not None: + setattr(obj, attr_name, value.to("cpu", non_blocking=non_blocking)) + + def get_lazy_load_file_path(lazy_load_file, weight_name_for_block=None): """Get the full file path for lazy loading, handling both file and directory inputs. @@ -161,10 +379,16 @@ def create_cpu_buffers(base_attrs, lazy_load_file, use_infer_dtype=False, scale_ """ result = {} - # Use get_source_tensor to load the tensor (weight_dict is not required when lazy_load=True) for name, attr_name, transpose in base_attrs: - tensor = get_source_tensor(name, {}, lazy_load=True, lazy_load_file=lazy_load_file, use_infer_dtype=use_infer_dtype, scale_force_fp32=scale_force_fp32, bias_force_fp32=bias_force_fp32) - result[attr_name] = create_pin_tensor(tensor, transpose=transpose) + tensor = get_source_tensor(name, {}, lazy_load=True, lazy_load_file=lazy_load_file, use_infer_dtype=False, scale_force_fp32=False, bias_force_fp32=False) + dtype = None + if use_infer_dtype: + dtype = GET_DTYPE() + elif scale_force_fp32 and "weight_scale" in name: + dtype = torch.float32 + elif bias_force_fp32 and "bias" in name: + dtype = torch.float32 + result[attr_name] = create_pin_tensor(tensor, transpose=transpose, dtype=dtype) return result @@ -223,7 +447,10 @@ def move_tensor_to_device(obj, attr_name, target_device, non_blocking=False, use if hasattr(obj, pin_attr_name) and getattr(obj, pin_attr_name) is not None: pin_tensor = getattr(obj, pin_attr_name) if hasattr(obj, attr_name) and getattr(obj, attr_name) is not None and use_copy: - setattr(obj, attr_name, pin_tensor.copy_(getattr(obj, attr_name), non_blocking=non_blocking).to(target_device)) + if target_device == "cpu": + move_tensor_back_to_cpu(obj, attr_name, non_blocking=non_blocking) + else: + setattr(obj, attr_name, pin_tensor.copy_(getattr(obj, attr_name), non_blocking=non_blocking).to(target_device)) else: setattr(obj, attr_name, pin_tensor.to(target_device, non_blocking=non_blocking)) elif hasattr(obj, attr_name) and getattr(obj, attr_name) is not None: diff --git a/lightx2v_platform/ops/mm/template.py b/lightx2v_platform/ops/mm/template.py index 28f519aa2..3d889c793 100644 --- a/lightx2v_platform/ops/mm/template.py +++ b/lightx2v_platform/ops/mm/template.py @@ -3,6 +3,7 @@ import torch +from lightx2v.common.ops.utils import create_pin_tensor, move_tensor_back_to_cpu from lightx2v_platform.base.global_var import AI_DEVICE @@ -39,11 +40,11 @@ def to_cuda(self, non_blocking=False): def to_cpu(self, non_blocking=False): if hasattr(self, "pin_weight"): - self.weight = self.pin_weight.copy_(self.weight, non_blocking=non_blocking).cpu() + move_tensor_back_to_cpu(self, "weight", non_blocking=non_blocking) if hasattr(self, "weight_scale_name"): - self.weight_scale = self.pin_weight_scale.copy_(self.weight_scale, non_blocking=non_blocking).cpu() + move_tensor_back_to_cpu(self, "weight_scale", non_blocking=non_blocking) if self.bias is not None: - self.bias = self.pin_bias.copy_(self.bias, non_blocking=non_blocking).cpu() + move_tensor_back_to_cpu(self, "bias", non_blocking=non_blocking) else: self.weight = self.weight.to("cpu", non_blocking=non_blocking) if hasattr(self, "weight_scale"): @@ -75,7 +76,7 @@ def load(self, weight_dict): if hasattr(self, "weight") and self.weight is not None: self.weight = self.weight.t() if hasattr(self, "pin_weight") and self.pin_weight is not None: - self.pin_weight = self.pin_weight.t() + self.pin_weight = create_pin_tensor(self.pin_weight, transpose=True) if hasattr(self, "weight_cuda_buffer") and self.weight_cuda_buffer is not None: self.weight_cuda_buffer = self.weight_cuda_buffer.t() @@ -140,20 +141,18 @@ def _get_cpu_pin_bias_tensor(self, source, is_lazy): return None if is_lazy: bias_tensor = source.get_tensor(self.bias_name) - if not self.bias_force_fp32: - bias_tensor = bias_tensor.to(self.infer_dtype) else: bias_tensor = source[self.bias_name] if self.bias_force_fp32: - bias_tensor = bias_tensor.to(torch.float32) - return self._create_pin_tensor(bias_tensor) + bias_dtype = torch.float32 + elif is_lazy: + bias_dtype = self.infer_dtype + else: + bias_dtype = None + return self._create_pin_tensor(bias_tensor, dtype=bias_dtype) def _create_pin_tensor(self, tensor, dtype=None): - dtype = dtype or tensor.dtype - pin_tensor = torch.empty(tensor.shape, pin_memory=True, dtype=dtype) - pin_tensor.copy_(tensor) - del tensor - return pin_tensor + return create_pin_tensor(tensor, dtype=dtype) def _load_default_tensors(self, weight_dict): if not self.lazy_load: @@ -221,15 +220,8 @@ def load_mxfp4(self, weight_dict): else: device = weight_dict[self.weight_name].device if device.type == "cpu": - weight_shape = weight_dict[self.weight_name].shape - weight_dtype = weight_dict[self.weight_name].dtype - self.pin_weight = torch.empty(weight_shape, pin_memory=True, dtype=weight_dtype) - self.pin_weight.copy_(weight_dict[self.weight_name]) - - weight_scale_shape = weight_dict[self.weight_scale_name].shape - weight_scale_dtype = weight_dict[self.weight_scale_name].dtype - self.pin_weight_scale = torch.empty(weight_scale_shape, pin_memory=True, dtype=weight_scale_dtype) - self.pin_weight_scale.copy_(weight_dict[self.weight_scale_name]) + self.pin_weight = create_pin_tensor(weight_dict[self.weight_name]) + self.pin_weight_scale = create_pin_tensor(weight_dict[self.weight_scale_name]) del weight_dict[self.weight_name] else: self.weight = weight_dict[self.weight_name] @@ -244,15 +236,8 @@ def load_mxfp6(self, weight_dict): else: device = weight_dict[self.weight_name].device if device.type == "cpu": - weight_shape = weight_dict[self.weight_name].shape - weight_dtype = weight_dict[self.weight_name].dtype - self.pin_weight = torch.empty(weight_shape, pin_memory=True, dtype=weight_dtype) - self.pin_weight.copy_(weight_dict[self.weight_name]) - - weight_scale_shape = weight_dict[self.weight_scale_name].shape - weight_scale_dtype = weight_dict[self.weight_scale_name].dtype - self.pin_weight_scale = torch.empty(weight_scale_shape, pin_memory=True, dtype=weight_scale_dtype) - self.pin_weight_scale.copy_(weight_dict[self.weight_scale_name]) + self.pin_weight = create_pin_tensor(weight_dict[self.weight_name]) + self.pin_weight_scale = create_pin_tensor(weight_dict[self.weight_scale_name]) del weight_dict[self.weight_name] else: self.weight = weight_dict[self.weight_name] @@ -267,15 +252,8 @@ def load_mxfp8(self, weight_dict): else: device = weight_dict[self.weight_name].device if device.type == "cpu": - weight_shape = weight_dict[self.weight_name].shape - weight_dtype = weight_dict[self.weight_name].dtype - self.pin_weight = torch.empty(weight_shape, pin_memory=True, dtype=weight_dtype) - self.pin_weight.copy_(weight_dict[self.weight_name]) - - weight_scale_shape = weight_dict[self.weight_scale_name].shape - weight_scale_dtype = weight_dict[self.weight_scale_name].dtype - self.pin_weight_scale = torch.empty(weight_scale_shape, pin_memory=True, dtype=weight_scale_dtype) - self.pin_weight_scale.copy_(weight_dict[self.weight_scale_name]) + self.pin_weight = create_pin_tensor(weight_dict[self.weight_name]) + self.pin_weight_scale = create_pin_tensor(weight_dict[self.weight_scale_name]) del weight_dict[self.weight_name] else: self.weight = weight_dict[self.weight_name] @@ -290,25 +268,10 @@ def load_nvfp4(self, weight_dict): alpha = 1.0 / (input_global_scale * weight_global_scale) if device.type == "cpu": - weight_shape = weight_dict[self.weight_name].shape - weight_dtype = weight_dict[self.weight_name].dtype - self.pin_weight = torch.empty(weight_shape, pin_memory=True, dtype=weight_dtype) - self.pin_weight.copy_(weight_dict[self.weight_name]) - - weight_scale_shape = weight_dict[self.weight_scale_name].shape - weight_scale_dtype = weight_dict[self.weight_scale_name].dtype - self.pin_weight_scale = torch.empty(weight_scale_shape, pin_memory=True, dtype=weight_scale_dtype) - self.pin_weight_scale.copy_(weight_dict[self.weight_scale_name]) - - input_global_scale_shape = input_global_scale.shape - input_global_scale_dtype = input_global_scale.dtype - self.pin_input_global_scale = torch.empty(input_global_scale_shape, pin_memory=True, dtype=input_global_scale_dtype) - self.pin_input_global_scale.copy_(input_global_scale) - - alpha_shape = alpha.shape - alpha_dtype = alpha.dtype - self.pin_alpha = torch.empty(alpha_shape, pin_memory=True, dtype=alpha_dtype) - self.pin_alpha.copy_(alpha) + self.pin_weight = create_pin_tensor(weight_dict[self.weight_name]) + self.pin_weight_scale = create_pin_tensor(weight_dict[self.weight_scale_name]) + self.pin_input_global_scale = create_pin_tensor(input_global_scale) + self.pin_alpha = create_pin_tensor(alpha) del weight_dict[self.weight_name] else: @@ -323,10 +286,7 @@ def load_nvfp4(self, weight_dict): else: device = weight_dict[self.bias_name].device if device.type == "cpu": - bias_shape = weight_dict[self.bias_name].shape - bias_dtype = weight_dict[self.bias_name].dtype - self.pin_bias = torch.empty(bias_shape, pin_memory=True, dtype=bias_dtype) - self.pin_bias.copy_(weight_dict[self.bias_name]) + self.pin_bias = create_pin_tensor(weight_dict[self.bias_name]) else: self.bias = weight_dict[self.bias_name] else: @@ -448,16 +408,11 @@ def load_state_dict_from_disk(self, block_index, adapter_block_index=None): self.weight_name = re.sub(r"\.\d+", lambda m: f".{block_index}", self.weight_name, count=1) self.weight_scale_name = re.sub(r"\.\d+", lambda m: f".{block_index}", self.weight_scale_name, count=1) - if self.weight_need_transpose: - weight_tensor = self.lazy_load_file.get_tensor(self.weight_name).t() - else: - weight_tensor = self.lazy_load_file.get_tensor(self.weight_name) - self.pin_weight = self.pin_weight.copy_(weight_tensor) + weight_tensor = self.lazy_load_file.get_tensor(self.weight_name) + self.pin_weight = create_pin_tensor(weight_tensor, transpose=self.weight_need_transpose) weight_scale_tensor = self.lazy_load_file.get_tensor(self.weight_scale_name) - self.pin_weight_scale = self.pin_weight_scale.copy_(weight_scale_tensor) - - del weight_tensor + self.pin_weight_scale = create_pin_tensor(weight_scale_tensor, dtype=torch.float32) if self.bias_name is not None: if self.is_post_adapter: @@ -467,5 +422,5 @@ def load_state_dict_from_disk(self, block_index, adapter_block_index=None): self.bias_name = re.sub(r"\.\d+", lambda m: f".{block_index}", self.bias_name, count=1) bias_tensor = self.lazy_load_file.get_tensor(self.bias_name) - self.pin_bias.copy_(bias_tensor) - del bias_tensor + bias_dtype = torch.float32 if self.bias_force_fp32 else self.infer_dtype + self.pin_bias = create_pin_tensor(bias_tensor, dtype=bias_dtype) diff --git a/lightx2v_platform/ops/norm/norm_template.py b/lightx2v_platform/ops/norm/norm_template.py index 3775e5d73..90cf205b4 100644 --- a/lightx2v_platform/ops/norm/norm_template.py +++ b/lightx2v_platform/ops/norm/norm_template.py @@ -7,6 +7,7 @@ import torch from safetensors import safe_open +from lightx2v.common.ops.utils import create_pin_tensor, move_tensor_back_to_cpu from lightx2v_platform.base.global_var import AI_DEVICE DTYPE_MAP = { @@ -78,25 +79,22 @@ def _get_weight_tensor(self, weight_dict=None, use_infer_dtype=False): lazy_load_file_path = os.path.join(self.lazy_load_file, f"block_{self.weight_name.split('.')[1]}.safetensors") with safe_open(lazy_load_file_path, framework="pt", device="cpu") as lazy_load_file: tensor = lazy_load_file.get_tensor(self.weight_name) - if use_infer_dtype: - tensor = tensor.to(self.infer_dtype) else: tensor = weight_dict[self.weight_name] + if use_infer_dtype: + tensor = tensor.to(self.infer_dtype) return tensor - def _create_cpu_pin_weight(self, tensor): - pin_tensor = torch.empty(tensor.shape, pin_memory=True, dtype=tensor.dtype) - pin_tensor.copy_(tensor) - del tensor - return pin_tensor + def _create_cpu_pin_weight(self, tensor, dtype=None): + return create_pin_tensor(tensor, dtype=dtype) def _load_cuda_buffer(self, weight_dict): weight_tensor = self._get_weight_tensor(weight_dict, use_infer_dtype=self.lazy_load) self.weight_cuda_buffer = weight_tensor.to(AI_DEVICE) def _load_cpu_pin_buffer(self): - weight_tensor = self._get_weight_tensor(use_infer_dtype=True) - self.pin_weight = self._create_cpu_pin_weight(weight_tensor) + weight_tensor = self._get_weight_tensor() + self.pin_weight = self._create_cpu_pin_weight(weight_tensor, dtype=self.infer_dtype) @abstractmethod def apply(self, input_tensor): @@ -111,7 +109,7 @@ def to_cuda(self, non_blocking=False): def to_cpu(self, non_blocking=False): if hasattr(self, "pin_weight"): - self.weight = self.pin_weight.copy_(self.weight, non_blocking=non_blocking).cpu() + move_tensor_back_to_cpu(self, "weight", non_blocking=non_blocking) else: self.weight = self.weight.to("cpu", non_blocking=non_blocking) @@ -143,9 +141,8 @@ def load_state_dict_from_disk(self, block_index, adapter_block_index=None): else: lazy_load_file_path = os.path.join(self.lazy_load_file, f"block_{block_index}.safetensors") with safe_open(lazy_load_file_path, framework="pt", device="cpu") as lazy_load_file: - weight_tensor = lazy_load_file.get_tensor(self.weight_name).to(self.infer_dtype) - self.pin_weight = self.pin_weight.copy_(weight_tensor) - del weight_tensor + weight_tensor = lazy_load_file.get_tensor(self.weight_name) + self.pin_weight = create_pin_tensor(weight_tensor, dtype=self.infer_dtype) class LayerNormWeightTemplate(metaclass=ABCMeta): @@ -197,19 +194,16 @@ def _get_tensor(self, name, weight_dict=None, use_infer_dtype=False): lazy_load_file_path = os.path.join(self.lazy_load_file, f"block_{name.split('.')[1]}.safetensors") with safe_open(lazy_load_file_path, framework="pt", device="cpu") as lazy_load_file: tensor = lazy_load_file.get_tensor(name) - if use_infer_dtype: - tensor = tensor.to(self.infer_dtype) else: tensor = weight_dict[name] + if use_infer_dtype: + tensor = tensor.to(self.infer_dtype) return tensor def _create_cpu_pin_tensor(self, tensor): if tensor is None: return None - pin_tensor = torch.empty(tensor.shape, pin_memory=True, dtype=tensor.dtype) - pin_tensor.copy_(tensor) - del tensor - return pin_tensor + return create_pin_tensor(tensor) def _load_cuda_buffers(self, weight_dict): weight_tensor = self._get_tensor(self.weight_name, weight_dict, use_infer_dtype=self.lazy_load) @@ -221,15 +215,15 @@ def _load_cuda_buffers(self, weight_dict): self.bias_cuda_buffer = bias_tensor.to(AI_DEVICE) def _load_cpu_pin_buffers(self): - weight_tensor = self._get_tensor(self.weight_name, use_infer_dtype=True) + weight_tensor = self._get_tensor(self.weight_name) if weight_tensor is not None: - self.pin_weight = self._create_cpu_pin_tensor(weight_tensor) + self.pin_weight = create_pin_tensor(weight_tensor, dtype=self.infer_dtype) else: self.weight = None - bias_tensor = self._get_tensor(self.bias_name, use_infer_dtype=True) + bias_tensor = self._get_tensor(self.bias_name) if bias_tensor is not None: - self.pin_bias = self._create_cpu_pin_tensor(bias_tensor) + self.pin_bias = create_pin_tensor(bias_tensor, dtype=self.infer_dtype) else: self.bias = None self.pin_bias = None @@ -254,11 +248,11 @@ def to_cuda(self, non_blocking=False): def to_cpu(self, non_blocking=False): if hasattr(self, "pin_weight"): - self.weight = self.pin_weight.copy_(self.weight, non_blocking=non_blocking).cpu() + move_tensor_back_to_cpu(self, "weight", non_blocking=non_blocking) else: self.weight = self.weight.to("cpu", non_blocking=non_blocking) if hasattr(self, "pin_bias"): - self.bias = self.pin_bias.copy_(self.bias, non_blocking=non_blocking).cpu() + move_tensor_back_to_cpu(self, "bias", non_blocking=non_blocking) else: if self.bias is not None: self.bias = self.bias.to("cpu", non_blocking=non_blocking) @@ -299,9 +293,8 @@ def load_state_dict_from_disk(self, block_index, adapter_block_index=None): else: lazy_load_file_path = os.path.join(self.lazy_load_file, f"block_{block_index}.safetensors") with safe_open(lazy_load_file_path, framework="pt", device="cpu") as lazy_load_file: - weight_tensor = lazy_load_file.get_tensor(self.weight_name).to(self.infer_dtype) - self.pin_weight = self.pin_weight.copy_(weight_tensor) + weight_tensor = lazy_load_file.get_tensor(self.weight_name) + self.pin_weight = create_pin_tensor(weight_tensor, dtype=self.infer_dtype) if self.bias_name is not None: - bias_tensor = lazy_load_file.get_tensor(self.bias_name).to(self.infer_dtype) - self.pin_bias = self.pin_bias.copy_(bias_tensor) - del weight_tensor + bias_tensor = lazy_load_file.get_tensor(self.bias_name) + self.pin_bias = create_pin_tensor(bias_tensor, dtype=self.infer_dtype) diff --git a/test_cases/test_mmap_pinned_weights.py b/test_cases/test_mmap_pinned_weights.py new file mode 100644 index 000000000..819c620d7 --- /dev/null +++ b/test_cases/test_mmap_pinned_weights.py @@ -0,0 +1,233 @@ +import os +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest import mock + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +os.environ.setdefault("SKIP_PLATFORM_CHECK", "1") + + +def ensure_lightx2v_pipeline_stub(): + if "lightx2v.pipeline" not in sys.modules: + pipeline_stub = types.ModuleType("lightx2v.pipeline") + pipeline_stub.LightX2VPipeline = object + sys.modules["lightx2v.pipeline"] = pipeline_stub + + +class MmapPinnedWeightsTest(unittest.TestCase): + def setUp(self): + ensure_lightx2v_pipeline_stub() + from lightx2v.common.ops import utils + + utils.reset_pinned_weight_stats() + + def _write_safetensors(self, tensors): + tmpdir = tempfile.TemporaryDirectory() + path = Path(tmpdir.name) / "weights.safetensors" + save_file(tensors, str(path)) + self.addCleanup(tmpdir.cleanup) + return path + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for cudaHostRegister") + def test_safetensors_tensor_registers_without_copy(self): + from lightx2v.common.ops.utils import create_pin_tensor, get_pinned_weight_stats + + path = self._write_safetensors({"w": torch.arange(16, dtype=torch.float32).reshape(4, 4)}) + with safe_open(str(path), framework="pt", device="cpu") as f: + tensor = f.get_tensor("w") + original_ptr = tensor.data_ptr() + pinned = create_pin_tensor(tensor) + + self.assertEqual(pinned.data_ptr(), original_ptr) + self.assertTrue(pinned.is_pinned()) + stats = get_pinned_weight_stats() + self.assertEqual(stats["fallback_count"], 0) + self.assertGreaterEqual(stats["registered_bytes"], pinned.numel() * pinned.element_size()) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for cudaHostRegister") + def test_registered_tensor_survives_safe_open_close_for_h2d(self): + from lightx2v.common.ops.utils import create_pin_tensor, is_readonly_pinned_source + + source = torch.arange(16, dtype=torch.float32).reshape(4, 4) + path = self._write_safetensors({"w": source}) + with safe_open(str(path), framework="pt", device="cpu") as f: + pinned = create_pin_tensor(f.get_tensor("w")) + + self.assertTrue(is_readonly_pinned_source(pinned)) + self.assertTrue(torch.equal(pinned, source)) + cuda_copy = pinned.to("cuda", non_blocking=True) + torch.cuda.synchronize() + self.assertTrue(torch.equal(cuda_copy.cpu(), source)) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for cudaHostRegister") + def test_transposed_view_keeps_registered_base(self): + from lightx2v.common.ops.utils import create_pin_tensor + + path = self._write_safetensors({"w": torch.arange(12, dtype=torch.float32).reshape(3, 4)}) + with safe_open(str(path), framework="pt", device="cpu") as f: + tensor = f.get_tensor("w") + original_ptr = tensor.data_ptr() + pinned = create_pin_tensor(tensor, transpose=True) + + self.assertEqual(tuple(pinned.shape), (4, 3)) + self.assertEqual(pinned.data_ptr(), original_ptr) + self.assertTrue(pinned.is_pinned()) + self.assertTrue(hasattr(pinned, "_lightx2v_cuda_host_registered_base")) + + def test_dtype_conversion_falls_back_to_copy(self): + from lightx2v.common.ops.utils import create_pin_tensor, get_pinned_weight_stats + + source = torch.arange(8, dtype=torch.float32) + pinned = create_pin_tensor(source, dtype=torch.float16) + + self.assertEqual(pinned.dtype, torch.float16) + self.assertNotEqual(pinned.data_ptr(), source.data_ptr()) + stats = get_pinned_weight_stats() + self.assertEqual(stats["registered_count"], 0) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for cudaHostRegister") + def test_safetensors_dtype_conversion_does_not_register_converted_copy(self): + from lightx2v.common.ops.utils import create_pin_tensor, get_pinned_weight_stats + + path = self._write_safetensors({"w": torch.arange(8, dtype=torch.float32)}) + with safe_open(str(path), framework="pt", device="cpu") as f: + tensor = f.get_tensor("w") + original_ptr = tensor.data_ptr() + pinned = create_pin_tensor(tensor, dtype=torch.float16) + + self.assertEqual(pinned.dtype, torch.float16) + self.assertNotEqual(pinned.data_ptr(), original_ptr) + stats = get_pinned_weight_stats() + self.assertEqual(stats["registered_count"], 0) + self.assertEqual(stats["fallback_count"], 0) + + def test_registration_failure_falls_back(self): + from lightx2v.common.ops import utils + + source = torch.arange(8, dtype=torch.float32) + with mock.patch.object(utils, "cuda_register_host_tensor", side_effect=RuntimeError("forced failure")): + pinned = utils.create_pin_tensor(source) + + self.assertNotEqual(pinned.data_ptr(), source.data_ptr()) + self.assertEqual(utils.get_pinned_weight_stats()["fallback_count"], 1) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for cudaHostRegister") + def test_to_cpu_does_not_write_back_readonly_pinned_source(self): + from lightx2v.common.ops.utils import create_pin_tensor, move_tensor_back_to_cpu + + class Holder: + pass + + source = torch.arange(8, dtype=torch.float32) + path = self._write_safetensors({"w": source}) + with safe_open(str(path), framework="pt", device="cpu") as f: + holder = Holder() + holder.pin_weight = create_pin_tensor(f.get_tensor("w")) + holder.weight = torch.full((8,), 99.0, device="cuda") + + original_ptr = holder.pin_weight.data_ptr() + move_tensor_back_to_cpu(holder, "weight") + + self.assertEqual(holder.weight.data_ptr(), original_ptr) + self.assertTrue(torch.equal(holder.weight, source)) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for cudaHostRegister") + def test_mm_weight_to_cpu_does_not_write_back_readonly_source(self): + import lightx2v.common.ops.mm.mm_weight as mm_weight + + source = torch.arange(12, dtype=torch.float32).reshape(3, 4) + path = self._write_safetensors({"blocks.0.ffn.0.weight": source}) + weight = mm_weight.MMWeight( + "blocks.0.ffn.0.weight", + None, + create_cpu_buffer=True, + lazy_load=True, + lazy_load_file=str(path), + ) + weight.load({}) + original = weight.pin_weight.clone() + original_ptr = weight.pin_weight.data_ptr() + weight.weight = torch.full(tuple(weight.pin_weight.shape), 99.0, device="cuda") + + weight.to_cpu() + + self.assertEqual(weight.weight.data_ptr(), original_ptr) + self.assertTrue(torch.equal(weight.weight, original)) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for pinned copy-back") + def test_to_cpu_preserves_copy_back_for_writable_pin_buffer(self): + from lightx2v.common.ops.utils import create_writable_pin_tensor, move_tensor_back_to_cpu + + class Holder: + pass + + holder = Holder() + holder.pin_weight = create_writable_pin_tensor(torch.zeros(4, dtype=torch.float32)) + holder.weight = torch.arange(4, dtype=torch.float32, device="cuda") + move_tensor_back_to_cpu(holder, "weight") + + self.assertEqual(holder.weight.data_ptr(), holder.pin_weight.data_ptr()) + self.assertTrue(torch.equal(holder.weight, torch.arange(4, dtype=torch.float32))) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for cudaHostRegister") + def test_mm_weight_lazy_reload_uses_registered_mmap_tensor(self): + import lightx2v.common.ops.mm.mm_weight as mm_weight + + path = self._write_safetensors( + { + "blocks.0.ffn.0.weight": torch.arange(12, dtype=torch.float32).reshape(3, 4), + "blocks.0.ffn.0.bias": torch.arange(4, dtype=torch.float32), + } + ) + weight = mm_weight.MMWeight( + "blocks.0.ffn.0.weight", + "blocks.0.ffn.0.bias", + create_cpu_buffer=True, + lazy_load=True, + lazy_load_file=str(path), + ) + weight.load({}) + initial_ptr = weight.pin_weight.data_ptr() + + weight.load_state_dict_from_disk(0) + + self.assertEqual(tuple(weight.pin_weight.shape), (4, 3)) + self.assertTrue(weight.pin_weight.is_pinned()) + self.assertNotEqual(weight.pin_weight.data_ptr(), initial_ptr) + self.assertTrue(hasattr(weight.pin_weight, "_lightx2v_cuda_host_registered_base")) + self.assertTrue(weight.pin_bias.is_pinned()) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for cudaHostRegister") + def test_quant_lazy_reload_preserves_scale_and_bias_dtypes(self): + import lightx2v.common.ops.mm.mm_weight as mm_weight + + path = self._write_safetensors( + { + "blocks.0.attn.q.weight": torch.arange(12, dtype=torch.float32).reshape(3, 4), + "blocks.0.attn.q.weight_scale": torch.arange(4, dtype=torch.float16), + "blocks.0.attn.q.bias": torch.arange(4, dtype=torch.float32), + } + ) + weight = mm_weight.MMWeightWfp8channelAfp8channeldynamicVllm( + "blocks.0.attn.q.weight", + "blocks.0.attn.q.bias", + create_cpu_buffer=True, + lazy_load=True, + lazy_load_file=str(path), + ) + weight.bias_force_fp32 = False + + weight.load_state_dict_from_disk(0) + + self.assertEqual(weight.pin_weight_scale.dtype, torch.float32) + self.assertEqual(weight.pin_bias.dtype, weight.infer_dtype) + + +if __name__ == "__main__": + unittest.main() From 586aa15ccfbf4508447180b091c5e9269bc1af58 Mon Sep 17 00:00:00 2001 From: xlycae Date: Wed, 1 Jul 2026 17:26:02 +0800 Subject: [PATCH 2/2] Make pinned weight unregister finalizer robust --- lightx2v/common/ops/utils.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/lightx2v/common/ops/utils.py b/lightx2v/common/ops/utils.py index 266b62af1..0c0b9142c 100755 --- a/lightx2v/common/ops/utils.py +++ b/lightx2v/common/ops/utils.py @@ -124,9 +124,12 @@ def cuda_register_host_tensor(tensor): if error != 0: raise RuntimeError(f"cudaHostRegister failed for {nbytes} bytes at 0x{ptr:x}: {_cuda_error_text(error)}") - def unregister(address): - if torch.cuda.is_available(): - torch.cuda.cudart().cudaHostUnregister(address) + def unregister(address, is_available=torch.cuda.is_available, cudart=torch.cuda.cudart): + try: + if is_available(): + cudart().cudaHostUnregister(address) + except Exception: + pass # Tie cudaHostUnregister to storage lifetime, not to a temporary tensor view. finalizer = weakref.finalize(storage, unregister, ptr)