Skip to content

Optimize offload weight pinning with mmap host registration #1206

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

Open
Fatemanx wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Optimize offload weight pinning with mmap host registration #1206

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3043bec
Optimize offload weight pinning with mmap host registration
Fatemanx Jul 1, 2026
586aa15
Make pinned weight unregister finalizer robust
Fatemanx Jul 1, 2026
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
24 changes: 10 additions & 14 deletions lightx2v/common/ops/embedding/embedding_weight.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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")
Expand Down
65 changes: 22 additions & 43 deletions lightx2v/common/ops/mm/mm_weight.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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"):
Expand Down Expand Up @@ -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")
Expand Down
9 changes: 4 additions & 5 deletions lightx2v/common/ops/norm/layer_norm_weight.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
5 changes: 2 additions & 3 deletions lightx2v/common/ops/norm/rms_norm_weight.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
24 changes: 10 additions & 14 deletions lightx2v/common/ops/tensor/tensor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"):
Expand All @@ -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)

Expand Down Expand Up @@ -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)
Loading
Loading