rohitg00/ai-engineering-from-scratch · error · ValueError

reduce_scatter needs numel divisible by world_size, got {n}

Error message

reduce_scatter needs numel divisible by world_size, got {n} / {w}

What it means

Error "reduce_scatter needs numel divisible by world_size, got {n} / {w}" thrown in rohitg00/ai-engineering-from-scratch.

Source

Thrown at phases/19-capstone-projects/76-collective-ops-from-scratch/code/main.py:184

        mesh.send(next_rank, shards[send_idx])
        shards[recv_idx] = mesh.recv(prev_rank)
    return torch.cat(shards)


def reduce_scatter(mesh: Mesh, tensor: torch.Tensor) -> torch.Tensor:
    """Reduce-scatter as the first half of ring allreduce.

    Input is a tensor of length world_size * T. Output is the rank's chunk of
    length T holding the sum across all ranks for that index range. The
    underlying ring algorithm parks the full sum at index (r + 1) % W; we
    return that chunk and label it as rank r's output to match
    torch.distributed's contract that rank r owns chunks[r].
    """
    w = mesh.world_size
    r = mesh.rank
    n = tensor.numel()
    if n % w != 0:
        raise ValueError(f"reduce_scatter needs numel divisible by world_size, got {n} / {w}")
    if w == 1:
        return tensor.clone()
    rotated = list(tensor.chunk(w))
    rotated = [rotated[(i - 1) % w].clone() for i in range(w)]
    chunks = rotated
    next_rank = (r + 1) % w
    prev_rank = (r - 1) % w
    for step in range(w - 1):
        send_idx = (r - step) % w
        recv_idx = (r - step - 1) % w
        mesh.send(next_rank, chunks[send_idx])
        incoming = mesh.recv(prev_rank)
        chunks[recv_idx] = chunks[recv_idx] + incoming
    return chunks[(r + 1) % w]


def _gloo_worker(rank: int, world_size: int, op: str, tensor_bytes: bytes,
                 shape, dtype_str: str, init_file: str,

View on GitHub (pinned to 39ea8a1c6d)

When it happens

Trigger: Thrown at phases/19-capstone-projects/76-collective-ops-from-scratch/code/main.py:184 when the library encounters an invalid state.

Common situations: See trigger scenarios.

AI-assisted analysis of rohitg00/ai-engineering-from-scratch@39ea8a1c6d (2026-08-26). Data as JSON: /api/errors/10ac0b178502411b. Report an issue: GitHub.