minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] feat: avoid direct GPU parameter copies (#36832)

Browse Source
This commit is contained in:
Mick
2026-08-29 22:43:25 +08:00
committed by GitHub
parent cdbfe90b4a
commit b24bd44556
4 changed files with 91 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py
+3 -2
View File
@@ -427,8 +427,9 @@ class TransformerLoader(ComponentLoader):
)
if direct_gpu_weight_loading:
logger.warning(
"Direct GPU weight loading is enabled for %s; the complete checkpoint "
"state dict and materialized model weights may coexist on GPU during startup",
"Direct GPU weight loading is enabled for %s; compatible checkpoint "
"tensors become model storage, while transformed tensors may still "
"require temporary GPU allocations",
component_name,
)
python/sglang/multimodal_gen/runtime/loader/fsdp_load.py
+16 -8
View File
@@ -98,14 +98,12 @@ def _make_param_like(
return new_param
def _can_assign_cpu_tensor_without_copy(
def _can_assign_tensor_without_copy(
actual_param: torch.nn.Parameter,
full_tensor: torch.Tensor,
target_param: torch.Tensor,
) -> bool:
"""Return whether a TP=1 linear loader would only copy this CPU tensor."""
if full_tensor.device.type != "cpu":
return False
"""Return whether a TP=1 linear loader would only copy this tensor."""
weight_loader = actual_param.__dict__.get("weight_loader")
if not isinstance(weight_loader, MethodType):
return False
@@ -134,6 +132,8 @@ def _can_assign_cpu_tensor_without_copy(
return (
full_tensor.shape == target_param.shape
and full_tensor.dtype == target_param.dtype
and full_tensor.layout == target_param.layout
and full_tensor.stride() == target_param.stride()
)
@@ -432,6 +432,9 @@ def maybe_load_fsdp_model(
cpu_offload=load_on_cpu,
param_names_mapping=param_names_mapping_fn,
keep_checkpoint_mapping=keep_checkpoint_mapping,
allow_device_tensor_assignment=(
weight_load_plan.load_full_state_dict_on_device
),
preconverted_state_dict=preconverted_state_dict,
)
if bnb_quant_states:
@@ -568,6 +571,7 @@ def load_model_from_full_model_state_dict(
]
| None
) = None,
allow_device_tensor_assignment: bool = False,
) -> _IncompatibleKeys:
"""
Converting full state dict into a sharded state dict
@@ -581,6 +585,10 @@ def load_model_from_full_model_state_dict(
cpu_offload (bool): flag to check if FSDP offload is enabled
param_names_mapping (Optional[Callable[[str], str]]): a function that maps full param name to sharded param name
keep_checkpoint_mapping (bool): retain compatible CPU checkpoint tensors instead of copying them
allow_device_tensor_assignment (bool): adopt compatible checkpoint tensors
already materialized on the target device. This is reserved for an
explicit full-state direct-device load; ordinary loading keeps its
established parameter materialization path.
Returns:
``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
* **missing_keys** is a list of str containing the missing keys
@@ -734,10 +742,10 @@ def load_model_from_full_model_state_dict(
sharded_tensor = full_tensor
elif weight_loader is not None:
assert actual_param is not None
if _can_assign_cpu_tensor_without_copy(
actual_param,
full_tensor,
meta_sharded_param,
if (
full_tensor.device.type == "cpu" or allow_device_tensor_assignment
) and _can_assign_tensor_without_copy(
actual_param, full_tensor, meta_sharded_param
):
sharded_tensor = full_tensor
else:
python/sglang/multimodal_gen/test/unit/test_fsdp_load.py
+71 -15
View File
@@ -9,7 +9,11 @@ import torch
from safetensors.torch import safe_open, save_file
from torch import nn
from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
from sglang.multimodal_gen.runtime.layers.linear import (
ColumnParallelLinear,
ReplicatedLinear,
UnquantizedLinearMethod,
)
from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import (
BitsAndBytesConfig,
)
@@ -180,6 +184,26 @@ class TestFSDPEntrypointRegistration(unittest.TestCase):
class TestOrdinaryWeightLoading(unittest.TestCase):
def _load_replicated_weight(
self, device: torch.device, *, allow_device_tensor_assignment: bool = False
) -> tuple[torch.nn.Parameter, torch.Tensor]:
with torch.device("meta"):
model = _ReplicatedLinearModel()
checkpoint_weight = torch.arange(
16, dtype=torch.float32, device=device
).reshape(4, 4)
fsdp_load.load_model_from_full_model_state_dict(
model,
iter((("proj.weight", checkpoint_weight),)),
checkpoint_load_device=device,
param_dtype=torch.float32,
strict=True,
param_names_mapping=fsdp_load.get_param_names_mapping({}),
allow_device_tensor_assignment=allow_device_tensor_assignment,
)
return model.proj.weight, checkpoint_weight
def test_direct_device_loading_skips_rank_local_cpu_checkpoint(self):
load_plan = WeightLoadPlan(
checkpoint_load_device=torch.device("cuda:0"),
@@ -217,22 +241,54 @@ class TestOrdinaryWeightLoading(unittest.TestCase):
weight_load_plan=load_plan,
)
def test_tp1_unquantized_linear_assigns_checkpoint_tensor_without_copy(self):
with torch.device("meta"):
model = _ReplicatedLinearModel()
checkpoint_weight = torch.arange(16, dtype=torch.float32).reshape(4, 4)
fsdp_load.load_model_from_full_model_state_dict(
model,
iter((("proj.weight", checkpoint_weight),)),
checkpoint_load_device=torch.device("cpu"),
param_dtype=torch.float32,
strict=True,
param_names_mapping=fsdp_load.get_param_names_mapping({}),
def test_tp1_unquantized_linear_adopts_cpu_checkpoint_storage(self):
model_weight, checkpoint_weight = self._load_replicated_weight(
torch.device("cpu")
)
self.assertEqual(model.proj.weight.data_ptr(), checkpoint_weight.data_ptr())
torch.testing.assert_close(model.proj.weight, checkpoint_weight)
self.assertEqual(model_weight.data_ptr(), checkpoint_weight.data_ptr())
torch.testing.assert_close(model_weight, checkpoint_weight)
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required")
def test_ordinary_cuda_loading_preserves_materialization_path(self):
model_weight, checkpoint_weight = self._load_replicated_weight(
torch.device("cuda:0")
)
self.assertNotEqual(model_weight.data_ptr(), checkpoint_weight.data_ptr())
torch.testing.assert_close(model_weight, checkpoint_weight)
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required")
def test_direct_cuda_loading_adopts_checkpoint_storage(self):
model_weight, checkpoint_weight = self._load_replicated_weight(
torch.device("cuda:0"), allow_device_tensor_assignment=True
)
self.assertEqual(model_weight.data_ptr(), checkpoint_weight.data_ptr())
torch.testing.assert_close(model_weight, checkpoint_weight)
def test_zero_copy_assignment_rejects_incompatible_layout_or_tp_weights(self):
with torch.device("meta"):
model = _ReplicatedLinearModel()
param = model.proj.weight
tensor = torch.empty(4, 4)
self.assertFalse(
fsdp_load._can_assign_tensor_without_copy(
param, tensor.as_strided((4, 4), (1, 4)), param
)
)
tp_owner = ColumnParallelLinear.__new__(ColumnParallelLinear)
nn.Module.__init__(tp_owner)
tp_owner.quant_method = UnquantizedLinearMethod()
tp_owner.tp_size = 1
tp_param = nn.Parameter(tensor)
tp_param.weight_loader = tp_owner.weight_loader
tp_owner.tp_size = 2
self.assertFalse(
fsdp_load._can_assign_tensor_without_copy(tp_param, tensor, tp_param)
)
class TestDevicePostprocessMove(unittest.TestCase):