diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index 5b3ed0ca8..4dff753e8 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -36,6 +36,16 @@ from sglang.multimodal_gen.utils import set_mixed_precision_policy logger = init_logger(__name__) +def _make_param_like( + actual_param: torch.nn.Parameter, tensor: torch.Tensor +) -> torch.nn.Parameter: + cls = actual_param.__class__ + new_param = cls.__new__(cls, tensor) + new_param.__dict__.update(actual_param.__dict__) + new_param.requires_grad = False + return new_param + + # TODO(PY): move this to utils elsewhere @contextlib.contextmanager def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]: @@ -270,13 +280,11 @@ def load_model_from_full_model_state_dict( else None ) if weight_loader is not None: + assert actual_param is not None sharded_tensor = torch.empty_like( meta_sharded_param, device=device, dtype=param_dtype ) - temp_param = nn.Parameter(sharded_tensor) - for attr in ["output_dim", "input_dim", "is_sharded_weight"]: - if hasattr(actual_param, attr): - setattr(temp_param, attr, getattr(actual_param, attr)) + temp_param = _make_param_like(actual_param, sharded_tensor) weight_loader(temp_param, full_tensor) sharded_tensor = temp_param.data else: