# How to Resolve the "Model Parallel Size Does Not Match Checkpoint Files" Error in LLaMA

> Resolve the model parallel size does not match checkpoint files error in LLaMA by matching WORLD_SIZE to your checkpoint shards. Learn how to properly launch inference with torchrun.

- Repository: [Meta Llama/llama](https://github.com/meta-llama/llama)
- Tags: how-to-guide
- Published: 2026-03-05

---

**To fix the "model parallel size does not match checkpoint files" error, ensure your `WORLD_SIZE` environment variable or `model_parallel_size` parameter equals the exact number of `.pth` checkpoint shards in your directory, then launch inference with `torchrun --nproc_per_node=N`.**

When running inference with Meta's LLaMA models, the framework validates that the number of parallel processes matches the number of checkpoint shards on disk. This validation occurs in [`llama/generation.py`](https://github.com/meta-llama/llama/blob/main/llama/generation.py) and raises an assertion error when the counts diverge. Understanding how the checkpoint loading logic works will help you quickly align your runtime configuration with your model artifacts.

## Understanding the Checkpoint Validation Logic

The error originates in **[`llama/generation.py`](https://github.com/meta-llama/llama/blob/main/llama/generation.py)**, where the library initializes the model parallel group and loads checkpoints. The code performs the following validation:

1. Determines the world size from either the `model_parallel_size` argument or the `WORLD_SIZE` environment variable.
2. Collects all checkpoint files (`*.pth`) from the specified `ckpt_dir` and sorts them.
3. Asserts that `model_parallel_size == len(checkpoints)`.

If this assertion fails, the runtime raises:

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

```

This strict one-to-one mapping exists because each parallel process loads exactly one shard containing a portion of the model weights.

## Common Causes of the Mismatch

Several configuration errors trigger this assertion failure:

- **Incorrect `WORLD_SIZE`**: Setting the environment variable to a value different from the number of `.pth` files in your checkpoint directory. For example, exporting `WORLD_SIZE=8` when only four shard files exist.
- **Incomplete or Corrupted Downloads**: Missing shard files or extra stray files (such as older checkpoints or temporary downloads) in the `ckpt_dir`.
- **Single-GPU Inference on Sharded Checkpoints**: Attempting to load a model-parallel checkpoint with a single process when the weights were saved using multiple GPUs.
- **Mismatched Parallelism Configuration**: Using a checkpoint sharded with a different world size than your current runtime configuration.

## Step-by-Step Resolution

Follow these steps to align your runtime configuration with your checkpoint shards:

1. **Count your checkpoint shards**

   ```bash
   ls /path/to/checkpoint/*.pth | wc -l
   ```

   Record this number as `N`.

2. **Set the world size environment variable**

   ```bash
   export WORLD_SIZE=$N
   ```

   Alternatively, pass `model_parallel_size=$N` directly to the generation function.

3. **Launch with the correct number of processes**

   ```bash
   torchrun --nproc_per_node=$N generate.py \
       --ckpt-dir /path/to/checkpoint \
       --tokenizer-path /path/to/tokenizer.model \
       --prompt "Your prompt here"
   ```

4. **Verify no extraneous files exist**

   Ensure your checkpoint directory contains exactly `N` `.pth` files and no other files matching that extension.

If you need to consolidate multiple shards into a single checkpoint for single-GPU inference, use the utility script [`tools/merge_checkpoints.py`](https://github.com/meta-llama/llama/blob/main/tools/merge_checkpoints.py) (available in the upstream repository) to merge the shards before loading.

## Code Examples

### Example 1: Explicit Model Parallel Size in Python

When calling the generation API directly, explicitly set `model_parallel_size` to match your shard count:

```python
from llama.generation import generate_llama

generate_llama(
    ckpt_dir="checkpoints/13B",
    tokenizer_path="tokenizer.model",
    prompt="Explain quantum computing",
    max_seq_len=512,
    model_parallel_size=2,  # Must equal the number of *.pth files

    seed=42,
)

```

### Example 2: Multi-Node Launch with torchrun

For distributed inference across multiple GPUs, ensure the environment variable matches the checkpoint sharding:

```bash

# Count shards: assume 4 files present

export WORLD_SIZE=4

torchrun --nproc_per_node=4 example_text_completion.py \
    --ckpt-dir checkpoints/7B \
    --tokenizer-path tokenizer.model \
    --max-seq-len 512 \
    --temperature 0.8

```

## Summary

- The "model parallel size does not match checkpoint files" error occurs in [`llama/generation.py`](https://github.com/meta-llama/llama/blob/main/llama/generation.py) when the number of parallel processes does not equal the number of checkpoint shards.
- **Always verify** that `WORLD_SIZE` or `model_parallel_size` matches the count of `.pth` files in your checkpoint directory.
- Launch inference using `torchrun --nproc_per_node=N` where `N` equals the shard count.
- For single-GPU inference on sharded checkpoints, either merge the shards using [`tools/merge_checkpoints.py`](https://github.com/meta-llama/llama/blob/main/tools/merge_checkpoints.py) or launch the appropriate number of processes.

## Frequently Asked Questions

### How do I check how many checkpoint shards my LLaMA model has?

Navigate to your checkpoint directory and count the `.pth` files using `ls /path/to/checkpoint/*.pth | wc -l`. The resulting number is the model parallel size you must use for inference. Each shard typically follows a naming pattern like `consolidated.00.pth`, `consolidated.01.pth`, etc.

### Can I run inference on a sharded checkpoint with a single GPU?

No, not directly. The LLaMA inference code requires the number of processes to match the number of shards. To run on a single GPU, you must first merge the checkpoint shards into a single file using the [`tools/merge_checkpoints.py`](https://github.com/meta-llama/llama/blob/main/tools/merge_checkpoints.py) utility script, then load the unified checkpoint with `model_parallel_size=1`.

### What is the difference between `WORLD_SIZE` and `model_parallel_size`?

`WORLD_SIZE` is an environment variable that sets the total number of processes participating in the distributed run, while `model_parallel_size` is a parameter passed directly to the generation functions in [`llama/generation.py`](https://github.com/meta-llama/llama/blob/main/llama/generation.py). Both serve the same purpose: telling the framework how many checkpoint shards to expect. If `model_parallel_size` is not provided, the code falls back to `WORLD_SIZE`.