{"record":{"id":"412414c0cbcb9139","repo":"meta-llama/llama","slug":"loading-a-checkpoint-for-mp-len-checkpoints-but","errorCode":null,"errorMessage":"Loading a checkpoint for MP={len(checkpoints)} but world size is {model_parallel_size}","messagePattern":"Loading a checkpoint for MP=(.+?) but world size is (.+?)","errorType":"validation","errorClass":"AssertionError","httpStatus":null,"severity":"critical","filePath":"llama/generation.py","lineNumber":103,"sourceCode":"            torch.distributed.init_process_group(\"nccl\")\n        if not model_parallel_is_initialized():\n            if model_parallel_size is None:\n                model_parallel_size = int(os.environ.get(\"WORLD_SIZE\", 1))\n            initialize_model_parallel(model_parallel_size)\n\n        local_rank = int(os.environ.get(\"LOCAL_RANK\", 0))\n        torch.cuda.set_device(local_rank)\n\n        # seed must be the same in all processes\n        torch.manual_seed(seed)\n\n        if local_rank > 0:\n            sys.stdout = open(os.devnull, \"w\")\n\n        start_time = time.time()\n        checkpoints = sorted(Path(ckpt_dir).glob(\"*.pth\"))\n        assert len(checkpoints) > 0, f\"no checkpoint files found in {ckpt_dir}\"\n        assert model_parallel_size == len(\n            checkpoints\n        ), f\"Loading a checkpoint for MP={len(checkpoints)} but world size is {model_parallel_size}\"\n        ckpt_path = checkpoints[get_model_parallel_rank()]\n        checkpoint = torch.load(ckpt_path, map_location=\"cpu\")\n        with open(Path(ckpt_dir) / \"params.json\", \"r\") as f:\n            params = json.loads(f.read())\n\n        model_args: ModelArgs = ModelArgs(\n            max_seq_len=max_seq_len,\n            max_batch_size=max_batch_size,\n            **params,\n        )\n        tokenizer = Tokenizer(model_path=tokenizer_path)\n        model_args.vocab_size = tokenizer.n_words\n        torch.set_default_tensor_type(torch.cuda.HalfTensor)\n        model = Transformer(model_args)\n        model.load_state_dict(checkpoint, strict=False)\n        print(f\"Loaded in {time.time() - start_time:.2f} seconds\")","sourceCodeStart":85,"sourceCodeEnd":121,"githubUrl":"https://github.com/meta-llama/llama/blob/689c7f261b9c5514636ecc3c5fefefcbb3e6eed7/llama/generation.py#L85-L121","documentation":"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.","triggerScenarios":"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).","commonSituations":"- Launching with `torchrun --nproc_per_node N` where N doesn't match the model's shard count (7B/13B = 1 shard; 70B = 8 shards).\n- Partial download of a sharded checkpoint: some consolidated.*.pth copied, so len(checkpoints) < the model's true MP size.\n- 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.\n- Copying only a subset of shards to a smaller node to 'save disk'.","solutions":["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)","If shards are missing, complete the download/copy so all consolidated.*.pth files are present, then relaunch with the matching process count","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)","Verify model_parallel_size resolves to what you expect: when torch.distributed is initialized it defaults to WORLD_SIZE — print it before calling build"],"exampleFix":"# before\n# 70B model (8 shards) launched with:\ntorchrun --nproc_per_node 4 example_chat_completion.py  # -> MP=8 != world 4? no: world=4, shards=8 -> assert\n\n# after  (match processes to shard count, and to available GPUs)\ntorchrun --nproc_per_node 8 example_chat_completion.py\n# or, single GPU with a 1-shard model:\npython example_chat_completion.py  # 7B: 1 shard, model_parallel_size defaults to WORLD_SIZE=1","handlingStrategy":"validation","validationCode":"from pathlib import Path\nimport os\n\ndef mp_size_matches_shards(ckpt_dir: str, model_parallel_size: int) -> bool:\n    return len(list(Path(ckpt_dir).glob(\"*.pth\"))) == model_parallel_size\n\n# if torch.distributed is up, WORLD_SIZE is what build will use:\nmp = int(os.environ.get(\"WORLD_SIZE\", 1))\nassert mp_size_matches_shards(ckpt_dir, mp), f\"{mp} ranks vs {len(list(Path(ckpt_dir).glob('*.pth')))} shards\"","typeGuard":"from pathlib import Path\n\ndef shards_match_world(ckpt_dir: str, world_size: int) -> bool:\n    return len(list(Path(ckpt_dir).glob(\"*.pth\"))) == world_size","tryCatchPattern":"try:\n    llama = Llama.build(ckpt_dir=ckpt_dir, tokenizer_path=tok, max_seq_len=512, max_batch_size=8)\nexcept AssertionError as e:\n    if \"world size\" in str(e):\n        n = len(list(Path(ckpt_dir).glob(\"*.pth\")))\n        raise RuntimeError(f\"relaunch with --nproc_per_node={n} to match {n} shards\") from e\n    raise","preventionTips":["Derive nproc_per_node from the shard count in your launch script instead of hardcoding it","Ensure the download is complete: shard count is part of the checkpoint's contract (7B/13B=1, 70B=8)","Log the resolved model_parallel_size (WORLD_SIZE when distributed) before build so mismatches are obvious"],"tags":["llama","model-parallel","checkpoint","torchrun","distributed"],"backgroundTag":null,"analyzedSha":"689c7f261b9c5514636ecc3c5fefefcbb3e6eed7","analyzedAt":"2026-08-15T02:36:29.698Z","schemaVersion":2},"datasetVersion":"2026-08-15T22:17:37.221Z"}