diff --git a/python/pyproject.toml b/python/pyproject.toml index 55ab27ee2..efc6cf46c 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -87,6 +87,7 @@ dependencies = [ "uvloop", "watchfiles", "xgrammar==0.2.1", + "xxhash", "zstandard", ] diff --git a/python/pyproject_cpu.toml b/python/pyproject_cpu.toml index 31ae12a42..4677c17c3 100644 --- a/python/pyproject_cpu.toml +++ b/python/pyproject_cpu.toml @@ -69,6 +69,7 @@ dependencies = [ "triton==3.7.0", "uvicorn", "uvloop", + "xxhash", "xgrammar==0.2.1", "zstandard", ] diff --git a/python/pyproject_npu.toml b/python/pyproject_npu.toml index 727c89de9..c2c078efe 100644 --- a/python/pyproject_npu.toml +++ b/python/pyproject_npu.toml @@ -65,6 +65,7 @@ dependencies = [ "transformers==5.12.1", "uvicorn", "uvloop", + "xxhash", "xgrammar==0.2.1", ] diff --git a/python/pyproject_other.toml b/python/pyproject_other.toml index 3264b161c..61aeb106e 100755 --- a/python/pyproject_other.toml +++ b/python/pyproject_other.toml @@ -66,6 +66,7 @@ runtime_common = [ "transformers==5.12.1", "uvicorn", "uvloop", + "xxhash", "xgrammar==0.2.1", ] diff --git a/python/pyproject_xpu.toml b/python/pyproject_xpu.toml index bcd3d1a72..9f14db014 100644 --- a/python/pyproject_xpu.toml +++ b/python/pyproject_xpu.toml @@ -67,6 +67,7 @@ dependencies = [ "tqdm", "transformers==5.12.1", "uvicorn", + "xxhash", "uvloop", # "xgrammar==0.2.1", xgrammar depends on CUDA PyTorch and Triton only ] diff --git a/python/sglang/srt/configs/load_config.py b/python/sglang/srt/configs/load_config.py index c70ca6bc8..9b63ab964 100644 --- a/python/sglang/srt/configs/load_config.py +++ b/python/sglang/srt/configs/load_config.py @@ -21,6 +21,7 @@ class LoadFormat(str, enum.Enum): NPCACHE = "npcache" DUMMY = "dummy" SHARDED_STATE = "sharded_state" + PRESHARDED = "presharded" GGUF = "gguf" BITSANDBYTES = "bitsandbytes" MISTRAL = "mistral" diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index b673cb670..8bafdd5cd 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -6,15 +6,18 @@ from __future__ import annotations # ruff: noqa: SIM117 import collections +import concurrent.futures import dataclasses import fnmatch import gc import glob +import hashlib import json import logging import math import os import re +import shutil import socket import tempfile import threading @@ -1497,6 +1500,15 @@ class ShardedStateLoader(BaseModelLoader): result: Dict[str, torch.Tensor] = {} for group in same_storage_groups.values(): for k, t in group: + if not t.is_contiguous(): + # End-pointer dedup assumes a flat view; non-contiguous + # tensors (e.g. produced by + # ``.transpose(...).contiguous().transpose(...)`` in some + # quant ``post_load_weights`` paths) cannot be flattened + # via ``view(-1)``. Include them directly; downstream + # writers call ``.contiguous()`` before save. + result[k] = t + continue a, b = t.data_ptr(), get_end_ptr(t) for k2, t2 in group: if not t2.is_contiguous(): @@ -1629,6 +1641,839 @@ class ShardedStateLoader(BaseModelLoader): ) +class PreshardedModelLoader(DefaultModelLoader): + """Dump/reload post-process weights under ``/presharded//``. + + Optional roots in ``model_loader_extra_config`` (subdir still appended): + ``presharded_path`` (target), ``draft_presharded_path`` (speculative draft). + Dump dir must be shared across ranks/nodes. + """ + + DEFAULT_SUBDIR = "presharded" + MAX_FILE_BYTES = 20 * (1024**3) + CHECKSUM_FILENAME = "checksum.json" + READY_FILENAME = "READY" + TMP_SUBDIR = "_tmp_presharding" + PLAN_VERSION = 1 + DEFAULT_HASH_NUM_THREADS = 8 + _CONTENT_HASH_HEX_LEN = 32 + + def __init__(self, load_config: LoadConfig): + extra = ( + {} + if load_config.model_loader_extra_config is None + else dict(load_config.model_loader_extra_config) + ) + self._presharded_path_override = extra.pop("presharded_path", None) + self._draft_presharded_path_override = extra.pop("draft_presharded_path", None) + self._max_file_bytes = int(extra.pop("max_file_bytes", self.MAX_FILE_BYTES)) + self._hash_num_threads = int( + extra.pop("hash_num_threads", self.DEFAULT_HASH_NUM_THREADS) + ) + self._verify_on_load = bool(extra.pop("verify_on_load", False)) + load_config.model_loader_extra_config = extra + load_config.load_format = LoadFormat.AUTO + super().__init__(load_config) + + def download_model(self, model_config: ModelConfig) -> None: + presharded_dir = self._presharded_dir(model_config) + if not self._presharded_ready(presharded_dir): + super().download_model(model_config) + + def load_model( + self, + *, + model_config: ModelConfig, + device_config: DeviceConfig, + ) -> nn.Module: + shard_config = self._collect_shard_config(model_config) + presharded_dir = self._presharded_dir(model_config, shard_config) + if self._presharded_ready(presharded_dir) and self._shard_config_matches( + presharded_dir, shard_config + ): + logger.info("Loading from presharded checkpoint at %s", presharded_dir) + return self._load_from_presharded( + model_config, device_config, presharded_dir + ) + logger.info( + "No presharded checkpoint at %s; doing first-time load and dump.", + presharded_dir, + ) + return self._first_time_load_and_dump( + model_config, device_config, presharded_dir, shard_config + ) + + @classmethod + def _presharded_ready(cls, presharded_dir: str) -> bool: + return os.path.isfile(os.path.join(presharded_dir, cls.READY_FILENAME)) + + def _presharded_dir( + self, + model_config: ModelConfig, + shard_config: Optional[Dict[str, Any]] = None, + ) -> str: + if shard_config is None: + shard_config = self._collect_shard_config(model_config) + subfolder = self._build_subfolder_name(shard_config) + if model_config.is_draft_model: + root = self._draft_presharded_path_override + else: + root = self._presharded_path_override + if root is None: + root = os.path.join(model_config.model_path, self.DEFAULT_SUBDIR) + return os.path.join(root, subfolder) + + def _collect_shard_config(self, model_config: ModelConfig) -> Dict[str, Any]: + def _safe(fn) -> int: + try: + return fn() + except (AssertionError, AttributeError, RuntimeError): + return 1 + + parallel = get_parallel() + server_args = get_server_args() + return { + "tp": _safe(lambda: parallel.tp_size), + "dp": _safe(lambda: parallel.moe_dp_size), + "ep": _safe(lambda: parallel.moe_ep_size), + "pp": _safe(lambda: parallel.pp_size), + "moe_dense_tp_size": server_args.moe_dense_tp_size, + "moe_dp_size": server_args.moe_dp_size, + "enable_dp_lm_head": server_args.enable_dp_lm_head, + "enable_fp32_lm_head": server_args.enable_fp32_lm_head, + "quantization": model_config.quantization, + "model_dtype": str(model_config.dtype), + "ep_num_redundant_experts": server_args.ep_num_redundant_experts, + "enable_eplb": server_args.enable_eplb, + "init_expert_location": self._normalize_init_expert_location( + server_args.init_expert_location + ), + "structural_signature": self._compute_structural_signature(model_config), + } + + @staticmethod + def _normalize_init_expert_location(value: Optional[str]) -> Optional[str]: + if value is None or value == "trivial": + return value + if value.endswith((".json", ".pt")) and os.path.isfile(value): + h = hashlib.sha1() + with open(value, "rb") as f: + for chunk in iter(lambda: f.read(1024 * 1024), b""): + h.update(chunk) + return f"file:{os.path.basename(value)}:sha1:{h.hexdigest()[:16]}" + return value + + def _build_subfolder_name(self, shard_config: Dict[str, Any]) -> str: + combined = hashlib.sha1( + json.dumps(shard_config, sort_keys=True).encode() + ).hexdigest()[:16] + return f"TP-{shard_config['tp']}-sig-{combined}" + + def _shard_config_matches( + self, presharded_dir: str, shard_config: Dict[str, Any] + ) -> bool: + try: + with open(os.path.join(presharded_dir, self.CHECKSUM_FILENAME)) as f: + stored = json.load(f).get("shard_config") + except (OSError, ValueError): + stored = None + current = json.loads(json.dumps(shard_config)) + if stored == current: + return True + logger.warning( + "Presharded checkpoint at %s was dumped with a different shard " + "config than the current launch (stored=%s, current=%s). " + "Treating as a cache miss and re-dumping.", + presharded_dir, + stored, + current, + ) + return False + + def _compute_structural_signature(self, model_config: ModelConfig) -> Optional[str]: + local_sig = self._compute_local_structural_signature(model_config) + return self._make_rank_invariant_structural_signature(local_sig) + + def _compute_local_structural_signature( + self, model_config: ModelConfig + ) -> Optional[str]: + from sglang.srt.layers.rotary_embedding.factory import _ROPE_DICT + + def _clear_meta_rope_cache() -> None: + meta_keys = [ + k + for k, v in _ROPE_DICT.items() + if any(p.device.type == "meta" for p in v.parameters()) + or any(b.device.type == "meta" for b in v.buffers()) + ] + for k in meta_keys: + del _ROPE_DICT[k] + + try: + quant_config = _get_quantization_config(model_config, self.load_config) + with set_default_torch_dtype(model_config.dtype): + with torch.device("meta"): + meta_model = _initialize_model( + model_config, self.load_config, quant_config + ) + state_dict = meta_model.state_dict() + sig_input = sorted( + (name, tuple(t.shape), str(t.dtype)) + for name, t in state_dict.items() + ) + del meta_model + return self._hash_structural_signature(sig_input) + except Exception as e: + logger.warning( + "Failed to build structural signature for presharded cache key " + "(model_type=%s): %s", + getattr( + getattr(model_config, "hf_config", None), "model_type", "unknown" + ), + e, + ) + return None + finally: + _clear_meta_rope_cache() + + @classmethod + def _make_rank_invariant_structural_signature( + cls, local_sig: Optional[str] + ) -> Optional[str]: + try: + from sglang.srt.distributed import get_world_group + + group = get_world_group() + if group.world_size <= 1: + return local_sig + all_sigs = group.all_gather_object(local_sig) + except (AssertionError, AttributeError, RuntimeError): + return local_sig + + if all(s is None for s in all_sigs): + return None + return hashlib.sha1(repr(all_sigs).encode()).hexdigest()[:16] + + @staticmethod + def _hash_structural_signature( + sig_input: List[Tuple[str, Tuple[int, ...], str]], + ) -> str: + h = hashlib.sha1(repr(sig_input).encode()) + return h.hexdigest()[:16] + + @staticmethod + def _world_rank_and_size() -> Tuple[int, int]: + from sglang.srt.distributed import get_world_group + + try: + g = get_world_group() + return g.rank_in_group, g.world_size + except (AssertionError, AttributeError): + return 0, 1 + + @staticmethod + def _world_barrier() -> None: + from sglang.srt.distributed import get_world_group + + try: + get_world_group().barrier() + except (AssertionError, AttributeError): + pass + + @staticmethod + def _new_content_hasher(): + import xxhash + + return xxhash.xxh3_128() + + @staticmethod + def _hash_tensor(tensor: torch.Tensor) -> str: + # CPU copy so concurrent dump workers cannot race CUDA D2H hashing. + t = tensor.detach() + prefix = str(tuple(t.shape)).encode() + str(t.dtype).encode() + h = PreshardedModelLoader._new_content_hasher() + h.update(prefix) + + if t.numel() == 0: + return h.hexdigest() + + cpu = t.contiguous().to(device="cpu", copy=True).contiguous() + flat_u8 = cpu.reshape(-1).view(torch.uint8) + h.update(memoryview(flat_u8.numpy())) + return h.hexdigest() + + def _verify_rank_checksum( + self, + verify_hashes: List[Tuple[str, str]], + plan: Dict[str, Any], + rank: int, + presharded_dir: str, + ) -> None: + expected = plan.get("rank_checksums", {}).get(str(rank)) + if expected is None: + raise ValueError( + f"Plan at {presharded_dir} has no rank_checksums entry for " + f"rank {rank}; cannot verify. Set " + f"--model-loader-extra-config '{{\"verify_on_load\": false}}' " + f"to skip verification, or re-dump the checkpoint." + ) + + total = 0 + for name, content_hash in verify_hashes: + d = PreshardedModelLoader._fold_name_content_digest(name, content_hash) + total = (total + int.from_bytes(d[:8], "big")) & 0xFFFFFFFFFFFFFFFF + actual = format(total, "016x") + + if actual != expected: + raise ValueError( + f"Rank-{rank} checksum mismatch for presharded checkpoint at " + f"{presharded_dir}: expected {expected}, got {actual}. The " + f"checkpoint files may be corrupted; re-dump or skip " + f"verification with --model-loader-extra-config " + f"'{{\"verify_on_load\": false}}'." + ) + + @staticmethod + def _fold_name_content_digest(name: str, content_hash: str) -> bytes: + h = PreshardedModelLoader._new_content_hasher() + h.update((name + ":" + content_hash).encode("utf-8")) + return h.digest() + + @staticmethod + def _collect_extra_tensors(model: nn.Module) -> Dict[str, torch.Tensor]: + seen: set = set() + param_storages: set = set() + for name, tensor in model.state_dict().items(): + seen.add(name) + if tensor.numel() > 0: + param_storages.add((tensor.device, tensor.untyped_storage().data_ptr())) + extras: Dict[str, torch.Tensor] = {} + for module_name, module in model.named_modules(): + prefix = f"{module_name}." if module_name else "" + for attr_name in list(vars(module).keys()): + if attr_name.startswith("_"): + continue + try: + val = getattr(module, attr_name) + except AttributeError: + continue + if isinstance(val, torch.Tensor) and not isinstance( + val, torch.nn.Parameter + ): + full_name = f"{prefix}{attr_name}" + if full_name in seen: + continue + if val.numel() > 0: + key = (val.device, val.untyped_storage().data_ptr()) + if key in param_storages: + continue + extras[full_name] = val + return extras + + @staticmethod + def _rebind_parameter_aliases(model: nn.Module) -> None: + for _, module in model.named_modules(): + gemma_w = getattr(module, "gemma_weight", None) + weight = getattr(module, "weight", None) + if ( + isinstance(gemma_w, torch.Tensor) + and isinstance(weight, torch.nn.Parameter) + and gemma_w.shape == weight.shape + ): + torch.add(weight.data, 1.0, out=gemma_w) + + attn = getattr(module, "attn", None) + conv1d = getattr(module, "conv1d", None) + if attn is None: + continue + if hasattr(module, "A_log") and hasattr(attn, "A_log"): + attn.A_log = module.A_log + if hasattr(module, "dt_bias") and hasattr(attn, "dt_bias"): + attn.dt_bias = module.dt_bias + if conv1d is None: + continue + cweight = getattr(conv1d, "weight", None) + if cweight is not None and hasattr(attn, "conv_weights"): + if cweight.dim() == 3 and cweight.size(1) == 1: + attn.conv_weights = cweight.view(cweight.size(0), cweight.size(2)) + else: + attn.conv_weights = ( + cweight.squeeze() if cweight.dim() > 2 else cweight + ) + if hasattr(conv1d, "bias") and hasattr(attn, "bias"): + attn.bias = conv1d.bias + + def _ensure_presharded_dir_writable(self, presharded_dir: str) -> None: + rank, _ = self._world_rank_and_size() + try: + os.makedirs(presharded_dir, exist_ok=True) + if rank == 0: + probe = os.path.join(presharded_dir, ".presharded_write_probe") + last_err: Optional[OSError] = None + for _ in range(5): + try: + with open(probe, "w") as f: + f.write("ok") + os.unlink(probe) + last_err = None + break + except OSError as e: + last_err = e + os.makedirs(presharded_dir, exist_ok=True) + time.sleep(0.05) + if last_err is not None: + raise last_err + except OSError as e: + raise RuntimeError( + f"Presharded dump directory is not writable: {presharded_dir}. " + "Set model_loader_extra_config " + '\'{"presharded_path": "..."}\' (or draft_presharded_path for ' + "the draft model) to a writable shared filesystem path. " + f"Original error: {e}" + ) from e + self._world_barrier() + + def _first_time_load_and_dump( + self, + model_config: ModelConfig, + device_config: DeviceConfig, + presharded_dir: str, + shard_config: Dict[str, Any], + ) -> nn.Module: + self._ensure_presharded_dir_writable(presharded_dir) + target_device = torch.device(device_config.device) + quant_config = _get_quantization_config(model_config, self.load_config) + with set_default_torch_dtype(model_config.dtype): + with target_device: + model = _initialize_model(model_config, self.load_config, quant_config) + self.load_weights_and_postprocess( + model, + self._get_all_weights(model_config, model), + target_device, + ) + + state_dict = dict(model.state_dict()) + extras = self._collect_extra_tensors(model) + self._dump_state_to_disk(state_dict, extras, presharded_dir, shard_config) + del state_dict + del extras + gc.collect() + + self.counter_after_loading_weights = time.perf_counter() + return model.eval() + + def _dump_state_to_disk( + self, + state_dict: Dict[str, torch.Tensor], + extras: Dict[str, torch.Tensor], + presharded_dir: str, + shard_config: Dict[str, Any], + ) -> None: + rank, world_size = self._world_rank_and_size() + tmp_dir = os.path.join(presharded_dir, self.TMP_SUBDIR) + if rank == 0: + ready_path = os.path.join(presharded_dir, self.READY_FILENAME) + if os.path.isfile(ready_path): + os.unlink(ready_path) + os.makedirs(tmp_dir, exist_ok=True) + self._world_barrier() + + items: List[Tuple[str, torch.Tensor, bool]] = [] + items.extend((n, t, False) for n, t in state_dict.items()) + items.extend((n, t, True) for n, t in extras.items()) + + def _entry(item: Tuple[str, torch.Tensor, bool]) -> Tuple[str, Dict[str, Any]]: + name, tensor, is_extra = item + return name, { + "checksum": self._hash_tensor(tensor), + "size": tensor.numel() * tensor.element_size(), + "dtype": str(tensor.dtype), + "shape": list(tensor.shape), + "is_extra": is_extra, + } + + manifest: Dict[str, Dict[str, Any]] = {} + num_workers = min(max(1, len(items)), self._hash_num_threads) + if num_workers <= 1: + for it in items: + name, info = _entry(it) + manifest[name] = info + else: + with concurrent.futures.ThreadPoolExecutor( + max_workers=num_workers, + thread_name_prefix="presharded-hash", + ) as ex: + for name, info in ex.map(_entry, items): + manifest[name] = info + + with open(os.path.join(tmp_dir, f"manifest_{rank:05d}.json"), "w") as f: + json.dump(manifest, f) + self._world_barrier() + + if rank == 0: + plan = self._build_dump_plan(world_size, tmp_dir, self._max_file_bytes) + plan["shard_config"] = shard_config + with open(os.path.join(presharded_dir, self.CHECKSUM_FILENAME), "w") as f: + json.dump(plan, f, indent=2) + self._world_barrier() + + with open(os.path.join(presharded_dir, self.CHECKSUM_FILENAME)) as f: + plan = json.load(f) + all_tensors = {**state_dict, **extras} + self._dump_files_for_rank(all_tensors, plan, rank, presharded_dir) + self._world_barrier() + + if rank == 0: + shutil.rmtree(tmp_dir, ignore_errors=True) + ready_path = os.path.join(presharded_dir, self.READY_FILENAME) + with open(ready_path, "w") as f: + json.dump( + { + "plan_version": self.PLAN_VERSION, + "world_size": world_size, + "created_at": time.time(), + }, + f, + ) + self._world_barrier() + + @staticmethod + def _make_filename( + file_id: int, rank_list: Tuple[int, ...], is_common: bool + ) -> str: + if is_common: + return f"model-{file_id:05d}-common.safetensor" + rank_str = ",".join(f"{r:03d}" for r in rank_list) + return f"model-{file_id:05d}-rank-{rank_str}.safetensor" + + @classmethod + def _build_dump_plan( + cls, world_size: int, tmp_dir: str, max_file_bytes: int + ) -> Dict[str, Any]: + rank_to_manifest: Dict[int, Dict[str, Dict[str, Any]]] = {} + for r in range(world_size): + manifest_path = os.path.join(tmp_dir, f"manifest_{r:05d}.json") + try: + with open(manifest_path) as f: + rank_to_manifest[r] = json.load(f) + except FileNotFoundError as e: + raise FileNotFoundError( + f"Rank {r} did not write {manifest_path}. The presharded " + "dump directory must be on a filesystem shared by all " + "ranks/nodes (set presharded_path / draft_presharded_path " + "to a shared path if model_path is node-local)." + ) from e + + checksum_to_entries: Dict[str, List[Tuple[int, str, Dict[str, Any]]]] = ( + collections.defaultdict(list) + ) + name_to_is_extra: Dict[Tuple[int, str], bool] = {} + for r, manifest in rank_to_manifest.items(): + for name, info in manifest.items(): + checksum_to_entries[info["checksum"]].append((r, name, info)) + name_to_is_extra[(r, name)] = bool(info.get("is_extra", False)) + + tensor_records: List[Dict[str, Any]] = [] + for checksum, entries in checksum_to_entries.items(): + sizes = {info["size"] for _, _, info in entries} + if len(sizes) != 1: + raise RuntimeError( + f"Checksum {checksum} maps to inconsistent sizes {sizes}; " + f"this indicates a hash collision or stale manifest." + ) + size = next(iter(sizes)) + ranks = sorted({r for r, _, _ in entries}) + rank_to_names: Dict[str, List[str]] = collections.defaultdict(list) + for r, n, _ in entries: + rank_to_names[str(r)].append(n) + tensor_records.append( + { + "checksum": checksum, + "size": size, + "rank_list": ranks, + "rank_to_names": {k: sorted(v) for k, v in rank_to_names.items()}, + } + ) + + by_rank_list: Dict[Tuple[int, ...], List[Dict[str, Any]]] = ( + collections.defaultdict(list) + ) + for rec in tensor_records: + by_rank_list[tuple(rec["rank_list"])].append(rec) + + files: List[Dict[str, Any]] = [] + next_file_id = 0 + for rank_tuple, recs in by_rank_list.items(): + recs.sort(key=lambda r: -r["size"]) + is_common = len(rank_tuple) == world_size and rank_tuple == tuple( + range(world_size) + ) + writer_load = {wr: 0 for wr in rank_tuple} + writer_records: Dict[int, List[Dict[str, Any]]] = { + wr: [] for wr in rank_tuple + } + for rec in recs: + wr = min(rank_tuple, key=lambda r: writer_load[r]) + writer_records[wr].append(rec) + writer_load[wr] += rec["size"] + + for wr, wr_recs in writer_records.items(): + cur_size = 0 + cur_tensors: List[Dict[str, Any]] = [] + for rec in wr_recs: + if cur_tensors and cur_size + rec["size"] > max_file_bytes: + files.append( + { + "filename": cls._make_filename( + next_file_id, rank_tuple, is_common + ), + "writer_rank": wr, + "rank_list": (None if is_common else list(rank_tuple)), + "is_common": is_common, + "tensors": cur_tensors, + } + ) + next_file_id += 1 + cur_size = 0 + cur_tensors = [] + cur_tensors.append( + { + "stored_key": rec["checksum"], + "size": rec["size"], + "rank_to_names": rec["rank_to_names"], + } + ) + cur_size += rec["size"] + if cur_tensors: + files.append( + { + "filename": cls._make_filename( + next_file_id, rank_tuple, is_common + ), + "writer_rank": wr, + "rank_list": (None if is_common else list(rank_tuple)), + "is_common": is_common, + "tensors": cur_tensors, + } + ) + next_file_id += 1 + + rank_to_reads: Dict[int, List[Dict[str, Any]]] = collections.defaultdict(list) + for f in files: + for t in f["tensors"]: + for r_str, names in t["rank_to_names"].items(): + for name in names: + rank_to_reads[int(r_str)].append( + { + "filename": f["filename"], + "stored_key": t["stored_key"], + "name": name, + "is_extra": name_to_is_extra.get( + (int(r_str), name), False + ), + } + ) + + rank_checksums: Dict[str, str] = {} + for r in range(world_size): + total = 0 + for rec in rank_to_reads.get(r, []): + d = cls._fold_name_content_digest(rec["name"], rec["stored_key"]) + total = (total + int.from_bytes(d[:8], "big")) & 0xFFFFFFFFFFFFFFFF + rank_checksums[str(r)] = format(total, "016x") + + return { + "version": cls.PLAN_VERSION, + "world_size": world_size, + "files": files, + "rank_to_reads": {str(r): v for r, v in rank_to_reads.items()}, + "rank_checksums": rank_checksums, + } + + def _dump_files_for_rank( + self, + state_dict: Dict[str, torch.Tensor], + plan: Dict[str, Any], + rank: int, + presharded_dir: str, + ) -> None: + from safetensors.torch import save_file + + for f in plan["files"]: + if f["writer_rank"] != rank: + continue + tensors_to_save: Dict[str, torch.Tensor] = {} + for t in f["tensors"]: + names_for_this_rank = t["rank_to_names"].get(str(rank)) + if not names_for_this_rank: + raise RuntimeError( + f"writer_rank {rank} is missing tensor {t['stored_key']} " + f"for file {f['filename']}; plan is inconsistent." + ) + name_for_this_rank = names_for_this_rank[0] + tensor = ( + state_dict[name_for_this_rank] + .detach() + .to(device="cpu", copy=False) + .contiguous() + ) + tensors_to_save[t["stored_key"]] = tensor + save_file(tensors_to_save, os.path.join(presharded_dir, f["filename"])) + + @staticmethod + def _read_presharded_file( + full_path: str, stored_keys: List[str] + ) -> Dict[str, torch.Tensor]: + from safetensors.torch import safe_open + + with safe_open(full_path, framework="pt") as fh: + return {key: fh.get_tensor(key) for key in stored_keys} + + def _apply_presharded_file( + self, + *, + items: List[Dict[str, Any]], + cached: Dict[str, torch.Tensor], + model: nn.Module, + state_dict: Dict[str, torch.Tensor], + target_device: torch.device, + loaded_param_keys: set, + verify_hashes: List[Tuple[str, str]], + ) -> None: + if self._verify_on_load: + keys = list(cached.keys()) + n_workers = min(max(1, len(keys)), self._hash_num_threads) + + def _hash_one(key, _cached=cached): + return key, self._hash_tensor(_cached[key]) + + if n_workers <= 1: + key_to_hash = dict(_hash_one(k) for k in keys) + else: + with concurrent.futures.ThreadPoolExecutor( + max_workers=n_workers, + thread_name_prefix="presharded-verify", + ) as ex: + key_to_hash = dict(ex.map(_hash_one, keys)) + for r in items: + verify_hashes.append((r["name"], key_to_hash[r["stored_key"]])) + + for r in items: + tensor = cached[r["stored_key"]] + if r.get("is_extra"): + module_path, _, attr_name = r["name"].rpartition(".") + module = model.get_submodule(module_path) if module_path else model + if hasattr(module, attr_name): + try: + delattr(module, attr_name) + except AttributeError: + pass + setattr(module, attr_name, tensor.to(target_device)) + continue + if r["name"] not in state_dict: + continue + param_data = state_dict[r["name"]].data + param_shape = state_dict[r["name"]].shape + for dim, size in enumerate(tensor.shape): + if size < param_shape[dim]: + param_data = param_data.narrow(dim, 0, size) + if tensor.shape != param_data.shape: + raise ValueError( + f"Presharded tensor shape mismatch for '{r['name']}': " + f"dumped {tuple(tensor.shape)} vs parameter slice " + f"{tuple(param_data.shape)} (full param {tuple(param_shape)}). " + "Re-dump with matching quant/parallel config, or set " + "verify_on_load and check process_weights_after_loading." + ) + param_data.copy_(tensor) + loaded_param_keys.add(r["name"]) + + cached.clear() + del cached + + def _load_from_presharded( + self, + model_config: ModelConfig, + device_config: DeviceConfig, + presharded_dir: str, + ) -> nn.Module: + target_device = torch.device(device_config.device) + quant_config = _get_quantization_config(model_config, self.load_config) + + with set_default_torch_dtype(model_config.dtype): + with target_device: + model = _initialize_model(model_config, self.load_config, quant_config) + + for _, module in model.named_modules(): + quant_method = getattr(module, "quant_method", None) + if quant_method is not None: + with device_loading_context(module, target_device): + quant_method.process_weights_after_loading(module) + + rank, _ = self._world_rank_and_size() + with open(os.path.join(presharded_dir, self.CHECKSUM_FILENAME)) as f: + plan = json.load(f) + if plan.get("version") != self.PLAN_VERSION: + raise ValueError( + f"Unsupported presharded plan version {plan.get('version')!r} " + f"at {presharded_dir}; expected {self.PLAN_VERSION}." + ) + + state_dict = dict(model.state_dict()) + reads = plan.get("rank_to_reads", {}).get(str(rank), []) + + by_file: Dict[str, List[Dict[str, Any]]] = collections.defaultdict(list) + for r in reads: + by_file[r["filename"]].append(r) + + loaded_param_keys: set = set() + verify_hashes: List[Tuple[str, str]] = [] + for filename, items in by_file.items(): + stored_keys = list(dict.fromkeys(r["stored_key"] for r in items)) + cached = self._read_presharded_file( + os.path.join(presharded_dir, filename), stored_keys + ) + self._apply_presharded_file( + items=items, + cached=cached, + model=model, + state_dict=state_dict, + target_device=target_device, + loaded_param_keys=loaded_param_keys, + verify_hashes=verify_hashes, + ) + if torch.cuda.is_available(): + torch.cuda.synchronize() + + loaded_storages: set = set() + for k in loaded_param_keys: + t = state_dict[k] + if t.numel() > 0: + loaded_storages.add((t.device, t.untyped_storage().data_ptr())) + missing = [] + for k, t in state_dict.items(): + if k in loaded_param_keys: + continue + if t.numel() == 0: + continue + storage_key = (t.device, t.untyped_storage().data_ptr()) + if storage_key not in loaded_storages: + missing.append(k) + if missing: + raise ValueError( + f"Missing keys {tuple(sorted(missing))} in presharded " + f"checkpoint at {presharded_dir}." + ) + + self._rebind_parameter_aliases(model) + + if self._verify_on_load: + self._verify_rank_checksum(verify_hashes, plan, rank, presharded_dir) + + self.counter_after_loading_weights = time.perf_counter() + return model.eval() + + class BitsAndBytesModelLoader(BaseModelLoader): """Model loader to load model weights with BitAndBytes quantization.""" @@ -3271,6 +4116,9 @@ def get_model_loader( if load_config.load_format == LoadFormat.SHARDED_STATE: return ShardedStateLoader(load_config) + if load_config.load_format == LoadFormat.PRESHARDED: + return PreshardedModelLoader(load_config) + if load_config.load_format == LoadFormat.BITSANDBYTES: return BitsAndBytesModelLoader(load_config) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 2adfb873f..0887fc71a 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -113,6 +113,7 @@ LOAD_FORMAT_CHOICES = [ "npcache", "dummy", "sharded_state", + "presharded", "gguf", "bitsandbytes", "mistral", @@ -517,14 +518,31 @@ class ServerArgs: "quantization." '"layered" loads weights layer by layer so that one can quantize a ' "layer before loading another to make the peak memory envelope " - "smaller.", + "smaller." + '"presharded" performs a normal first-time load (with quantization), ' + "then dumps a per-rank/per-tensor sharded checkpoint with content " + "deduplication into " + "/presharded//. " + "Subsequent runs with the same parallelism+quantization config " + "load directly from this presharded checkpoint and skip " + "re-quantization. " + "The dump directory must be on a shared filesystem across all " + "ranks/nodes. Optional model_loader_extra_config roots: " + "presharded_path (target) and draft_presharded_path (speculative " + "draft); each replaces /presharded and still gets a " + "config subfolder appended. Use a writable path when model_path " + "is read-only (e.g. HF cache mounts).", choices=LOAD_FORMAT_CHOICES, ), NS("model"), ] = "auto" model_loader_extra_config: A[ str, - "Extra config for model loader. This will be passed to the model loader corresponding to the chosen load_format.", + "Extra config for model loader. This will be passed to the model loader " + "corresponding to the chosen load_format. For load_format=presharded, " + "JSON may include presharded_path (target cache root), " + "draft_presharded_path (draft cache root), max_file_bytes, " + "hash_num_threads, and verify_on_load.", NS("model"), ] = "{}" trust_remote_code: A[ diff --git a/test/registered/unit/model_loader/test_presharded_loader.py b/test/registered/unit/model_loader/test_presharded_loader.py new file mode 100644 index 000000000..b2a8dda86 --- /dev/null +++ b/test/registered/unit/model_loader/test_presharded_loader.py @@ -0,0 +1,960 @@ +"""Unit tests for PreshardedModelLoader's pure helpers. + +These tests exercise the deterministic pieces of the presharding algorithm +(tensor hashing, plan construction, file naming, dedup, file-size cap, and +per-rank workload balance) without needing a GPU or distributed setup. +""" + +import json +import os +import tempfile +import unittest +from types import SimpleNamespace +from unittest import mock + +import torch + +from sglang.srt.model_loader.loader import PreshardedModelLoader +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +class TestPreshardedHashTensor(unittest.TestCase): + def test_same_content_same_hash(self): + a = torch.arange(100, dtype=torch.float32).reshape(10, 10) + b = torch.arange(100, dtype=torch.float32).reshape(10, 10) + self.assertEqual( + PreshardedModelLoader._hash_tensor(a), + PreshardedModelLoader._hash_tensor(b), + ) + + def test_different_content_different_hash(self): + a = torch.arange(100, dtype=torch.float32) + b = a.clone() + b[0] = 999 + self.assertNotEqual( + PreshardedModelLoader._hash_tensor(a), + PreshardedModelLoader._hash_tensor(b), + ) + + def test_dtype_changes_hash(self): + a = torch.zeros(8, dtype=torch.float32) + b = torch.zeros(16, dtype=torch.float16) # same byte content (zeros) + self.assertNotEqual( + PreshardedModelLoader._hash_tensor(a), + PreshardedModelLoader._hash_tensor(b), + ) + + def test_shape_changes_hash(self): + a = torch.zeros(16, dtype=torch.float32) + b = torch.zeros((4, 4), dtype=torch.float32) + self.assertNotEqual( + PreshardedModelLoader._hash_tensor(a), + PreshardedModelLoader._hash_tensor(b), + ) + + def test_empty_tensor_is_hashable(self): + a = torch.empty(0, dtype=torch.float32) + digest = PreshardedModelLoader._hash_tensor(a) + self.assertIsInstance(digest, str) + self.assertEqual( + len(digest), PreshardedModelLoader._CONTENT_HASH_HEX_LEN + ) # xxh3-128 hex + + def test_cuda_and_cpu_digests_agree(self): + # Streaming D2H + host xxh3 must match a full CPU hash of the same + # bytes; otherwise dump/reload verify would false-fail across devices. + if not torch.cuda.is_available(): + self.skipTest("CUDA required") + cpu = torch.arange(10_000, dtype=torch.float32) + gpu = cpu.cuda() + self.assertEqual( + PreshardedModelLoader._hash_tensor(cpu), + PreshardedModelLoader._hash_tensor(gpu), + ) + + +class TestPreshardedFilename(unittest.TestCase): + def test_common_filename(self): + self.assertEqual( + PreshardedModelLoader._make_filename(0, (0, 1, 2, 3), is_common=True), + "model-00000-common.safetensor", + ) + + def test_rank_filename_three_digit_padding(self): + self.assertEqual( + PreshardedModelLoader._make_filename(5, (1, 3, 5, 7), is_common=False), + "model-00005-rank-001,003,005,007.safetensor", + ) + + def test_file_id_zero_padding(self): + self.assertTrue( + PreshardedModelLoader._make_filename(42, (0,), is_common=False).startswith( + "model-00042-" + ) + ) + + +class TestBuildDumpPlan(unittest.TestCase): + def _write_manifests(self, tmp_dir, manifests): + for r, m in manifests.items(): + with open(os.path.join(tmp_dir, f"manifest_{r:05d}.json"), "w") as f: + json.dump(m, f) + + def test_dedup_across_ranks(self): + # Same content (same checksum) on all 4 ranks → single file marked common. + with tempfile.TemporaryDirectory() as tmp: + shared = { + "checksum": "deadbeef", + "size": 1024, + "dtype": "torch.float32", + "shape": [256], + } + manifests = {r: {"shared.weight": shared} for r in range(4)} + self._write_manifests(tmp, manifests) + plan = PreshardedModelLoader._build_dump_plan( + world_size=4, tmp_dir=tmp, max_file_bytes=10**12 + ) + self.assertEqual(len(plan["files"]), 1) + self.assertTrue(plan["files"][0]["is_common"]) + self.assertIn("common", plan["files"][0]["filename"]) + # Each rank should still have a read entry pointing at the file. + for r in range(4): + reads = plan["rank_to_reads"][str(r)] + self.assertEqual(len(reads), 1) + self.assertEqual(reads[0]["name"], "shared.weight") + self.assertEqual(reads[0]["filename"], plan["files"][0]["filename"]) + self.assertEqual(reads[0]["stored_key"], "deadbeef") + self.assertIn("rank_checksums", plan) + self.assertEqual(set(plan["rank_checksums"].keys()), {"0", "1", "2", "3"}) + + def test_per_rank_unique_tensors(self): + # Each rank has its own tensor (different content). 4 distinct files, + # filenames should be rank-{rrr}. + with tempfile.TemporaryDirectory() as tmp: + manifests = { + r: { + "layer.weight": { + "checksum": f"hash_{r}", + "size": 2048, + "dtype": "torch.float32", + "shape": [512], + } + } + for r in range(4) + } + self._write_manifests(tmp, manifests) + plan = PreshardedModelLoader._build_dump_plan( + world_size=4, tmp_dir=tmp, max_file_bytes=10**12 + ) + self.assertEqual(len(plan["files"]), 4) + for f in plan["files"]: + self.assertFalse(f["is_common"]) + self.assertEqual(len(f["rank_list"]), 1) + # Writer is the only rank in the rank_list. + self.assertEqual(f["writer_rank"], f["rank_list"][0]) + self.assertIn(f"-rank-{f['rank_list'][0]:03d}.safetensor", f["filename"]) + + def test_partial_share_has_correct_rank_list(self): + # Tensor shared by ranks 1,3,5,7 only. + with tempfile.TemporaryDirectory() as tmp: + shared = { + "checksum": "shared_hash", + "size": 1024, + "dtype": "torch.float32", + "shape": [256], + } + manifests = { + r: ({"x": shared} if r in (1, 3, 5, 7) else {}) for r in range(8) + } + self._write_manifests(tmp, manifests) + plan = PreshardedModelLoader._build_dump_plan( + world_size=8, tmp_dir=tmp, max_file_bytes=10**12 + ) + self.assertEqual(len(plan["files"]), 1) + f = plan["files"][0] + self.assertFalse(f["is_common"]) + self.assertEqual(f["rank_list"], [1, 3, 5, 7]) + self.assertIn(f["writer_rank"], (1, 3, 5, 7)) + self.assertEqual(f["filename"], "model-00000-rank-001,003,005,007.safetensor") + # Only the 4 sharing ranks have read entries. + for r in range(8): + reads = plan["rank_to_reads"].get(str(r), []) + if r in (1, 3, 5, 7): + self.assertEqual(len(reads), 1) + else: + self.assertEqual(len(reads), 0) + + def test_max_file_size_caps_files(self): + # Two tensors of 1 MiB each shared by all 2 ranks; cap = 1.5 MiB → + # two files. + with tempfile.TemporaryDirectory() as tmp: + t1 = { + "checksum": "h1", + "size": 1024 * 1024, + "dtype": "torch.float32", + "shape": [256, 1024], + } + t2 = dict(t1, checksum="h2") + manifests = {0: {"a": t1, "b": t2}, 1: {"a": t1, "b": t2}} + self._write_manifests(tmp, manifests) + plan = PreshardedModelLoader._build_dump_plan( + world_size=2, + tmp_dir=tmp, + max_file_bytes=int(1.5 * 1024 * 1024), + ) + # Both tensors share the same rank_list (0,1). With balanced packing, + # each writer (rank 0 and rank 1) gets one tensor → 2 files. + self.assertEqual(len(plan["files"]), 2) + for f in plan["files"]: + self.assertEqual(len(f["tensors"]), 1) + # rank_list is full world ⇒ marked common + self.assertTrue(f["is_common"]) + + def test_workload_balanced_within_rank_list(self): + # 4 same-size tensors all shared by ranks (0,1,2,3) → with balanced + # packing each writer rank should get exactly one tensor. + with tempfile.TemporaryDirectory() as tmp: + tensors = { + f"t{i}": { + "checksum": f"h{i}", + "size": 1024, + "dtype": "torch.float32", + "shape": [256], + } + for i in range(4) + } + manifests = {r: dict(tensors) for r in range(4)} + self._write_manifests(tmp, manifests) + plan = PreshardedModelLoader._build_dump_plan( + world_size=4, + tmp_dir=tmp, + max_file_bytes=10**12, + ) + # 4 tensors, 4 writers, 4 files (one per writer). + self.assertEqual(len(plan["files"]), 4) + writers = sorted(f["writer_rank"] for f in plan["files"]) + self.assertEqual(writers, [0, 1, 2, 3]) + + def test_round_trip_dump_and_read_back(self): + # End-to-end on disk for world_size=1: build manifests for tensors, + # construct plan, write safetensors per the plan, read back and + # verify both checksums and bit-identity. + from safetensors.torch import safe_open, save_file + + with tempfile.TemporaryDirectory() as tmp: + tensors = { + "embed.weight": torch.arange(64, dtype=torch.float32).reshape(8, 8), + "norm.weight": torch.full((16,), 0.5, dtype=torch.float32), + "head.weight": torch.linspace(-1, 1, 32, dtype=torch.float32), + } + manifest_dir = os.path.join(tmp, "manifests") + presharded_dir = os.path.join(tmp, "presharded") + os.makedirs(manifest_dir) + os.makedirs(presharded_dir) + + checksums = { + name: PreshardedModelLoader._hash_tensor(t) + for name, t in tensors.items() + } + manifest = { + name: { + "checksum": checksums[name], + "size": t.numel() * t.element_size(), + "dtype": str(t.dtype), + "shape": list(t.shape), + } + for name, t in tensors.items() + } + with open(os.path.join(manifest_dir, "manifest_00000.json"), "w") as f: + json.dump(manifest, f) + + plan = PreshardedModelLoader._build_dump_plan( + world_size=1, tmp_dir=manifest_dir, max_file_bytes=10**12 + ) + + # Single rank → all tensors live in one common file. + self.assertEqual(len(plan["files"]), 1) + f = plan["files"][0] + self.assertTrue(f["is_common"]) + + # Write the file with stored_keys mapped to tensor content. + to_save = {} + for t_entry in f["tensors"]: + name = t_entry["rank_to_names"]["0"][0] + to_save[t_entry["stored_key"]] = tensors[name] + save_file(to_save, os.path.join(presharded_dir, f["filename"])) + + # Read back: verify each tensor's checksum and content. + with safe_open( + os.path.join(presharded_dir, f["filename"]), framework="pt" + ) as fh: + for r in plan["rank_to_reads"]["0"]: + loaded = fh.get_tensor(r["stored_key"]) + self.assertEqual( + PreshardedModelLoader._hash_tensor(loaded), r["stored_key"] + ) + torch.testing.assert_close(loaded, tensors[r["name"]]) + + def test_dedup_with_multiple_names_per_rank(self): + # Same checksum can appear under MULTIPLE param names on the same + # rank (e.g., k_scale and v_scale both default to 1.0). Both names + # must end up in rank_to_reads. + with tempfile.TemporaryDirectory() as tmp: + shared = { + "checksum": "scale_hash", + "size": 4, + "dtype": "torch.float32", + "shape": [], + } + manifests = { + 0: { + "layers.0.attn.k_scale": shared, + "layers.0.attn.v_scale": shared, + "layers.1.attn.k_scale": shared, + }, + 1: { + "layers.0.attn.k_scale": shared, + "layers.0.attn.v_scale": shared, + "layers.1.attn.k_scale": shared, + }, + } + self._write_manifests(tmp, manifests) + plan = PreshardedModelLoader._build_dump_plan( + world_size=2, tmp_dir=tmp, max_file_bytes=10**12 + ) + # All 6 (rank,name) pairs must be readable, even though there is one + # underlying tensor stored on disk. + self.assertEqual(len(plan["files"]), 1) + for r in (0, 1): + reads = plan["rank_to_reads"][str(r)] + names = sorted(rd["name"] for rd in reads) + self.assertEqual( + names, + [ + "layers.0.attn.k_scale", + "layers.0.attn.v_scale", + "layers.1.attn.k_scale", + ], + ) + # All point at the same stored_key (deduplicated content). + self.assertEqual({rd["stored_key"] for rd in reads}, {"scale_hash"}) + + def test_collision_size_mismatch_raises(self): + # Same checksum but different sizes ⇒ plan builder rejects. + with tempfile.TemporaryDirectory() as tmp: + manifests = { + 0: { + "x": { + "checksum": "same", + "size": 1024, + "dtype": "torch.float32", + "shape": [256], + } + }, + 1: { + "x": { + "checksum": "same", + "size": 2048, + "dtype": "torch.float32", + "shape": [512], + } + }, + } + self._write_manifests(tmp, manifests) + with self.assertRaises(RuntimeError): + PreshardedModelLoader._build_dump_plan( + world_size=2, tmp_dir=tmp, max_file_bytes=10**12 + ) + + def test_rank_checksum_deterministic(self): + # rank_checksums must be reproducible from the same manifest input + # and depend on (name, content-SHA) pairs of every tensor a rank + # owns. Permuting the manifest's insertion order must not change + # the rank checksum. + with tempfile.TemporaryDirectory() as tmp_a, tempfile.TemporaryDirectory() as tmp_b: + base_entries = { + "alpha.weight": { + "checksum": "h_alpha", + "size": 16, + "dtype": "torch.float32", + "shape": [4], + }, + "beta.weight": { + "checksum": "h_beta", + "size": 16, + "dtype": "torch.float32", + "shape": [4], + }, + } + self._write_manifests(tmp_a, {0: dict(base_entries)}) + # Insertion-order-permuted copy. + permuted = {k: base_entries[k] for k in reversed(list(base_entries))} + self._write_manifests(tmp_b, {0: permuted}) + plan_a = PreshardedModelLoader._build_dump_plan( + world_size=1, tmp_dir=tmp_a, max_file_bytes=10**12 + ) + plan_b = PreshardedModelLoader._build_dump_plan( + world_size=1, tmp_dir=tmp_b, max_file_bytes=10**12 + ) + self.assertEqual(plan_a["rank_checksums"], plan_b["rank_checksums"]) + + def test_rank_checksum_distinguishes_content(self): + # Changing one tensor's content-SHA must change the rank checksum. + with tempfile.TemporaryDirectory() as tmp_a, tempfile.TemporaryDirectory() as tmp_b: + entries_a = { + "x.weight": { + "checksum": "ha", + "size": 16, + "dtype": "torch.float32", + "shape": [4], + }, + } + entries_b = { + "x.weight": { + "checksum": "hb", # different content + "size": 16, + "dtype": "torch.float32", + "shape": [4], + }, + } + self._write_manifests(tmp_a, {0: entries_a}) + self._write_manifests(tmp_b, {0: entries_b}) + plan_a = PreshardedModelLoader._build_dump_plan( + world_size=1, tmp_dir=tmp_a, max_file_bytes=10**12 + ) + plan_b = PreshardedModelLoader._build_dump_plan( + world_size=1, tmp_dir=tmp_b, max_file_bytes=10**12 + ) + self.assertNotEqual( + plan_a["rank_checksums"]["0"], + plan_b["rank_checksums"]["0"], + ) + + def test_presharded_ready_sentinel(self): + # Loader treats a dir as a valid presharded ckpt only when the + # READY sentinel exists. A bare checksum.json (or partial files) + # is NOT enough. + with tempfile.TemporaryDirectory() as tmp: + self.assertFalse(PreshardedModelLoader._presharded_ready(tmp)) + # checksum.json alone is not sufficient. + with open( + os.path.join(tmp, PreshardedModelLoader.CHECKSUM_FILENAME), "w" + ) as f: + f.write("{}") + self.assertFalse(PreshardedModelLoader._presharded_ready(tmp)) + # READY makes it ready. + with open( + os.path.join(tmp, PreshardedModelLoader.READY_FILENAME), "w" + ) as f: + f.write("{}") + self.assertTrue(PreshardedModelLoader._presharded_ready(tmp)) + + def test_separate_presharded_path_overrides_for_target_and_draft(self): + # Target and draft get distinct roots via presharded_path vs + # draft_presharded_path. Failure mode if this regresses: draft + # re-dump collides with / wipes target READY under a shared path. + loader = object.__new__(PreshardedModelLoader) + with tempfile.TemporaryDirectory() as tmp: + target_root = os.path.join(tmp, "target_cache") + draft_root = os.path.join(tmp, "draft_cache") + loader._presharded_path_override = target_root + loader._draft_presharded_path_override = draft_root + cfg = {"tp": 8, "structural_signature": "same_sig"} + sub = loader._build_subfolder_name(cfg) + target_mc = SimpleNamespace(model_path="/models/dsv3", is_draft_model=False) + draft_mc = SimpleNamespace(model_path="/models/dsv3", is_draft_model=True) + target_dir = loader._presharded_dir(target_mc, cfg) + draft_dir = loader._presharded_dir(draft_mc, cfg) + self.assertEqual(target_dir, os.path.join(target_root, sub)) + self.assertEqual(draft_dir, os.path.join(draft_root, sub)) + self.assertNotEqual(target_dir, draft_dir) + + # Target-only override: draft must not fall back into the target + # root (same model_path is common for DeepSeek MTP). + loader._draft_presharded_path_override = None + draft_fallback = loader._presharded_dir(draft_mc, cfg) + self.assertEqual( + draft_fallback, + os.path.join( + "/models/dsv3", + PreshardedModelLoader.DEFAULT_SUBDIR, + sub, + ), + ) + self.assertFalse(draft_fallback.startswith(target_root)) + + # No overrides: both use model_path/presharded/. + loader._presharded_path_override = None + self.assertEqual( + loader._presharded_dir(target_mc, cfg), + os.path.join( + "/models/dsv3", + PreshardedModelLoader.DEFAULT_SUBDIR, + sub, + ), + ) + + def test_apply_shape_mismatch_raises(self): + # Reload must not silently copy_ into a wrong layout when process + # shapes and dumped tensors disagree (previously only warned). + # Use a *larger* dumped tensor so prefix-narrow does not paper over it. + loader = object.__new__(PreshardedModelLoader) + loader._verify_on_load = False + param = torch.nn.Parameter(torch.zeros(2, 2)) + state_dict = {"w": param} + items = [{"name": "w", "stored_key": "w", "is_extra": False}] + cached = {"w": torch.ones(4, 4)} + loaded: set = set() + with self.assertRaises(ValueError) as ctx: + loader._apply_presharded_file( + items=items, + cached=cached, + model=torch.nn.Module(), + state_dict=state_dict, + target_device=torch.device("cpu"), + loaded_param_keys=loaded, + verify_hashes=[], + ) + self.assertIn("shape mismatch", str(ctx.exception).lower()) + + def test_build_dump_plan_missing_manifest_mentions_shared_fs(self): + # Multi-node without a shared dump dir fails at plan build with a + # clear message (not a bare FileNotFoundError path). + with tempfile.TemporaryDirectory() as tmp: + with self.assertRaises(FileNotFoundError) as ctx: + PreshardedModelLoader._build_dump_plan( + world_size=2, tmp_dir=tmp, max_file_bytes=1024 + ) + msg = str(ctx.exception) + self.assertIn("Rank 0", msg) + self.assertIn("shared", msg.lower()) + + def test_ensure_presharded_dir_writable_rejects_readonly(self): + # Guard against spending a full source load before discovering a + # read-only dump root (HF cache mounts). Mock OSError because root + # can often still write to mode-0555 dirs under CAP_DAC_OVERRIDE. + loader = object.__new__(PreshardedModelLoader) + with mock.patch.object( + loader, "_world_rank_and_size", return_value=(0, 1) + ), mock.patch.object(loader, "_world_barrier"), mock.patch( + "os.makedirs", side_effect=OSError("Read-only file system") + ): + with self.assertRaises(RuntimeError) as ctx: + loader._ensure_presharded_dir_writable("/ro/presharded") + self.assertIn("not writable", str(ctx.exception).lower()) + self.assertIn("presharded_path", str(ctx.exception)) + + def test_ensure_presharded_dir_writable_ok_rank0(self): + loader = object.__new__(PreshardedModelLoader) + with tempfile.TemporaryDirectory() as tmp: + leaf = os.path.join(tmp, "TP-8-sig-test") + with mock.patch.object( + loader, "_world_rank_and_size", return_value=(0, 1) + ), mock.patch.object(loader, "_world_barrier") as barrier: + loader._ensure_presharded_dir_writable(leaf) + self.assertTrue(os.path.isdir(leaf)) + barrier.assert_called_once() + + +class TestStructuralSignature(unittest.TestCase): + """The structural signature is the long-term fix for the underlying + problem `moe_dense_tp_size` was an instance of: instead of hand- + enumerating every parallelism knob that might change a rank's tensor + shapes, hash the (name, shape, dtype) of every parameter in a + meta-device model skeleton built under the live parallel state. Any + future sharding knob that changes shapes is automatically caught + without touching `_build_subfolder_name`.""" + + def test_hash_changes_with_shape_dtype_or_added_tensor(self): + # Regression: a future sharding knob that changes shapes/dtypes or + # adds per-rank buffers must change the digest, otherwise cache + # collisions load wrong weights. One case covers the three axes that + # `_hash_structural_signature` is responsible for. + base = [("a.weight", (4, 4), "torch.float32")] + shape = [("a.weight", (4, 8), "torch.float32")] + dtype = [("a.weight", (4, 4), "torch.float16")] + extra = [ + ("a.weight", (4, 4), "torch.float32"), + ("a.extra_buf", (4,), "torch.float32"), + ] + h = PreshardedModelLoader._hash_structural_signature + self.assertNotEqual(h(base), h(shape)) + self.assertNotEqual(h(base), h(dtype)) + self.assertNotEqual(h(base), h(extra)) + + def test_local_signature_sorts_state_dict_order(self): + # Production path sorts state_dict items before hashing. If that + # sorted(...) is dropped, two identical models whose state_dict + # iteration order differs would get different signatures and thrash + # the cache. Guard the sort, not just the pure hash helper. + import torch.nn as nn + + class OrderedModule(nn.Module): + def __init__(self, order): + super().__init__() + for name in order: + self.register_parameter( + name, nn.Parameter(torch.zeros(2, 2), requires_grad=False) + ) + + def _init_stub(model_config, load_config, quant_config): + return OrderedModule(model_config.order) + + loader = object.__new__(PreshardedModelLoader) + loader.load_config = SimpleNamespace() + with mock.patch( + "sglang.srt.model_loader.loader._get_quantization_config", + return_value=None, + ), mock.patch( + "sglang.srt.model_loader.loader._initialize_model", + side_effect=_init_stub, + ), mock.patch( + "sglang.srt.model_loader.loader.set_default_torch_dtype", + return_value=mock.MagicMock(__enter__=mock.Mock(), __exit__=mock.Mock()), + ): + sig_ab = loader._compute_local_structural_signature( + SimpleNamespace( + quantization=None, dtype=torch.float32, order=["a", "b"] + ) + ) + sig_ba = loader._compute_local_structural_signature( + SimpleNamespace( + quantization=None, dtype=torch.float32, order=["b", "a"] + ) + ) + self.assertIsNotNone(sig_ab) + self.assertEqual(sig_ab, sig_ba) + + def test_rank_invariant_signature_aggregates_per_rank_locals(self): + # Under PP, ranks build different local digests. The shared cache key + # must still agree across ranks: all-gather then hash the ordered + # list. Without this, each PP stage would pick a different + # presharded subfolder and the multi-rank dump protocol breaks. + fake_group = mock.Mock() + fake_group.world_size = 2 + # Simulate two ranks each calling with their own local sig; both + # must see the same gathered list and thus the same aggregate. + fake_group.all_gather_object.side_effect = lambda local: ["sig-pp0", "sig-pp1"] + + with mock.patch( + "sglang.srt.distributed.get_world_group", return_value=fake_group + ): + agg_from_rank0 = ( + PreshardedModelLoader._make_rank_invariant_structural_signature( + "sig-pp0" + ) + ) + agg_from_rank1 = ( + PreshardedModelLoader._make_rank_invariant_structural_signature( + "sig-pp1" + ) + ) + self.assertIsNotNone(agg_from_rank0) + self.assertEqual(agg_from_rank0, agg_from_rank1) + # Changing any rank's local contribution must change the aggregate. + fake_group.all_gather_object.side_effect = lambda local: [ + "sig-pp0", + "sig-pp1-changed", + ] + with mock.patch( + "sglang.srt.distributed.get_world_group", return_value=fake_group + ): + agg_changed = ( + PreshardedModelLoader._make_rank_invariant_structural_signature( + "sig-pp0" + ) + ) + self.assertNotEqual(agg_from_rank0, agg_changed) + + def test_compute_structural_signature_returns_none_on_failure(self): + # No self.load_config (and no real model class behind + # SimpleNamespace) means _get_quantization_config / _initialize_model + # will raise; this must be swallowed and return None, never raise, + # since a model class that can't be built on meta device must not + # break the overall load. + loader = object.__new__(PreshardedModelLoader) + model_config = SimpleNamespace(quantization=None) + self.assertIsNone(loader._compute_structural_signature(model_config)) + + def test_compute_structural_signature_picks_up_meta_model_shapes(self): + # End-to-end on a tiny real nn.Module standing in for the model + # class, to prove the meta-device construction + hashing wiring + # actually reflects shapes that depend on the live parallel state + # (here simulated via a width captured at construction time). + import torch.nn as nn + + class FakeModelLoader(PreshardedModelLoader): + def __init__(self, width): + self._width = width + self.load_config = SimpleNamespace() + + def _initialize_model_stub(model_config, load_config, quant_config): + return nn.Linear(4, model_config.width, bias=False) + + with mock.patch( + "sglang.srt.model_loader.loader._get_quantization_config", + return_value=None, + ), mock.patch( + "sglang.srt.model_loader.loader._initialize_model", + side_effect=_initialize_model_stub, + ), mock.patch( + "sglang.srt.model_loader.loader.set_default_torch_dtype", + return_value=mock.MagicMock(__enter__=mock.Mock(), __exit__=mock.Mock()), + ): + loader_narrow = FakeModelLoader(width=2) + loader_wide = FakeModelLoader(width=8) + model_config_narrow = SimpleNamespace( + quantization=None, dtype=torch.float32, width=2 + ) + model_config_wide = SimpleNamespace( + quantization=None, dtype=torch.float32, width=8 + ) + sig_narrow = loader_narrow._compute_structural_signature( + model_config_narrow + ) + sig_wide = loader_wide._compute_structural_signature(model_config_wide) + self.assertIsNotNone(sig_narrow) + self.assertIsNotNone(sig_wide) + self.assertNotEqual(sig_narrow, sig_wide) + + def test_meta_rope_cache_cleared_even_on_failure(self): + # If _initialize_model partially populates _ROPE_DICT with meta-device + # entries before raising, _compute_structural_signature must still + # clean them up (via finally), otherwise the real model init reuses + # the meta module and fails with "Cannot copy out of meta tensor". + import torch.nn as nn + + from sglang.srt.layers.rotary_embedding.factory import _ROPE_DICT + + # Plant a fake meta-device rotary module in the global cache. + fake_key = ("_test_meta_rope_cleanup_sentinel",) + with torch.device("meta"): + fake_module = nn.Linear(4, 4) + _ROPE_DICT[fake_key] = fake_module + + loader = object.__new__(PreshardedModelLoader) + loader.load_config = SimpleNamespace() + + with mock.patch( + "sglang.srt.model_loader.loader._get_quantization_config", + return_value=None, + ), mock.patch( + "sglang.srt.model_loader.loader._initialize_model", + side_effect=RuntimeError("simulated init failure"), + ), mock.patch( + "sglang.srt.model_loader.loader.set_default_torch_dtype", + return_value=mock.MagicMock(__enter__=mock.Mock(), __exit__=mock.Mock()), + ): + result = loader._compute_structural_signature( + SimpleNamespace(quantization=None, dtype=torch.float32) + ) + + self.assertIsNone(result) + self.assertNotIn(fake_key, _ROPE_DICT) + + +class TestShardConfig(unittest.TestCase): + """Bookkeeping guards for the enumerated cache-key fields and the + load-time match path. Dropping a field from `_collect_shard_config` or + breaking `_shard_config_matches` would silently collide caches; these + cases pin the failure modes the PR is meant to prevent.""" + + def _base_config(self, **overrides): + cfg = { + "tp": 8, + "dp": 1, + "ep": 1, + "pp": 1, + "moe_dense_tp_size": None, + "moe_dp_size": 1, + "enable_dp_lm_head": False, + "enable_fp32_lm_head": False, + "quantization": None, + "model_dtype": "torch.bfloat16", + "ep_num_redundant_experts": 0, + "enable_eplb": False, + "init_expert_location": "trivial", + "structural_signature": "deadbeef", + } + cfg.update(overrides) + return cfg + + def test_collect_shard_config_includes_required_keys(self): + # Registry completeness: dropping a key from the dict literal in + # `_collect_shard_config` is the exact failure mode that left + # moe_dense_tp_size / LM-head flags out of the cache key before. + loader = object.__new__(PreshardedModelLoader) + server_args = SimpleNamespace( + moe_dense_tp_size=1, + moe_dp_size=2, + enable_dp_lm_head=True, + enable_fp32_lm_head=True, + ep_num_redundant_experts=4, + enable_eplb=True, + init_expert_location="trivial", + ) + model_config = SimpleNamespace(quantization="fp8", dtype=torch.bfloat16) + required = { + "tp", + "dp", + "ep", + "pp", + "moe_dense_tp_size", + "moe_dp_size", + "enable_dp_lm_head", + "enable_fp32_lm_head", + "quantization", + "model_dtype", + "ep_num_redundant_experts", + "enable_eplb", + "init_expert_location", + "structural_signature", + } + parallel = SimpleNamespace(tp_size=8, moe_dp_size=2, moe_ep_size=4, pp_size=1) + with mock.patch( + "sglang.srt.model_loader.loader.get_server_args", + return_value=server_args, + ), mock.patch( + "sglang.srt.model_loader.loader.get_parallel", + return_value=parallel, + ), mock.patch.object( + loader, "_compute_structural_signature", return_value="sig16" + ): + cfg = loader._collect_shard_config(model_config) + self.assertEqual(required, set(cfg.keys())) + self.assertEqual(cfg["tp"], 8) + self.assertEqual(cfg["dp"], 2) + self.assertEqual(cfg["ep"], 4) + self.assertEqual(cfg["pp"], 1) + self.assertEqual(cfg["moe_dense_tp_size"], 1) + self.assertEqual(cfg["moe_dp_size"], 2) + self.assertTrue(cfg["enable_dp_lm_head"]) + self.assertTrue(cfg["enable_fp32_lm_head"]) + self.assertEqual(cfg["init_expert_location"], "trivial") + + def test_enumerated_fields_change_subfolder_hash(self): + # Each enumerated content/shape knob must feed the subfolder name. + # A silent drop from `_collect_shard_config` would leave this red. + loader = object.__new__(PreshardedModelLoader) + base = self._base_config() + base_name = loader._build_subfolder_name(base) + field_variants = { + "moe_dense_tp_size": 1, + "moe_dp_size": 2, + "enable_dp_lm_head": True, + "enable_fp32_lm_head": True, + "ep_num_redundant_experts": 8, + "enable_eplb": True, + "init_expert_location": "file:map.json:sha1:abcd", + "structural_signature": "cafebabe", + "quantization": "fp8", + } + for field, value in field_variants.items(): + with self.subTest(field=field): + other = self._base_config(**{field: value}) + other_name = loader._build_subfolder_name(other) + self.assertNotEqual( + base_name, + other_name, + f"changing {field} must change the cache subfolder name", + ) + + def test_shard_config_matches_equality_and_missing(self): + # Match must require exact stored equality; missing/mismatched + # shard_config is a cache miss (never raises) so reload re-dumps. + loader = object.__new__(PreshardedModelLoader) + cfg = self._base_config() + with tempfile.TemporaryDirectory() as tmp: + # No checksum.json → miss. + self.assertFalse(loader._shard_config_matches(tmp, cfg)) + # Matching stored config → hit. + with open( + os.path.join(tmp, PreshardedModelLoader.CHECKSUM_FILENAME), "w" + ) as f: + json.dump({"shard_config": cfg}, f) + self.assertTrue(loader._shard_config_matches(tmp, cfg)) + # One field differs → miss. + mismatched = self._base_config(moe_dense_tp_size=1) + self.assertFalse(loader._shard_config_matches(tmp, mismatched)) + # Plan without shard_config key → miss (upgrade path). + with open( + os.path.join(tmp, PreshardedModelLoader.CHECKSUM_FILENAME), "w" + ) as f: + json.dump({"version": 1}, f) + self.assertFalse(loader._shard_config_matches(tmp, cfg)) + + def test_init_expert_location_hashes_file_contents(self): + # Overwriting the same path with a different expert map must bust + # the cache; keying on the path alone would silently reuse weights. + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "experts.json") + with open(path, "w") as f: + json.dump({"logical_count": [[1, 0], [0, 1]]}, f) + key_a = PreshardedModelLoader._normalize_init_expert_location(path) + with open(path, "w") as f: + json.dump({"logical_count": [[0, 1], [1, 0]]}, f) + key_b = PreshardedModelLoader._normalize_init_expert_location(path) + self.assertNotEqual(key_a, key_b) + self.assertTrue(key_a.startswith("file:experts.json:sha1:")) + self.assertEqual( + PreshardedModelLoader._normalize_init_expert_location("trivial"), + "trivial", + ) + + def test_redump_clears_ready_before_rewrite(self): + # Config-mismatch re-dump into an already-ready dir must drop READY + # before mutating files, otherwise concurrent readers can observe + # READY + partial checksum/safetensors. + loader = object.__new__(PreshardedModelLoader) + loader._hash_num_threads = 1 + loader._max_file_bytes = 10**12 + + with tempfile.TemporaryDirectory() as tmp: + ready_path = os.path.join(tmp, PreshardedModelLoader.READY_FILENAME) + with open(ready_path, "w") as f: + f.write("{}") + self.assertTrue(os.path.isfile(ready_path)) + + # Force rank 0 / world 1 so the method runs the rank-0 prologue + # and then fails early on empty state (no need for full dump). + with mock.patch.object( + PreshardedModelLoader, + "_world_rank_and_size", + return_value=(0, 1), + ), mock.patch.object( + PreshardedModelLoader, "_world_barrier", return_value=None + ), mock.patch.object( + PreshardedModelLoader, + "_build_dump_plan", + return_value={ + "version": PreshardedModelLoader.PLAN_VERSION, + "files": [], + "rank_to_reads": {"0": []}, + "rank_checksums": {"0": "0"}, + }, + ), mock.patch.object( + PreshardedModelLoader, "_dump_files_for_rank", return_value=None + ): + loader._dump_state_to_disk( + state_dict={}, + extras={}, + presharded_dir=tmp, + shard_config=self._base_config(), + ) + # Dump rewrote READY at the end; the critical property is that + # the prologue unlinked the *previous* READY before rewriting + # checksum.json. Assert checksum was written and READY exists + # only as the fresh sentinel from this dump. + self.assertTrue(os.path.isfile(ready_path)) + with open(ready_path) as f: + sentinel = json.load(f) + self.assertIn("plan_version", sentinel) + with open(os.path.join(tmp, PreshardedModelLoader.CHECKSUM_FILENAME)) as f: + plan = json.load(f) + self.assertEqual(plan["shard_config"]["tp"], 8) + + +if __name__ == "__main__": + unittest.main()