diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index bd0e175cb..f31f69039 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/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,20 +277,55 @@ 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) - if preprocess_loaded_state_dict is not None: - weight_iterator = preprocess_loaded_state_dict(weight_iterator) bnb_quant_states = None - if _is_bitsandbytes_quant_config(init_params.get("quant_config")): - normal_weights, raw_quant_state = split_bitsandbytes_4bit_state(weight_iterator) - bnb_quant_states = build_bitsandbytes_4bit_quant_states( - [name for name, _ in normal_weights], - raw_quant_state, - device, - param_names_mapping_fn, + 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, + ) ) - weight_iterator = iter(normal_weights) + 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) + if is_bnb_quantized: + normal_weights, raw_quant_state = split_bitsandbytes_4bit_state( + weight_iterator + ) + bnb_quant_states = build_bitsandbytes_4bit_quant_states( + [name for name, _ in normal_weights], + raw_quant_state, + device, + 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 - 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 + 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 - - full_tensor = _maybe_dequantize_fp8( - full_tensor, target_dtype, target_param_name, custom_param_sd + 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, # type: ignore[arg-type] + ) if full_tensor.dtype != target_dtype: mismatch_key = (full_tensor.dtype, target_dtype) if ( - full_tensor.dtype in _QUANTIZED_DTYPES - or target_dtype in _QUANTIZED_DTYPES + full_tensor.dtype in rank_local_checkpoint.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 = ( diff --git a/python/sglang/multimodal_gen/runtime/loader/rank_local_checkpoint.py b/python/sglang/multimodal_gen/runtime/loader/rank_local_checkpoint.py new file mode 100644 index 000000000..63b93343c --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py index 4d48eb1d4..d3f2ef47e 100644 --- a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py +++ b/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: - """ - Validate that a safetensors file is readable and not corrupted. - - Args: - 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 - +def _scan_safetensors_files( + hf_weights_files: list[str], +) -> tuple[list[str], dict[str, set[str]]]: + """Validate headers and detect cross-file duplicate keys in one pass.""" + corrupted_files: list[str] = [] key_to_file: dict[str, str] = {} duplicate_files_by_key: dict[str, set[str]] = defaultdict(set) for st_file in hf_weights_files: - 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 - if previous_file == st_file: - continue - duplicate_files_by_key[name].update((previous_file, st_file)) + 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 + elif previous_file != st_file: + 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 = [ - st_file - for st_file in hf_weights_files - if not _validate_safetensors_file(st_file) - ] + corrupted_files, duplicate_files_by_key = _scan_safetensors_files(hf_weights_files) 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( diff --git a/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py b/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py index bf60422e6..5feb49723 100644 --- a/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py +++ b/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) + )