[diffusion] UX: speed up tp and fsdp checkpoint loading (#33960)

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

@@ -12,6 +12,7 @@ from itertools import chain
 from typing import Any
 
 import torch
+import torch.distributed.tensor as dist_tensor
 from torch import nn
 from torch.distributed import DeviceMesh, init_device_mesh
 from torch.distributed._tensor import distribute_tensor
@@ -31,6 +32,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import (
     build_bitsandbytes_4bit_quant_states,
     split_bitsandbytes_4bit_state,
 )
+from sglang.multimodal_gen.runtime.loader import rank_local_checkpoint
 from sglang.multimodal_gen.runtime.loader.utils import (
     get_param_names_mapping,
     hf_to_custom_state_dict,
@@ -49,12 +51,6 @@ _is_npu = is_npu()
 
 logger = init_logger(__name__)
 
-_QUANTIZED_DTYPES = (
-    torch.uint8,
-    torch.float8_e4m3fn,
-    torch.float8_e5m2,
-    torch.int8,
-)
 _DTYPE_MISMATCH_EXAMPLE_LIMIT = 3
 
 
@@ -97,23 +93,6 @@ def _make_param_like(
     return new_param
 
 
-def _get_param_for_weight_loading(
-    model: torch.nn.Module,
-    param_dict: dict[str, torch.nn.Parameter],
-    param_name: str,
-) -> torch.nn.Parameter | None:
-    actual_param = param_dict.get(param_name)
-    if actual_param is not None and getattr(actual_param, "weight_loader", None):
-        return actual_param
-
-    pre_fsdp_weight_loader_params = getattr(model, "_pre_fsdp_weight_loader_params", {})
-    pre_fsdp_param = pre_fsdp_weight_loader_params.get(param_name)
-    if pre_fsdp_param is not None:
-        return pre_fsdp_param
-
-    return actual_param
-
-
 def _make_class_name_shard_condition(class_names: set[str]):
     def shard_condition(n: str, m: nn.Module) -> bool:
         return type(m).__name__ in class_names
@@ -298,13 +277,45 @@ def maybe_load_fsdp_model(
     param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping)
 
     # 2. load model from disk
-    weight_iterator = safetensors_weights_iterator(weight_dir_list)
     preprocess_loaded_state_dict = getattr(model, "preprocess_loaded_state_dict", None)
+    bnb_quant_states = None
+    preconverted_state_dict = None
+    is_bnb_quantized = _is_bitsandbytes_quant_config(init_params.get("quant_config"))
+    if (
+        use_fsdp
+        and weight_dir_list
+        and preprocess_loaded_state_dict is None
+        and not is_bnb_quantized
+    ):
+        preconverted_state_dict = (
+            rank_local_checkpoint.try_load_rank_local_fsdp_state_dict(
+                model,
+                weight_dir_list,
+                param_names_mapping_fn,
+            )
+        )
+    elif (
+        not use_fsdp
+        and weight_dir_list
+        and preprocess_loaded_state_dict is None
+        and not is_bnb_quantized
+    ):
+        preconverted_state_dict = (
+            rank_local_checkpoint.try_load_rank_local_tp_state_dict(
+                model,
+                weight_dir_list,
+                param_names_mapping_fn,
+            )
+        )
+
+    if preconverted_state_dict is None:
+        weight_iterator = safetensors_weights_iterator(weight_dir_list)
         if preprocess_loaded_state_dict is not None:
             weight_iterator = preprocess_loaded_state_dict(weight_iterator)
-    bnb_quant_states = None
+        if is_bnb_quantized:
-    if _is_bitsandbytes_quant_config(init_params.get("quant_config")):
+            normal_weights, raw_quant_state = split_bitsandbytes_4bit_state(
-        normal_weights, raw_quant_state = split_bitsandbytes_4bit_state(weight_iterator)
+                weight_iterator
+            )
             bnb_quant_states = build_bitsandbytes_4bit_quant_states(
                 [name for name, _ in normal_weights],
                 raw_quant_state,
@@ -312,6 +323,9 @@ def maybe_load_fsdp_model(
                 param_names_mapping_fn,
             )
             weight_iterator = iter(normal_weights)
+    else:
+        weight_iterator = iter(())
+
     load_model_from_full_model_state_dict(
         model,
         weight_iterator,
@@ -320,6 +334,7 @@ def maybe_load_fsdp_model(
         strict=strict,
         cpu_offload=load_cpu_offload,
         param_names_mapping=param_names_mapping_fn,
+        preconverted_state_dict=preconverted_state_dict,
     )
     if bnb_quant_states:
         attach_bitsandbytes_4bit_quant_states(
@@ -442,6 +457,18 @@ def load_model_from_full_model_state_dict(
     strict: bool = False,
     cpu_offload: bool = False,
     param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None,
+    preconverted_state_dict: (
+        tuple[
+            dict[
+                str,
+                torch.Tensor
+                | rank_local_checkpoint.LocalFSDPShard
+                | rank_local_checkpoint.LocalTPShard,
+            ],
+            dict[str, tuple[str, Any, Any]],
+        ]
+        | None
+    ) = None,
 ) -> _IncompatibleKeys:
     """
     Converting full state dict into a sharded state dict
@@ -464,14 +491,17 @@ def load_model_from_full_model_state_dict(
     param_dict = dict(model.named_parameters())
 
     # map names from checkpoint to customized names
+    if preconverted_state_dict is None:
         custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(
             full_sd_iterator,
             param_names_mapping,
             valid_target_names=set(meta_sd.keys()),
         )  # type: ignore
+    else:
+        custom_param_sd, reverse_param_names_mapping = preconverted_state_dict
 
     is_fsdp_model = isinstance(model, FSDPModule) or any(
-        hasattr(p, "device_mesh") for p in meta_sd.values()
+        isinstance(param, dist_tensor.DTensor) for param in meta_sd.values()
     )
 
     # sort parameter names to ensure all ranks process parameters in the same order
@@ -494,7 +524,7 @@ def load_model_from_full_model_state_dict(
 
     # shard from loaded state_dict, custom_param_sd -> sharded_sd
     for target_param_name in sorted_param_names:
-        full_tensor = custom_param_sd[target_param_name]
+        loaded_tensor = custom_param_sd[target_param_name]
         meta_sharded_param = meta_sd.get(target_param_name)
 
         if meta_sharded_param is None:
@@ -507,23 +537,29 @@ def load_model_from_full_model_state_dict(
                 skipped_checkpoint_keys.append(target_param_name)
                 continue
 
-        # use meta param dtype so quantized params (e.g. FP8) keep their dtype;
-        # for non-quantized models meta dtype equals param_dtype anyway
-        if meta_sharded_param is None:
-            # for nunchaku, some scales are patched later
-            target_dtype = full_tensor.dtype
-        else:
         target_dtype = meta_sharded_param.dtype
+        is_rank_local_fsdp_shard = isinstance(
+            loaded_tensor, rank_local_checkpoint.LocalFSDPShard
+        )
+        is_rank_local_tp_shard = isinstance(
+            loaded_tensor, rank_local_checkpoint.LocalTPShard
+        )
+        is_rank_local_shard = is_rank_local_fsdp_shard or is_rank_local_tp_shard
+        full_tensor = loaded_tensor.tensor if is_rank_local_shard else loaded_tensor
 
+        if not is_rank_local_shard:
             full_tensor = _maybe_dequantize_fp8(
-            full_tensor, target_dtype, target_param_name, custom_param_sd
+                full_tensor,
+                target_dtype,
+                target_param_name,
+                custom_param_sd,  # type: ignore[arg-type]
             )
 
         if full_tensor.dtype != target_dtype:
             mismatch_key = (full_tensor.dtype, target_dtype)
             if (
-                full_tensor.dtype in _QUANTIZED_DTYPES
+                full_tensor.dtype in rank_local_checkpoint.QUANTIZED_DTYPES
-                or target_dtype in _QUANTIZED_DTYPES
+                or target_dtype in rank_local_checkpoint.QUANTIZED_DTYPES
             ):
                 quantized_dtype_mismatch_counts[mismatch_key] += 1
                 if (
@@ -543,11 +579,41 @@ def load_model_from_full_model_state_dict(
                         target_param_name
                     )
 
-        if not hasattr(meta_sharded_param, "device_mesh"):
+        if is_rank_local_fsdp_shard:
+            if not isinstance(meta_sharded_param, dist_tensor.DTensor):
+                raise TypeError(
+                    f"Rank-local FSDP shard produced for non-DTensor parameter {target_param_name}"
+                )
+            local_tensor = full_tensor.to(
+                device=checkpoint_load_device,
+                dtype=target_dtype,
+            )
+            sharded_tensor = dist_tensor.DTensor.from_local(
+                local_tensor,
+                meta_sharded_param.device_mesh,
+                meta_sharded_param.placements,
+                run_check=False,
+                shape=meta_sharded_param.shape,
+                stride=meta_sharded_param.stride(),
+            )
+            if cpu_offload:
+                sharded_tensor = sharded_tensor.to("cpu")
+        elif is_rank_local_tp_shard:
+            if isinstance(meta_sharded_param, dist_tensor.DTensor):
+                raise TypeError(
+                    f"Rank-local TP shard produced for DTensor parameter {target_param_name}"
+                )
+            sharded_tensor = full_tensor.to(
+                device=checkpoint_load_device,
+                dtype=target_dtype,
+            )
+            if cpu_offload:
+                sharded_tensor = sharded_tensor.cpu()
+        elif not isinstance(meta_sharded_param, dist_tensor.DTensor):
             full_tensor = full_tensor.to(
                 device=checkpoint_load_device, dtype=target_dtype
             )
-            actual_param = _get_param_for_weight_loading(
+            actual_param = rank_local_checkpoint.get_param_for_weight_loading(
                 model, param_dict, target_param_name
             )
             weight_loader = (
@@ -600,7 +666,7 @@ def load_model_from_full_model_state_dict(
             full_tensor = full_tensor.to(
                 device=checkpoint_load_device, dtype=target_dtype
             )
-            actual_param = _get_param_for_weight_loading(
+            actual_param = rank_local_checkpoint.get_param_for_weight_loading(
                 model, param_dict, target_param_name
             )
             weight_loader = (

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

@@ -0,0 +1,476 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from collections import defaultdict
+from collections.abc import Callable
+from contextlib import ExitStack
+from dataclasses import dataclass
+from types import MethodType
+from typing import Any
+
+import torch
+import torch.distributed.tensor as dist_tensor
+from safetensors.torch import safe_open
+from torch.distributed.fsdp import FSDPModule
+
+from sglang.multimodal_gen.runtime.distributed import get_tp_rank, get_tp_world_size
+from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.loader.weight_utils import _scan_safetensors_files
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+QUANTIZED_DTYPES = (
+    torch.uint8,
+    torch.float8_e4m3fn,
+    torch.float8_e5m2,
+    torch.int8,
+)
+_QUANTIZED_SAFETENSORS_DTYPES = {
+    "F8_E4M3",
+    "F8_E5M2",
+    "I8",
+    "U8",
+}
+
+
+@dataclass(frozen=True)
+class SafetensorsSource:
+    file_path: str
+    param_name: str
+    shape: tuple[int, ...]
+    dtype: str
+    merge_index: int | None
+    num_params_to_merge: int | None
+
+
+@dataclass(frozen=True)
+class LocalFSDPShard:
+    tensor: torch.Tensor
+
+
+@dataclass(frozen=True)
+class LocalTPShard:
+    tensor: torch.Tensor
+
+
+def get_param_for_weight_loading(
+    model: torch.nn.Module,
+    param_dict: dict[str, torch.nn.Parameter],
+    param_name: str,
+) -> torch.nn.Parameter | None:
+    actual_param = param_dict.get(param_name)
+    if actual_param is not None and actual_param.__dict__.get("weight_loader"):
+        return actual_param
+
+    pre_fsdp_weight_loader_params = model.__dict__.get(
+        "_pre_fsdp_weight_loader_params", {}
+    )
+    pre_fsdp_param = pre_fsdp_weight_loader_params.get(param_name)
+    if pre_fsdp_param is not None:
+        return pre_fsdp_param
+
+    return actual_param
+
+
+def _mapped_param_name(
+    source_param_name: str,
+    param_names_mapping: Callable[[str], tuple[str, Any, Any]],
+    valid_target_names: set[str],
+) -> tuple[str, int | None, int | None]:
+    target_param_name, merge_index, num_params_to_merge = param_names_mapping(
+        source_param_name
+    )
+    if (
+        target_param_name != source_param_name
+        and source_param_name in valid_target_names
+        and target_param_name not in valid_target_names
+    ):
+        return source_param_name, None, None
+    return target_param_name, merge_index, num_params_to_merge
+
+
+def assembled_source_shape(
+    sources: list[SafetensorsSource],
+) -> tuple[int, ...] | None:
+    if len(sources) == 1 and sources[0].merge_index is None:
+        return sources[0].shape
+    if any(source.merge_index is None for source in sources):
+        return None
+
+    expected_count = sources[0].num_params_to_merge
+    if expected_count is None or len(sources) != expected_count:
+        return None
+    if any(source.num_params_to_merge != expected_count for source in sources):
+        return None
+    if {source.merge_index for source in sources} != set(range(expected_count)):
+        return None
+
+    reference_shape = sources[0].shape
+    if not reference_shape:
+        return None
+    if any(source.shape[1:] != reference_shape[1:] for source in sources):
+        return None
+    return (sum(source.shape[0] for source in sources), *reference_shape[1:])
+
+
+def _collect_safetensors_sources(
+    weight_files: list[str],
+    param_names_mapping: Callable[[str], tuple[str, Any, Any]],
+    valid_target_names: set[str],
+) -> (
+    tuple[
+        dict[str, list[SafetensorsSource]],
+        dict[str, tuple[str, Any, Any]],
+    ]
+    | None
+):
+    corrupted_files, duplicate_files_by_key = _scan_safetensors_files(weight_files)
+    if corrupted_files or duplicate_files_by_key:
+        return None
+
+    sources_by_target: dict[str, list[SafetensorsSource]] = defaultdict(list)
+    reverse_param_names_mapping: dict[str, tuple[str, Any, Any]] = {}
+    for file_path in weight_files:
+        with safe_open(file_path, framework="pt", device="cpu") as handle:
+            for source_param_name in handle.keys():  # noqa: SIM118
+                target_param_name, merge_index, num_params_to_merge = (
+                    _mapped_param_name(
+                        source_param_name,
+                        param_names_mapping,
+                        valid_target_names,
+                    )
+                )
+                if not target_param_name:
+                    continue
+                source_slice = handle.get_slice(source_param_name)
+                sources_by_target[target_param_name].append(
+                    SafetensorsSource(
+                        file_path=file_path,
+                        param_name=source_param_name,
+                        shape=tuple(source_slice.get_shape()),
+                        dtype=source_slice.get_dtype(),
+                        merge_index=merge_index,
+                        num_params_to_merge=num_params_to_merge,
+                    )
+                )
+                reverse_param_names_mapping[target_param_name] = (
+                    source_param_name,
+                    merge_index,
+                    num_params_to_merge,
+                )
+
+    return sources_by_target, reverse_param_names_mapping
+
+
+def read_rank_local_tensor(
+    sources: list[SafetensorsSource],
+    handles: dict[str, Any],
+    local_shape: tuple[int, ...],
+    global_offset: tuple[int, ...],
+) -> torch.Tensor:
+    if not local_shape:
+        if len(sources) != 1:
+            raise RuntimeError("Scalar checkpoint parameters cannot be merged")
+        source = sources[0]
+        return handles[source.file_path].get_tensor(source.param_name)
+
+    ordered_sources = sorted(
+        sources,
+        key=lambda source: (
+            source.merge_index is None,
+            source.merge_index if source.merge_index is not None else 0,
+        ),
+    )
+    if any(size == 0 for size in local_shape):
+        source = ordered_sources[0]
+        empty_slices = (slice(0, 0),) + (slice(None),) * (len(source.shape) - 1)
+        return (
+            handles[source.file_path]
+            .get_slice(source.param_name)[empty_slices]
+            .contiguous()
+            .reshape(local_shape)
+        )
+
+    local_start = global_offset[0]
+    local_end = local_start + local_shape[0]
+    source_start = 0
+    local_parts: list[torch.Tensor] = []
+    for source in ordered_sources:
+        source_end = source_start + source.shape[0]
+        intersection_start = max(local_start, source_start)
+        intersection_end = min(local_end, source_end)
+        if intersection_start < intersection_end:
+            slices = [
+                slice(
+                    intersection_start - source_start,
+                    intersection_end - source_start,
+                )
+            ]
+            slices.extend(
+                slice(offset, offset + size)
+                for offset, size in zip(global_offset[1:], local_shape[1:])
+            )
+            local_parts.append(
+                handles[source.file_path]
+                .get_slice(source.param_name)[tuple(slices)]
+                .contiguous()
+            )
+        source_start = source_end
+
+    if not local_parts:
+        raise RuntimeError(
+            f"No checkpoint slice overlaps local FSDP shard at offset={global_offset}, shape={local_shape}"
+        )
+    local_tensor = (
+        local_parts[0] if len(local_parts) == 1 else torch.cat(local_parts, dim=0)
+    )
+    if tuple(local_tensor.shape) != local_shape:
+        raise RuntimeError(
+            "Rank-local checkpoint slice shape mismatch: "
+            f"loaded={tuple(local_tensor.shape)}, expected={local_shape}"
+        )
+    return local_tensor
+
+
+def _resolve_tp_shard_dim(
+    actual_param: torch.nn.Parameter,
+) -> tuple[bool, int | None]:
+    weight_loader = actual_param.__dict__.get("weight_loader")
+    if weight_loader is None:
+        return True, None
+    if not isinstance(weight_loader, MethodType):
+        return False, None
+
+    owner = weight_loader.__self__
+    if isinstance(owner, ReplicatedLinear):
+        return True, None
+    if isinstance(owner, ColumnParallelLinear):
+        output_dim = actual_param.__dict__.get("output_dim")
+        return output_dim is not None, output_dim
+    if isinstance(owner, RowParallelLinear):
+        input_dim = actual_param.__dict__.get("input_dim")
+        return True, input_dim
+    return False, None
+
+
+def tp_local_shape(
+    sources: list[SafetensorsSource],
+    shard_dim: int | None,
+    tp_size: int,
+) -> tuple[int, ...] | None:
+    assembled_shape = assembled_source_shape(sources)
+    if assembled_shape is None or shard_dim is None:
+        return assembled_shape
+    if shard_dim >= len(assembled_shape):
+        return None
+    if any(source.shape[shard_dim] % tp_size != 0 for source in sources):
+        return None
+
+    local_shape = list(assembled_shape)
+    local_shape[shard_dim] //= tp_size
+    return tuple(local_shape)
+
+
+def read_tp_local_tensor(
+    sources: list[SafetensorsSource],
+    handles: dict[str, Any],
+    shard_dim: int | None,
+    tp_rank: int,
+    tp_size: int,
+) -> torch.Tensor:
+    if shard_dim is None:
+        assembled_shape = assembled_source_shape(sources)
+        if assembled_shape is None:
+            raise RuntimeError("Invalid checkpoint sources for replicated TP parameter")
+        return read_rank_local_tensor(
+            sources,
+            handles,
+            assembled_shape,
+            (0,) * len(assembled_shape),
+        )
+
+    ordered_sources = sorted(
+        sources,
+        key=lambda source: (
+            source.merge_index is None,
+            source.merge_index if source.merge_index is not None else 0,
+        ),
+    )
+    local_parts = []
+    for source in ordered_sources:
+        shard_size = source.shape[shard_dim] // tp_size
+        slices = [slice(None)] * len(source.shape)
+        slices[shard_dim] = slice(
+            tp_rank * shard_size,
+            (tp_rank + 1) * shard_size,
+        )
+        local_parts.append(
+            handles[source.file_path]
+            .get_slice(source.param_name)[tuple(slices)]
+            .contiguous()
+        )
+    return local_parts[0] if len(local_parts) == 1 else torch.cat(local_parts, dim=0)
+
+
+def try_load_rank_local_tp_state_dict(
+    model: torch.nn.Module,
+    weight_files: list[str],
+    param_names_mapping: Callable[[str], tuple[str, Any, Any]],
+) -> (
+    tuple[
+        dict[str, LocalTPShard],
+        dict[str, tuple[str, Any, Any]],
+    ]
+    | None
+):
+    tp_size = get_tp_world_size()
+    if tp_size == 1:
+        return None
+
+    meta_sd = model.state_dict()
+    param_dict = dict(model.named_parameters())
+    checkpoint_sources = _collect_safetensors_sources(
+        weight_files,
+        param_names_mapping,
+        set(meta_sd),
+    )
+    if checkpoint_sources is None:
+        return None
+    sources_by_target, reverse_param_names_mapping = checkpoint_sources
+
+    shard_dims: dict[str, int | None] = {}
+    for target_param_name, sources in sources_by_target.items():
+        meta_param = meta_sd.get(target_param_name)
+        if meta_param is None or isinstance(meta_param, dist_tensor.DTensor):
+            return None
+        if meta_param.dtype in QUANTIZED_DTYPES:
+            return None
+        if any(source.dtype in _QUANTIZED_SAFETENSORS_DTYPES for source in sources):
+            return None
+
+        actual_param = get_param_for_weight_loading(
+            model,
+            param_dict,
+            target_param_name,
+        )
+        if actual_param is None:
+            supported, shard_dim = True, None
+        else:
+            supported, shard_dim = _resolve_tp_shard_dim(actual_param)
+        if not supported:
+            return None
+        if tp_local_shape(sources, shard_dim, tp_size) != tuple(meta_param.shape):
+            return None
+        shard_dims[target_param_name] = shard_dim
+
+    local_param_sd: dict[str, LocalTPShard] = {}
+    local_bytes = 0
+    tp_rank = get_tp_rank()
+    with ExitStack() as stack:
+        handles = {
+            file_path: stack.enter_context(
+                safe_open(file_path, framework="pt", device="cpu")
+            )
+            for file_path in weight_files
+        }
+        for target_param_name in sorted(sources_by_target):
+            tensor = read_tp_local_tensor(
+                sources_by_target[target_param_name],
+                handles,
+                shard_dims[target_param_name],
+                tp_rank,
+                tp_size,
+            )
+            local_param_sd[target_param_name] = LocalTPShard(tensor)
+            local_bytes += tensor.numel() * tensor.element_size()
+
+    logger.info(
+        "Loaded rank-local TP checkpoint slices: rank=%d, tensors=%d, bytes=%.2f GiB",
+        torch.distributed.get_rank(),
+        len(local_param_sd),
+        local_bytes / (1024**3),
+    )
+    return local_param_sd, reverse_param_names_mapping
+
+
+def try_load_rank_local_fsdp_state_dict(
+    model: FSDPModule,
+    weight_files: list[str],
+    param_names_mapping: Callable[[str], tuple[str, Any, Any]],
+) -> (
+    tuple[
+        dict[str, torch.Tensor | LocalFSDPShard],
+        dict[str, tuple[str, Any, Any]],
+    ]
+    | None
+):
+    if get_tp_world_size() != 1:
+        return None
+
+    meta_sd = model.state_dict()
+    checkpoint_sources = _collect_safetensors_sources(
+        weight_files,
+        param_names_mapping,
+        set(meta_sd),
+    )
+    if checkpoint_sources is None:
+        return None
+    sources_by_target, reverse_param_names_mapping = checkpoint_sources
+
+    for target_param_name, sources in sources_by_target.items():
+        meta_param = meta_sd.get(target_param_name)
+        assembled_shape = assembled_source_shape(sources)
+        if meta_param is None or assembled_shape != tuple(meta_param.shape):
+            return None
+        if meta_param.dtype in QUANTIZED_DTYPES:
+            return None
+        if any(source.dtype in _QUANTIZED_SAFETENSORS_DTYPES for source in sources):
+            return None
+
+    local_param_sd: dict[str, torch.Tensor | LocalFSDPShard] = {}
+    local_bytes = 0
+    with ExitStack() as stack:
+        handles = {
+            file_path: stack.enter_context(
+                safe_open(file_path, framework="pt", device="cpu")
+            )
+            for file_path in weight_files
+        }
+        for target_param_name in sorted(sources_by_target):
+            meta_param = meta_sd[target_param_name]
+            if isinstance(meta_param, dist_tensor.DTensor):
+                local_shape, global_offset = (
+                    dist_tensor._utils.compute_local_shape_and_global_offset(
+                        meta_param.shape,
+                        meta_param.device_mesh,
+                        meta_param.placements,
+                    )
+                )
+                tensor = read_rank_local_tensor(
+                    sources_by_target[target_param_name],
+                    handles,
+                    tuple(local_shape),
+                    tuple(global_offset),
+                )
+                local_param_sd[target_param_name] = LocalFSDPShard(tensor)
+            else:
+                tensor = read_rank_local_tensor(
+                    sources_by_target[target_param_name],
+                    handles,
+                    tuple(meta_param.shape),
+                    (0,) * meta_param.ndim,
+                )
+                local_param_sd[target_param_name] = tensor
+            local_bytes += tensor.numel() * tensor.element_size()
+
+    logger.info(
+        "Loaded rank-local FSDP checkpoint slices: rank=%d, tensors=%d, bytes=%.2f GiB",
+        torch.distributed.get_rank(),
+        len(local_param_sd),
+        local_bytes / (1024**3),
+    )
+    return local_param_sd, reverse_param_names_mapping

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

@@ -117,55 +117,38 @@ def filter_files_not_needed_for_inference(hf_weights_files: list[str]) -> list[s
 _BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]\n"  # noqa: E501
 
 
-def _validate_safetensors_file(file_path: str) -> bool:
+def _scan_safetensors_files(
-    """
+    hf_weights_files: list[str],
-    Validate that a safetensors file is readable and not corrupted.
+) -> tuple[list[str], dict[str, set[str]]]:
-
+    """Validate headers and detect cross-file duplicate keys in one pass."""
-    Args:
+    corrupted_files: list[str] = []
-        file_path: Path to the safetensors file
-
-    Returns:
-        True if file is valid, False if corrupted
-    """
-    try:
-        with safe_open(file_path, framework="pt", device="cpu") as f:
-            _ = list(f.keys())
-        return True
-    except Exception as e:
-        logger.error(
-            "Corrupted safetensors file detected: %s - %s: %s",
-            file_path,
-            type(e).__name__,
-            str(e),
-        )
-        return False
-
-
-def _raise_if_duplicate_safetensors_keys(hf_weights_files: list[str]) -> None:
-    """Fail fast when multiple safetensors files define the same tensor name. Make sure runtime behavior is deterministic
-
-    Duplicate keys across files are almost always a packaging error for inference:
-    for example shipping both full and fp16 variants, or mixing consolidated and
-    sharded checkpoints. Continuing would make the final loaded value depend on
-    file iteration or streamer delivery order.
-    """
-    if len(hf_weights_files) <= 1:
-        return
-
     key_to_file: dict[str, str] = {}
     duplicate_files_by_key: dict[str, set[str]] = defaultdict(set)
 
     for st_file in hf_weights_files:
+        try:
             with safe_open(st_file, framework="pt", device="cpu") as f:
                 for name in f.keys():  # noqa: SIM118
                     previous_file = key_to_file.get(name)
                     if previous_file is None:
                         key_to_file[name] = st_file
-                    continue
+                    elif previous_file != st_file:
-                if previous_file == st_file:
-                    continue
                         duplicate_files_by_key[name].update((previous_file, st_file))
+        except Exception as e:
+            logger.error(
+                "Corrupted safetensors file detected: %s - %s: %s",
+                st_file,
+                type(e).__name__,
+                str(e),
+            )
+            corrupted_files.append(st_file)
 
+    return corrupted_files, duplicate_files_by_key
+
+
+def _raise_if_duplicate_safetensors_keys(
+    duplicate_files_by_key: dict[str, set[str]],
+) -> None:
     if not duplicate_files_by_key:
         return
 
@@ -210,11 +193,7 @@ def safetensors_weights_iterator(
         )
 
     # Validate files before loading
-    corrupted_files = [
+    corrupted_files, duplicate_files_by_key = _scan_safetensors_files(hf_weights_files)
-        st_file
-        for st_file in hf_weights_files
-        if not _validate_safetensors_file(st_file)
-    ]
 
     if corrupted_files:
         # Delete corrupted files (both symlink and blob if applicable)
@@ -245,7 +224,7 @@ def safetensors_weights_iterator(
             "Please retry - the files will be re-downloaded automatically."
         )
 
-    _raise_if_duplicate_safetensors_keys(hf_weights_files)
+    _raise_if_duplicate_safetensors_keys(duplicate_files_by_key)
 
     if use_runai_model_streamer:
         logger.info(

python/sglang/multimodal_gen/test/unit/test_fsdp_load.py

@@ -1,12 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 
+import tempfile
 import unittest
+from pathlib import Path
 from unittest.mock import patch
 
 import torch
+from safetensors.torch import safe_open, save_file
 from torch import nn
 
-from sglang.multimodal_gen.runtime.loader import fsdp_load
+from sglang.multimodal_gen.runtime.loader import fsdp_load, rank_local_checkpoint
 
 
 class _UniformDtypeModel(nn.Module):
@@ -92,3 +95,129 @@ class TestFSDPMixedPrecisionPolicy(unittest.TestCase):
         self.assertEqual(policy.param_dtype, torch.bfloat16)
         self.assertEqual(state_kwargs["param_dtype"], torch.bfloat16)
         shard_model.assert_not_called()
+
+
+class TestRankLocalSafetensorsRead(unittest.TestCase):
+    def _source(
+        self,
+        file_path: str,
+        param_name: str,
+        shape: tuple[int, ...],
+        merge_index: int | None = None,
+        num_params_to_merge: int | None = None,
+    ) -> rank_local_checkpoint.SafetensorsSource:
+        return rank_local_checkpoint.SafetensorsSource(
+            file_path=file_path,
+            param_name=param_name,
+            shape=shape,
+            dtype="BF16",
+            merge_index=merge_index,
+            num_params_to_merge=num_params_to_merge,
+        )
+
+    def test_reads_rank_local_slice(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            file_path = str(Path(temp_dir) / "model.safetensors")
+            weight = torch.arange(24, dtype=torch.bfloat16).reshape(6, 4)
+            save_file({"weight": weight}, file_path)
+
+            with safe_open(file_path, framework="pt", device="cpu") as handle:
+                tensor = rank_local_checkpoint.read_rank_local_tensor(
+                    [self._source(file_path, "weight", (6, 4))],
+                    {file_path: handle},
+                    local_shape=(2, 4),
+                    global_offset=(2, 0),
+                )
+
+            torch.testing.assert_close(tensor, weight[2:4])
+
+    def test_reads_rank_local_slice_across_merged_sources(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            file_path = str(Path(temp_dir) / "model.safetensors")
+            first = torch.arange(12, dtype=torch.bfloat16).reshape(3, 4)
+            second = torch.arange(12, 24, dtype=torch.bfloat16).reshape(3, 4)
+            save_file({"first": first, "second": second}, file_path)
+            sources = [
+                self._source(file_path, "first", (3, 4), 0, 2),
+                self._source(file_path, "second", (3, 4), 1, 2),
+            ]
+
+            with safe_open(file_path, framework="pt", device="cpu") as handle:
+                tensor = rank_local_checkpoint.read_rank_local_tensor(
+                    sources,
+                    {file_path: handle},
+                    local_shape=(4, 4),
+                    global_offset=(1, 0),
+                )
+
+            torch.testing.assert_close(tensor, torch.cat((first, second))[1:5])
+
+    def test_reads_zero_sized_rank_local_shard(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            file_path = str(Path(temp_dir) / "model.safetensors")
+            weight = torch.ones((1, 4), dtype=torch.bfloat16)
+            save_file({"weight": weight}, file_path)
+
+            with safe_open(file_path, framework="pt", device="cpu") as handle:
+                tensor = rank_local_checkpoint.read_rank_local_tensor(
+                    [self._source(file_path, "weight", (1, 4))],
+                    {file_path: handle},
+                    local_shape=(0, 4),
+                    global_offset=(1, 0),
+                )
+
+            self.assertEqual(tensor.shape, (0, 4))
+            self.assertEqual(tensor.dtype, torch.bfloat16)
+            self.assertEqual(tensor.numel(), 0)
+
+    def test_reads_tp_local_merged_column_slice_per_source(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            file_path = str(Path(temp_dir) / "model.safetensors")
+            first = torch.arange(16, dtype=torch.bfloat16).reshape(4, 4)
+            second = torch.arange(16, 32, dtype=torch.bfloat16).reshape(4, 4)
+            save_file({"first": first, "second": second}, file_path)
+            sources = [
+                self._source(file_path, "first", (4, 4), 0, 2),
+                self._source(file_path, "second", (4, 4), 1, 2),
+            ]
+
+            with safe_open(file_path, framework="pt", device="cpu") as handle:
+                tensor = rank_local_checkpoint.read_tp_local_tensor(
+                    sources,
+                    {file_path: handle},
+                    shard_dim=0,
+                    tp_rank=1,
+                    tp_size=2,
+                )
+
+            expected = torch.cat((first[2:4], second[2:4]))
+            torch.testing.assert_close(tensor, expected)
+            self.assertEqual(
+                rank_local_checkpoint.tp_local_shape(sources, 0, 2), (4, 4)
+            )
+
+    def test_reads_tp_local_merged_row_slice(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            file_path = str(Path(temp_dir) / "model.safetensors")
+            first = torch.arange(16, dtype=torch.bfloat16).reshape(4, 4)
+            second = torch.arange(16, 32, dtype=torch.bfloat16).reshape(4, 4)
+            save_file({"first": first, "second": second}, file_path)
+            sources = [
+                self._source(file_path, "first", (4, 4), 0, 2),
+                self._source(file_path, "second", (4, 4), 1, 2),
+            ]
+
+            with safe_open(file_path, framework="pt", device="cpu") as handle:
+                tensor = rank_local_checkpoint.read_tp_local_tensor(
+                    sources,
+                    {file_path: handle},
+                    shard_dim=1,
+                    tp_rank=1,
+                    tp_size=2,
+                )
+
+            expected = torch.cat((first[:, 2:4], second[:, 2:4]))
+            torch.testing.assert_close(tensor, expected)
+            self.assertEqual(
+                rank_local_checkpoint.tp_local_shape(sources, 1, 2), (8, 2)
+            )