[diffusion] fix: fix fsdp tp load make param miss parallel meta data (#17058)

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)
+                temp_param = _make_param_like(actual_param, 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))
                 weight_loader(temp_param, full_tensor)
                 sharded_tensor = temp_param.data
             else: