diff --git a/python/pyproject.toml b/python/pyproject.toml index b5919cd0a..751584a49 100755 --- a/python/pyproject.toml +++ b/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] diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index 38c73c902..ae40df3f1 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py index 2bda6ee6e..6577319a0 100644 --- a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py +++ b/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( - 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 + if use_runai_model_streamer: + with SafetensorsStreamer() as streamer: + streamer.stream_files(hf_weights_files) + for name, tensor in streamer.get_tensors(): + if to_cpu: + yield name, tensor.clone().detach() + else: + yield name, tensor.to(device) + else: + 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(