meta-llama/llama · critical · AssertionError

Loading a checkpoint for MP={len(checkpoints)} but world siz

Error message

Loading a checkpoint for MP={len(checkpoints)} but world size is {model_parallel_size}

What it means

This assertion in Llama.build (llama/generation.py:103) requires model_parallel_size to exactly equal the number of .pth shards found in ckpt_dir. Each rank loads exactly one shard (checkpoints[get_model_parallel_rank()]), so a 7B checkpoint (1 shard) with model_parallel_size=2 (or a 70B checkpoint with 2 shards on a world size of 8) makes the mapping undefined, and the loader aborts instead of loading wrong/duplicate weights.

Source

Thrown at llama/generation.py:103

            torch.distributed.init_process_group("nccl")
        if not model_parallel_is_initialized():
            if model_parallel_size is None:
                model_parallel_size = int(os.environ.get("WORLD_SIZE", 1))
            initialize_model_parallel(model_parallel_size)

        local_rank = int(os.environ.get("LOCAL_RANK", 0))
        torch.cuda.set_device(local_rank)

        # seed must be the same in all processes
        torch.manual_seed(seed)

        if local_rank > 0:
            sys.stdout = open(os.devnull, "w")

        start_time = time.time()
        checkpoints = sorted(Path(ckpt_dir).glob("*.pth"))
        assert len(checkpoints) > 0, f"no checkpoint files found in {ckpt_dir}"
        assert model_parallel_size == len(
            checkpoints
        ), f"Loading a checkpoint for MP={len(checkpoints)} but world size is {model_parallel_size}"
        ckpt_path = checkpoints[get_model_parallel_rank()]
        checkpoint = torch.load(ckpt_path, map_location="cpu")
        with open(Path(ckpt_dir) / "params.json", "r") as f:
            params = json.loads(f.read())

        model_args: ModelArgs = ModelArgs(
            max_seq_len=max_seq_len,
            max_batch_size=max_batch_size,
            **params,
        )
        tokenizer = Tokenizer(model_path=tokenizer_path)
        model_args.vocab_size = tokenizer.n_words
        torch.set_default_tensor_type(torch.cuda.HalfTensor)
        model = Transformer(model_args)
        model.load_state_dict(checkpoint, strict=False)
        print(f"Loaded in {time.time() - start_time:.2f} seconds")

View on GitHub (pinned to 689c7f261b)

Solutions

1. Count the shards and align the world size: `ls ckpt_dir/*.pth | wc -l` must equal the nproc_per_node used to launch (and the model's true MP size)
2. If shards are missing, complete the download/copy so all consolidated.*.pth files are present, then relaunch with the matching process count
3. For single-GPU work, use a model whose checkpoint is 1 shard (7B/13B) or set nproc_per_node to the shard count with enough GPUs (each rank also needs the memory for its shard)
4. Verify model_parallel_size resolves to what you expect: when torch.distributed is initialized it defaults to WORLD_SIZE — print it before calling build

Example fix

# before
# 70B model (8 shards) launched with:
torchrun --nproc_per_node 4 example_chat_completion.py  # -> MP=8 != world 4? no: world=4, shards=8 -> assert

# after  (match processes to shard count, and to available GPUs)
torchrun --nproc_per_node 8 example_chat_completion.py
# or, single GPU with a 1-shard model:
python example_chat_completion.py  # 7B: 1 shard, model_parallel_size defaults to WORLD_SIZE=1

from pathlib import Path
import os

def mp_size_matches_shards(ckpt_dir: str, model_parallel_size: int) -> bool:
    return len(list(Path(ckpt_dir).glob("*.pth"))) == model_parallel_size

# if torch.distributed is up, WORLD_SIZE is what build will use:
mp = int(os.environ.get("WORLD_SIZE", 1))
assert mp_size_matches_shards(ckpt_dir, mp), f"{mp} ranks vs {len(list(Path(ckpt_dir).glob('*.pth')))} shards"

from pathlib import Path

def shards_match_world(ckpt_dir: str, world_size: int) -> bool:
    return len(list(Path(ckpt_dir).glob("*.pth"))) == world_size

try:
    llama = Llama.build(ckpt_dir=ckpt_dir, tokenizer_path=tok, max_seq_len=512, max_batch_size=8)
except AssertionError as e:
    if "world size" in str(e):
        n = len(list(Path(ckpt_dir).glob("*.pth")))
        raise RuntimeError(f"relaunch with --nproc_per_node={n} to match {n} shards") from e
    raise

When it happens

Trigger: Llama.build(...) where model_parallel_size (defaulting to int(os.environ['WORLD_SIZE']) via torch.distributed when initialized) differs from the count of consolidated.*.pth files. Examples: running torchrun with 2 processes against the 7B model (1 shard); running single-process against a 13B/70B download whose shards were only partially copied (e.g. 6 of 8 shards present, world size 8).

Common situations: - Launching with `torchrun --nproc_per_node N` where N doesn't match the model's shard count (7B/13B = 1 shard; 70B = 8 shards). - Partial download of a sharded checkpoint: some consolidated.*.pth copied, so len(checkpoints) < the model's true MP size. - Passing model_parallel_size explicitly while a torch.distributed world is already up (or vice versa), so the value used is not the one you think. - Copying only a subset of shards to a smaller node to 'save disk'.

Related errors

• no checkpoint files found in {ckpt_dir}
• Error: special tags are not allowed as part of the prompt.
• model only supports 'system', 'user' and 'assistant' roles,
• Last message must be from user, got {dialog[-1]['role']}

AI-assisted analysis of meta-llama/llama@689c7f261b (2026-08-15). Data as JSON: /api/errors/412414c0cbcb9139. Report an issue: GitHub.