[diffusion] improve: further optimize model load (#13836)

python/pyproject.toml

@@ -94,6 +94,7 @@ diffusion = [
   "st_attn ==0.0.7",
   "vsa==0.0.4",
   "yunchang==0.6.3.post1",
+  "runai_model_streamer",
 ]
 
 [tool.uv.extra-build-dependencies]

python/sglang/multimodal_gen/runtime/loader/fsdp_load.py

@@ -269,7 +269,7 @@ def load_model_from_full_model_state_dict(
                 meta_sharded_param.placements,
             )
             if cpu_offload:
-                sharded_tensor = sharded_tensor.cpu()
+                sharded_tensor = sharded_tensor.to("cpu", pin_memory=True)
         sharded_sd[target_param_name] = nn.Parameter(sharded_tensor)
 
     model.reverse_param_names_mapping = reverse_param_names_mapping

python/sglang/multimodal_gen/runtime/loader/weight_utils.py

@@ -16,6 +16,13 @@ import torch
 from safetensors.torch import safe_open
 from tqdm.auto import tqdm
 
+try:
+    from runai_model_streamer import SafetensorsStreamer
+
+    HAS_RUNAI_MODEL_STREAMER = True
+except ImportError:
+    HAS_RUNAI_MODEL_STREAMER = False
+
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
@@ -142,6 +149,7 @@ def _validate_safetensors_file(file_path: str) -> bool:
 def safetensors_weights_iterator(
     hf_weights_files: list[str],
     to_cpu: bool = True,
+    use_runai_model_streamer: bool = HAS_RUNAI_MODEL_STREAMER,
 ) -> Generator[tuple[str, torch.Tensor], None, None]:
     """Iterate over the weights in the model safetensor files."""
     enable_tqdm = (
@@ -185,16 +193,25 @@ def safetensors_weights_iterator(
             "Please retry - the files will be re-downloaded automatically."
         )
 
-    for st_file in tqdm(
+    if use_runai_model_streamer:
-        hf_weights_files,
+        with SafetensorsStreamer() as streamer:
-        desc="Loading safetensors checkpoint shards",
+            streamer.stream_files(hf_weights_files)
-        disable=not enable_tqdm,
+            for name, tensor in streamer.get_tensors():
-        bar_format=_BAR_FORMAT,
+                if to_cpu:
-    ):
+                    yield name, tensor.clone().detach()
-        with safe_open(st_file, framework="pt", device=device) as f:
+                else:
-            for name in f.keys():  # noqa: SIM118
+                    yield name, tensor.to(device)
-                param = f.get_tensor(name)
+    else:
-                yield name, param
+        for st_file in tqdm(
+            hf_weights_files,
+            desc="Loading safetensors checkpoint shards",
+            disable=not enable_tqdm,
+            bar_format=_BAR_FORMAT,
+        ):
+            with safe_open(st_file, framework="pt", device=device) as f:
+                for name in f.keys():  # noqa: SIM118
+                    param = f.get_tensor(name)
+                    yield name, param
 
 
 def pt_weights_iterator(