Lightning-AI/pytorch-lightning · error · RuntimeError

The sizes `data_parallel_size={data_parallel_size}` and `ten

Error message

The sizes `data_parallel_size={data_parallel_size}` and `tensor_parallel_size={tensor_parallel_size}` multiplied should equal the world size ({world_size}).

What it means

_setup_device_mesh validates that data_parallel_size * tensor_parallel_size equals the distributed world size before calling torch.distributed.device_mesh.init_device_mesh. A mismatch means the parallel layout cannot map onto the available ranks.

Source

Thrown at src/lightning/fabric/strategies/model_parallel.py:526

        # return the remaining metadata that wasn't requested as part of `state`
        return checkpoint

    raise ValueError(
        f"The path {str(path)!r} does not point to a valid checkpoint. Make sure the path points to either a"
        " directory with distributed checkpoint shards, or a single file with a full checkpoint."
    )


def _setup_device_mesh(
    data_parallel_size: int,
    tensor_parallel_size: int,
    world_size: int,
    device: torch.device,
) -> "DeviceMesh":
    from torch.distributed.device_mesh import init_device_mesh

    if data_parallel_size * tensor_parallel_size != world_size:
        raise RuntimeError(
            f"The sizes `data_parallel_size={data_parallel_size}` and"
            f" `tensor_parallel_size={tensor_parallel_size}` multiplied should equal the world size"
            f" ({world_size})."
        )
    return init_device_mesh(
        device_type=device.type,
        mesh_shape=(data_parallel_size, tensor_parallel_size),
        mesh_dim_names=("data_parallel", "tensor_parallel"),
    )


def _has_dtensor_modules(module: object) -> TypeGuard[Module]:
    from torch.distributed._tensor import DTensor

    return isinstance(module, Module) and any(isinstance(t, DTensor) for t in module.parameters())


def _load_raw_module_state_from_path(path: _PATH, module: Module, world_size: int, strict: bool = True) -> None:

View on GitHub (pinned to 9fed5c27d2)

Solutions

Set data_parallel_size * tensor_parallel_size == world_size (e.g. world_size=4 → dp=2, tp=2, or dp=1, tp=4)
Or change the launcher/process count to match the configured sizes (e.g. torchrun --nproc_per_node=4)
Compute sizes dynamically: tensor_parallel_size = min(world_size, desired_tp); data_parallel_size = world_size // tensor_parallel_size

Example fix

# before
strategy = ModelParallelStrategy(tensor_parallel_size=4)  # run with world_size=2
# after
world_size = torch.distributed.get_world_size()
tp = min(4, world_size)
strategy = ModelParallelStrategy(tensor_parallel_size=tp, data_parallel_size=world_size // tp)

Defensive patterns

Strategy: validation

Validate before calling

ws = torch.distributed.get_world_size()
assert data_parallel_size * tensor_parallel_size == ws, (
    f'dp({data_parallel_size}) * tp({tensor_parallel_size}) != world_size({ws})'
)

Prevention

Derive dp/tp sizes from world_size instead of hardcoding
Parameterize sizes via config tied to the launcher's nproc_per_node

When it happens

Trigger: Creating ModelParallelStrategy(data_parallel_size=2, tensor_parallel_size=2) and launching with world_size != 4 (e.g. 2 GPUs, or 8 GPUs); using a different number of processes/devices than the strategy's parallel sizes imply.

Common situations: Scaling a TP config from one machine size to another (e.g. 8xGPU config run on 2xGPU) without adjusting sizes; env var overrides of world size; mixed use of devices=['auto'] and hardcoded sizes.

Related errors

AI-assisted analysis of Lightning-AI/pytorch-lightning@9fed5c27d2 (2026-08-28). Data as JSON: /api/errors/8420d43eee2c4462. Report an issue: GitHub.