inference/convert.py

"""Convert the original Step-5 checkpoint into validated TP shards.

The converter writes ``OUT/tpN/model-r{rank}.safetensors`` and matching
``config-r{rank}.json`` files, plus one shared ``OUT/tokenizer_files`` directory.
The only supported production layout is ``gqa-provider-shared-tp-v2``.

Example::

    python3 convert.py --checkpoint /path/to/Step-5 \
        --out-dir /path/to/Step-5-TP --tp-size 8 --ep-size 8
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
from typing import Any, Callable, TYPE_CHECKING

import torch
from safetensors import safe_open
from safetensors.torch import save_file

if TYPE_CHECKING:
    from model import Step5Config


STEP5_TP_LAYOUT_VERSION = "gqa-provider-shared-tp-v2"


# Each rule callable maps (full_tensor, layer_idx_or_None, tp_size, rank, config) to the
# per-rank tensor. Rules are matched by attempting them in order; the first hit wins.
Rule = Callable[[torch.Tensor, int | None, int, int, Any], torch.Tensor]


def _slice_dim(dim: int) -> Rule:
    """Column or Row: slice tensor along ``dim`` by rank's window, then clone."""

    def apply(t, _layer, tp, rank, _config):
        size = t.shape[dim]
        assert size % tp == 0, f"{size} not divisible by tp={tp}"
        per = size // tp
        return t.narrow(dim, rank * per, per).clone()

    return apply


def _group_window(num_groups: int, tp: int, rank: int) -> tuple[int, int]:
    """Return ``(first_group, count)`` consumed by rank's contiguous Q-head slice.

    For TP below the group count, groups partition across ranks. For TP above it, one group
    is duplicated across ``tp / num_groups`` adjacent ranks. The latter is the production
    TP=8 case: ``_group_window(4, 8, r) == (r // 2, 1)``.
    """
    if not 0 <= rank < tp:
        raise ValueError(f"rank {rank} outside tp_size={tp}")
    if num_groups % tp == 0:
        count = num_groups // tp
        return rank * count, count
    if tp % num_groups == 0:
        replicas = tp // num_groups
        return rank // replicas, 1
    raise ValueError(
        f"tp_size={tp} and group count {num_groups} are incompatible: "
        "one must divide the other"
    )


def _slice_grouped_dim(dim: int, config_attr: str) -> Rule:
    """Slice whole GQA/provider groups, duplicating a group when TP exceeds group count."""

    def apply(t, _layer, tp, rank, config):
        owner = config
        for component in config_attr.split("."):
            owner = getattr(owner, component)
        num_groups = int(owner)
        size = t.shape[dim]
        if size % num_groups:
            raise ValueError(
                f"tensor dim {dim} size {size} is not divisible by {num_groups} groups"
            )
        group_width = size // num_groups
        first, count = _group_window(num_groups, tp, rank)
        return t.narrow(dim, first * group_width, count * group_width).clone()

    return apply


def _replicate(t, _layer, _tp, _rank, _config):
    return t.clone()


# First matching key substring selects the rule. Attention and dense/shared
# experts use TP dimension slices; routed experts use EP slices of the expert
# axis. Routers, normalization parameters and vision weights are replicated.
_SHARD_TABLE: list[tuple[str, Rule]] = [
    # Attention: q_proj/g_proj ColumnParallel (slice head dim, dim 0); o_proj RowParallel
    # (slice head dim on dim 1 of [hidden, num_heads*head_dim] = the in dim). K/V follow
    # the GQA group consumed by the local Q slice. At TP=8 each rank gets one 192-d head,
    # duplicated across adjacent rank pairs rather than replicating all four heads.
    (".self_attn.q_proj.weight", _slice_dim(0)),
    (".self_attn.g_proj.weight", _slice_dim(0)),
    (".self_attn.o_proj.weight", _slice_dim(1)),
    (
        ".self_attn.k_proj.weight",
        _slice_grouped_dim(0, "num_attention_groups"),
    ),
    (
        ".self_attn.v_proj.weight",
        _slice_grouped_dim(0, "num_attention_groups"),
    ),
    # Dense MLP: gate/up ColumnParallel (slice intermediate, dim 0); down RowParallel
    # (slice intermediate on dim 1 of [hidden, intermediate] = the in dim).
    (".mlp.gate_proj.weight", _slice_dim(0)),
    (".mlp.up_proj.weight", _slice_dim(0)),
    (".mlp.down_proj.weight", _slice_dim(1)),
    # MoE routed experts -- EP, not TP-of-inner-dim. Each rank owns n_experts // tp_size
    # contiguous experts (full inner dim per expert). Slicing the experts dim on dim 0. Block
    # scales follow their weight onto the experts dim. ``moe.gate_proj.weight`` matches
    # ``model.layers.N.moe.gate_proj.weight``; the leading ``.`` excludes ``share_expert`` keys.
    (".moe.gate_proj.weight", _slice_dim(0)),
    (".moe.up_proj.weight", _slice_dim(0)),
    (".moe.down_proj.weight", _slice_dim(0)),
    (".moe.gate_proj.weight_scale_inv", _slice_dim(0)),
    (".moe.up_proj.weight_scale_inv", _slice_dim(0)),
    (".moe.down_proj.weight_scale_inv", _slice_dim(0)),
    # ``moe.gate.weight`` (the router) and ``moe.router_bias`` are replicated so every rank
    # selects over all experts and sees the same router scores.
    (".moe.gate.weight", _replicate),
    (".moe.router_bias", _replicate),
    # Shared expert is ordinary TP, not a replicate: gate/up slice their output dim and
    # down slices its input dim. The runtime deliberately leaves the down result local,
    # combines it with the local EP-routed result in FP32, then all-reduces once.
    (".share_expert.gate_proj.weight", _slice_dim(0)),
    (".share_expert.up_proj.weight", _slice_dim(0)),
    (".share_expert.down_proj.weight", _slice_dim(1)),
    # The indexer's Q/W heads are grouped by DSA provider group. Provider groups align with
    # GQA KV groups, so TP=8 gives every rank one group (four Q/W heads), duplicated across
    # the same adjacent rank pair. K/Z are single-head MQA projections and remain replicated.
    # These exact matches must precede the generic sparse-indexer replicate rule.
    (
        ".sparse_indexer_q.weight",
        _slice_grouped_dim(0, "sparse_config.num_provider_groups"),
    ),
    (
        ".sparse_indexer_w.weight",
        _slice_grouped_dim(0, "sparse_config.num_provider_groups"),
    ),
    (".sparse_indexer_", _replicate),
    (".ssmax_s", _replicate),
    # Attention + indexer RMSNorms: each is a small per-head scale on the routed tensor
    # (the residual stream is full, so the norm acts on the full vector). Replicated.
    (".q_norm.weight", _replicate),
    (".k_norm.weight", _replicate),
    (".k_norm.bias", _replicate),
    # Per-layer input/post-attention RMSNorms and the model's final RMSNorm: replicated.
    (".input_layernorm.weight", _replicate),
    (".post_attention_layernorm.weight", _replicate),
    # Embedding, lm_head: replicated (vocab is split at runtime by VocabParallelLinear).
    ("model.embed_tokens.weight", _replicate),
    ("model.norm.weight", _replicate),
    ("lm_head.weight", _replicate),
]


def _shard_for(key: str) -> Rule:
    if key.startswith(("vision_model.", "vit_large_projector.")):
        return _replicate
    for substr, rule in _SHARD_TABLE:
        if substr in key:
            return rule
    raise KeyError(f"no sharding rule for {key}")


def _layer_idx_of(key: str) -> int | None:
    """``model.layers.N.`` -- extract N if present and in-range, else None."""
    marker = "model.layers."
    if key.startswith(marker):
        rest = key[len(marker):]
        layer_str = rest.split(".", 1)[0]
        try:
            return int(layer_str)
        except ValueError:
            return None
    return None


def _order_for_inspection(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    """Put embed / lm_head / norm first, then layers, then any stragglers."""
    ordered: dict[str, torch.Tensor] = {}
    for key in ("model.embed_tokens.weight", "model.norm.weight", "lm_head.weight"):
        if key in state:
            ordered[key] = state[key]
    for key in sorted(state.keys()):
        if key.startswith("model.layers."):
            ordered[key] = state[key]
    for key in sorted(state.keys()):
        if key not in ordered:
            ordered[key] = state[key]
    return ordered


def _load_config_for(checkpoint: str, tp_size: int) -> Step5Config:
    from model import config_from_dict

    with open(os.path.join(checkpoint, "config.json")) as handle:
        raw = json.load(handle)
    raw["tp_size"] = tp_size
    raw["expert_parallel_size"] = tp_size
    raw["tp_layout_version"] = STEP5_TP_LAYOUT_VERSION
    return config_from_dict(raw)


def _write_config(
    existing_config_path: str,
    out_path: str,
    tp_size: int,
    expert_parallel_size: int,
) -> None:
    with open(existing_config_path) as handle:
        cfg = json.load(handle)
    cfg["tp_size"] = tp_size
    cfg["expert_parallel_size"] = expert_parallel_size
    cfg["tp_layout_version"] = STEP5_TP_LAYOUT_VERSION
    cfg.pop("use_triton_sliding_attention", None)
    from model import Step5VLConfig, config_from_dict
    config = config_from_dict(cfg)
    cfg = config.to_dict()
    if isinstance(config, Step5VLConfig):
        cfg["architectures"] = ["Step5ForConditionalGeneration"]
        cfg["auto_map"] = {
            "AutoConfig": "model.Step5VLConfig",
            "AutoModelForCausalLM": "model.Step5ForConditionalGeneration",
            "AutoModelForImageTextToText": "model.Step5ForConditionalGeneration",
        }
    else:
        cfg["architectures"] = ["Step5ForCausalLM"]
        cfg["auto_map"] = {
            "AutoConfig": "model.Step5Config",
            "AutoModelForCausalLM": "model.Step5ForCausalLM",
        }
    # ``num_attention_heads`` stays at the full value -- the model derives the per-rank view via
    # ``num_attention_heads // tp_size``. ``tp_layout_version`` deliberately makes shards
    # produced by either the old replicated-KV/provider rule or the v1 replicated-shared
    # rule fail closed instead of loading with plausible but wrong shapes/semantics.
    with open(out_path, "w") as handle:
        json.dump(cfg, handle, indent=2)
        handle.write("\n")


def _write_readme(
    out_path: str,
    tp_size: int,
    expert_parallel_size: int,
    checkpoint: str,
    config: Step5Config,
) -> None:
    local_kv = config.num_kv_groups_per_rank
    local_index_heads = config.sparse_indexer_num_heads_per_rank
    text = f"""# step5 TP-{tp_size} shards

This directory holds one `model-r{{0..{tp_size-1}}}.safetensors` and one
`config-r{{0..{tp_size-1}}}.json` file per rank. The sharding matches the tensor-parallel convention in
`model.py`. Layout marker: `{STEP5_TP_LAYOUT_VERSION}`. Shards without this
marker use an obsolete replicated-KV/provider or replicated-shared layout and must not be
resumed or loaded.

The same {tp_size} ranks form the dense tensor-parallel group and the routed-expert
placement group (`expert_parallel_size={expert_parallel_size}`). Independent TP and EP
process-group axes are not implemented.

| Tensor | Rule | Per-rank shape (tp={tp_size}) |
|--------|------|-------------------------------|
| `self_attn.q_proj.weight` | ColumnParallel, slice dim 0 by `num_heads / tp_size` | `[{64 // tp_size} * 192, 4096]` |
| `self_attn.g_proj.weight` | ColumnParallel, slice dim 0 by `num_heads / tp_size` | `[{64 // tp_size}, 4096]` |
| `self_attn.k_proj.weight` | GQA group slice; duplicate only for ranks consuming the same group | `[{local_kv} * 192, 4096]` |
| `self_attn.v_proj.weight` | GQA group slice; duplicate only for ranks consuming the same group | `[{local_kv} * 192, 4096]` |
| `self_attn.o_proj.weight` | RowParallel, slice dim 1 by `num_heads * head_dim / tp_size` | `[4096, {64 // tp_size} * 192]` |
| `mlp.gate_proj.weight` | ColumnParallel, slice dim 0 by `intermediate / tp_size` | `[{13824 // tp_size}, 4096]` |
| `mlp.up_proj.weight` | ColumnParallel, slice dim 0 by `intermediate / tp_size` | `[{13824 // tp_size}, 4096]` |
| `mlp.down_proj.weight` | RowParallel, slice dim 1 by `intermediate / tp_size` | `[4096, {13824 // tp_size}]` |
| `moe.gate.weight` | replicated (router, every rank selects over all experts) | `[352, 4096]` |
| `moe.gate_proj.weight` | EP, slice dim 0 (experts) by `n_experts / tp_size` | `[{352 // tp_size}, 1536, 4096]` |
| `moe.up_proj.weight` | EP, slice dim 0 (experts) by `n_experts / tp_size` | `[{352 // tp_size}, 1536, 4096]` |
| `moe.down_proj.weight` | EP, slice dim 0 (experts) by `n_experts / tp_size` | `[{352 // tp_size}, 4096, 1536]` |
| `moe.*.weight_scale_inv` | slice the experts dim (dim 0), inner block grid unchanged | `[{352 // tp_size}, ceil(out/128), ceil(in/128)]` |
| `moe.router_bias` | replicated | `[352]` |
| `share_expert.gate/up_proj.weight` | ColumnParallel, slice dim 0 by `share_expert_dim / tp_size` | `[{1536 // tp_size}, 4096]` |
| `share_expert.down_proj.weight` | RowParallel, slice dim 1; reduce deferred until shared+routed FP32 combine | `[4096, {1536 // tp_size}]` |
| `sparse_indexer_q.weight` | provider-group slice (4 heads/group) | `[{local_index_heads} * 256, 4096]` |
| `sparse_indexer_w.weight` | provider-group slice (4 heads/group) | `[{local_index_heads}, 4096]` |
| `sparse_indexer_k/z`, `ssmax_s`, `*_norm.weight`, `k_norm.bias` | replicated | full |
| `model.embed_tokens.weight` | replicated | `[128896, 4096]` |
| `lm_head.weight` | slices the vocab dim at runtime (VocabParallelLinear) | `[128896, 4096]` full per rank |
| `model.norm.weight` | replicated | `[4096]` |

Reference geometry: `num_attention_heads={64}`, `num_attention_groups={4}`, `head_dim={192}`,
`hidden={4096}`, `intermediate={13824}`, `moe_intermediate={1536}`, `vocab={128896}`,
`moe_num_experts={352}`, `topk={8}`. Produced from `{checkpoint}` by
`inference/convert.py`.
"""
    with open(out_path, "w") as handle:
        handle.write(text)


_TOKENIZER_FILES = (
    "tokenizer.json",
    "tokenizer_config.json",
    "special_tokens_map.json",
    "chat_template.jinja",
    "vocab.json",
    "merges.txt",
    "added_tokens.json",
    "preprocessor_config.json",
    "processor_config.json",
)


def _build_rank_state(
    checkpoint: str,
    keys: list[str],
    weight_map: dict[str, str],
    tp_size: int,
    rank: int,
    config: Step5Config,
) -> dict[str, torch.Tensor]:
    state: dict[str, torch.Tensor] = {}
    chunks_by_shard: dict[str, list[str]] = {}
    for key in keys:
        chunks_by_shard.setdefault(weight_map[key], []).append(key)
    for shard in sorted(chunks_by_shard):
        print(f"  rank {rank}: shard {shard}", flush=True)
        with safe_open(os.path.join(checkpoint, shard), framework="pt", device="cpu") as handle:
            for key in chunks_by_shard[shard]:
                full = handle.get_tensor(key)
                rule = _shard_for(key)
                state[key] = rule(full, _layer_idx_of(key), tp_size, rank, config)
    return state


def validate_rank_shard(path, config, rank, *, expected=None) -> None:
    """Reject incompatible or incomplete shards before HF can initialize missing weights.

    Inspect headers only. In particular, text-only v2 shards must not load under
    a VL config just because their parallel-layout metadata is still current.
    """
    from model import model_class_for_config

    if expected is None:
        with torch.device("meta"):
            expected = model_class_for_config(config)(config).state_dict()
    text_config = config.get_text_config()
    required_metadata = {
        "tp_size": str(text_config.tp_size),
        "expert_parallel_size": str(getattr(text_config, "expert_parallel_size", text_config.tp_size)),
        "rank": str(rank),
    }
    with safe_open(str(path), framework="pt", device="cpu") as handle:
        metadata = handle.metadata() or {}
        # Tensor layout is unchanged; accept either serialized spelling, but
        # reject conflicting markers rather than silently choosing one.
        layout_fields = [name for name in ("step5_tp_layout", "step4_tp_layout") if name in metadata]
        if not layout_fields:
            raise ValueError(f"{path}: missing step5_tp_layout metadata")
        for field in layout_fields:
            if metadata[field] != STEP5_TP_LAYOUT_VERSION:
                raise ValueError(
                    f"{path}: {field} metadata is {metadata[field]!r}, "
                    f"expected {STEP5_TP_LAYOUT_VERSION!r}"
                )
        for field, value in required_metadata.items():
            if metadata.get(field) != value:
                raise ValueError(f"{path}: {field} metadata is {metadata.get(field)!r}, expected {value!r}")
        keys = set(handle.keys())
        missing, unexpected = set(expected) - keys, keys - set(expected)
        if missing or unexpected:
            raise ValueError(
                f"{path}: shard key mismatch: missing={sorted(missing)[:8]}, "
                f"unexpected={sorted(unexpected)[:8]}; regenerate complete TP shards"
            )
        for key, tensor in expected.items():
            shape = tuple(handle.get_slice(key).get_shape())
            if shape != tuple(tensor.shape):
                raise ValueError(f"{path}: {key} shape {shape} != model {tuple(tensor.shape)}")


def _validate_resume_shards(tp_dir, start_rank, *, config, expected) -> None:
    """Validate retained ranks with the same contract as the runtime loader."""
    for rank in range(start_rank):
        path = os.path.join(tp_dir, f"model-r{rank}.safetensors")
        if not os.path.isfile(path):
            raise FileNotFoundError(
                f"cannot resume at rank {start_rank}: missing completed shard {path}"
            )
        validate_rank_shard(path, config, rank, expected=expected)


def pre_shard(
    checkpoint: str,
    out_dir: str,
    tp_size: int,
    *,
    expert_parallel_size: int | None = None,
    dry_run: bool = False,
    start_rank: int = 0,
    check_only: bool = False,
) -> int:
    """Write ``tp_size`` per-rank safetensors (plus configs and a README) to ``out_dir``.

    ``start_rank`` resumes a partial conversion after checking earlier ranks for
    matching layout, topology and complete parameter keys/shapes.
    """
    expert_parallel_size = (
        tp_size if expert_parallel_size is None else expert_parallel_size
    )
    if tp_size < 1 or expert_parallel_size < 1:
        raise ValueError(
            f"parallel sizes must be positive, got TP={tp_size}, "
            f"EP={expert_parallel_size}"
        )
    if expert_parallel_size != tp_size:
        raise ValueError(
            "independent TP and EP groups are not implemented: "
            f"expert_parallel_size={expert_parallel_size} must equal "
            f"tp_size={tp_size}. The same ranks shard the dense backbone and "
            "the routed experts."
        )
    if not 0 <= start_rank < tp_size:
        raise ValueError(
            f"start_rank must be in [0, {tp_size}), got {start_rank}"
        )
    index_path = os.path.join(checkpoint, "model.safetensors.index.json")
    with open(index_path) as handle:
        weight_map = json.load(handle)["weight_map"]

    config = _load_config_for(checkpoint, tp_size)
    text_config = config.get_text_config()
    keys = sorted(k for k in weight_map if (
        _layer_idx_of(k) is None or _layer_idx_of(k) < text_config.num_hidden_layers
    ))
    from model import model_class_for_config

    with torch.device("meta"):
        expected = model_class_for_config(config)(config).state_dict()
    missing = set(expected) - set(keys)
    unexpected = set(keys) - set(expected)
    if missing or unexpected:
        raise ValueError(
            f"checkpoint key mismatch: missing={sorted(missing)[:8]}, "
            f"unexpected={sorted(unexpected)[:8]}"
        )
    # Check every shape before reading weight data or creating outputs. This also
    # detects vision geometry/configuration drift on the original checkpoint.
    by_shard = {}
    for key in keys:
        by_shard.setdefault(weight_map[key], []).append(key)
    for shard, shard_keys in sorted(by_shard.items()):
        with safe_open(os.path.join(checkpoint, shard), framework="pt", device="cpu") as handle:
            for key in shard_keys:
                shape = handle.get_slice(key).get_shape()
                full = torch.empty(shape, device="meta")
                actual = _shard_for(key)(full, _layer_idx_of(key), tp_size, 0, text_config)
                if actual.shape != expected[key].shape:
                    raise ValueError(f"{key}: shard shape {tuple(actual.shape)} != model {tuple(expected[key].shape)}")
    print(f"validated {len(keys)} tensors; excluded {len(weight_map) - len(keys)} MTP tensors", flush=True)
    if check_only:
        print("PASS (checkpoint keys and TP shapes; no weight data read or files written)")
        return 0

    tp_dir = os.path.join(out_dir, f"tp{tp_size}")
    os.makedirs(tp_dir, exist_ok=True)
    if start_rank:
        _validate_resume_shards(
            tp_dir,
            start_rank,
            config=config,
            expected=expected,
        )
    tokenizer_dst = os.path.join(out_dir, "tokenizer_files")
    if not dry_run and not os.path.isdir(tokenizer_dst):
        os.makedirs(tokenizer_dst, exist_ok=True)
        for name in _TOKENIZER_FILES:
            src = os.path.join(checkpoint, name)
            if os.path.exists(src):
                shutil.copy(src, os.path.join(tokenizer_dst, name))

    sample_keys = [
        "model.embed_tokens.weight",
        "model.norm.weight",
        "lm_head.weight",
        "model.layers.0.self_attn.q_proj.weight",
        "model.layers.0.self_attn.k_proj.weight",
        "model.layers.0.self_attn.o_proj.weight",
        "model.layers.0.mlp.gate_proj.weight",
        "model.layers.0.mlp.down_proj.weight",
        "model.layers.10.moe.gate_proj.weight",
        "model.layers.10.moe.gate_proj.weight_scale_inv",
        "model.layers.10.moe.down_proj.weight",
        "model.layers.10.moe.down_proj.weight_scale_inv",
        "model.layers.3.self_attn.sparse_indexer_q.weight",
        "model.layers.3.self_attn.sparse_indexer_w.weight",
    ]
    for rank in range(start_rank, tp_size):
        print(f"=== rank {rank} ===", flush=True)
        state = _build_rank_state(checkpoint, keys, weight_map, tp_size, rank, text_config)
        if dry_run:
            print(f"rank {rank}:")
            for key in sample_keys:
                if key in state:
                    print(f"  {key}: {tuple(state[key].shape)} {state[key].dtype}")
        else:
            ordered = _order_for_inspection(state)
            path = os.path.join(tp_dir, f"model-r{rank}.safetensors")
            save_file(
                ordered,
                path,
                metadata={
                    "format": "pt",
                    "format_version": "1",
                    "step5_tp_layout": STEP5_TP_LAYOUT_VERSION,
                    "tp_size": str(tp_size),
                    "expert_parallel_size": str(expert_parallel_size),
                    "rank": str(rank),
                },
            )
            print(f"  wrote {path}  ({os.path.getsize(path) / 1e9:.2f} GB)")

    if dry_run:
        print(f"(dry_run: no files written, would produce tp{tp_size}/ in {out_dir})")
        return 0

    config_path = os.path.join(tp_dir, "config-r0.json")
    if start_rank > 0:
        print(f"resuming from rank {start_rank} (ranks 0..{start_rank - 1} skipped)", flush=True)

    _write_config(
        os.path.join(checkpoint, "config.json"),
        config_path,
        tp_size,
        expert_parallel_size,
    )
    for rank in range(tp_size):
        dst = os.path.join(tp_dir, f"config-r{rank}.json")
        if dst != config_path:
            shutil.copy(config_path, dst)
    _write_readme(
        os.path.join(tp_dir, "README.md"),
        tp_size,
        expert_parallel_size,
        checkpoint,
        text_config,
    )
    print(f"  wrote {len(_TOKENIZER_FILES)} tokenizer files to {tokenizer_dst}")
    print(f"  wrote {tp_size} per-rank configs to {tp_dir}/config-r{{0..{tp_size-1}}}.json")
    print("PASS")
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--out-dir", "--out_dir", dest="out_dir", default=None)
    parser.add_argument("--tp-size", "--tp_size", dest="tp_size", type=int, default=8)
    parser.add_argument(
        "--ep-size",
        "--expert-parallel-size",
        dest="ep_size",
        type=int,
        default=8,
        help="Routed-expert sharding degree (default: 8). The current "
        "co-located topology requires EP size to equal TP size.",
    )
    parser.add_argument("--dry-run", "--dry_run", dest="dry_run", action="store_true")
    parser.add_argument(
        "--start-rank",
        "--start_rank",
        dest="start_rank",
        type=int,
        default=0,
        help="Resume from this rank; ranks 0..start_rank-1 are assumed already on disk.",
    )
    parser.add_argument("--check-only", action="store_true", help="Validate all checkpoint keys and TP shapes without reading tensor data or writing shards")
    args = parser.parse_args()
    if not args.check_only and args.out_dir is None:
        parser.error("--out-dir is required unless --check-only is set")
    return pre_shard(
        checkpoint=args.checkpoint,
        out_dir=args.out_dir,
        tp_size=args.tp_size,
        expert_parallel_size=args.ep_size,
        dry_run=args.dry_run,
        start_rank=args.start_rank,
        check_only=args.check_only,
    )


if __name__ == "__main__":
    raise SystemExit(main())