xai-org/x-algorithm · error · ValueError

async_emb tokens_per_batch={tokens_per_batch} does not shard

Error message

async_emb tokens_per_batch={tokens_per_batch} does not shard evenly over the {group_size}-rank communicator

What it means

make_context_handle requires tokens_per_batch to be divisible by group_size, the number of ranks in the flattened embedding-sharding communicator. The async_emb kernels split the batch's tokens evenly across communicator ranks; a non-divisible batch leaves ranks with unequal token counts and the collective buffer layout becomes invalid.

Source

Thrown at phoenix/xrex/cuda/async_emb/async_emb.py:121

            raise ValueError(f"async_emb axis {axis!r} is not a mesh axis of {mesh}")
    group_size, flatten_replicas = get_flatten_replica_groups(mesh, table_axis)
    missing_axes = [axis for axis in table_axis if axis not in data_axis]
    if missing_axes:
        raise ValueError(
            f"async_emb requires token shards to vary across the communicator: "
            f"table_axis {missing_axes} missing from data_axis {data_axis}"
        )
    off_communicator_shards = math.prod(
        mesh.shape[axis] for axis in data_axis if axis not in table_axis
    )
    if off_communicator_shards != 1:
        raise ValueError(
            f"async_emb requires exactly one token shard per communicator rank: "
            f"data_axis {data_axis} shards tokens over {off_communicator_shards} "
            f"positions outside table_axis {table_axis}"
        )
    if tokens_per_batch % group_size != 0:
        raise ValueError(
            f"async_emb tokens_per_batch={tokens_per_batch} does not shard evenly "
            f"over the {group_size}-rank communicator"
        )
    if emb_width % group_size != 0:
        raise ValueError(
            f"async_emb emb_width={emb_width} does not shard evenly over the "
            f"{group_size}-rank communicator"
        )
    device_ids = [d.id for d in mesh.devices.flatten()]
    flatten_replicas = tuple(device_ids[pos] for pos in flatten_replicas)
    group_key = get_context_id(group_size, flatten_replicas)
    context_id = get_context_id(
        group_key,
        (
            tokens_per_batch // group_size,
            emb_width // group_size,
            emb_width,
            num_unique,

View on GitHub (pinned to 24c60942c5)

Solutions

Round tokens_per_batch up to the next multiple of group_size (pad tokens) before calling make_context_handle.
Or change the mesh/table_axis layout so group_size divides tokens_per_batch.
Add an early assert tokens_per_batch % group_size == 0 in your training entry point with a clear message including both values.

Example fix

# before
tokens_per_batch = 1_000_003  # odd count
handle = make_context_handle(mesh, ..., tokens_per_batch=tokens_per_batch, ...)

# after: pad to a multiple of group_size
group_size = math.prod(mesh.shape[a] for a in table_axis)
pad = (-tokens_per_batch) % group_size
tokens_per_batch += pad  # padded tokens, masked out in the loss
handle = make_context_handle(mesh, ..., tokens_per_batch=tokens_per_batch, ...)

Defensive patterns

Strategy: validation

Validate before calling

import math

def group_size_for(mesh, table_axis) -> int:
    return math.prod(mesh.shape[a] for a in table_axis)

gs = group_size_for(mesh, table_axis)
assert tokens_per_batch % gs == 0, f"tokens_per_batch {tokens_per_batch} not divisible by group_size {gs}"
# or pad: tokens_per_batch += (-tokens_per_batch) % gs

Prevention

Choose token batch sizes that are multiples of the table-axis parallel degree.
Re-assert divisibility whenever the mesh layout or packing strategy changes.

When it happens

Trigger: Calling make_context_handle with tokens_per_batch=1_000_000 and group_size=8 (or any non-divisible combination); group_size derives from get_flatten_replica_groups(mesh, table_axis), i.e. the product of the table_axis mesh sizes.

Common situations: Changing batch size / sequence packing to a value not divisible by the embedding-sharding degree; resizing the mesh (e.g. table axis 8→6 ranks) without re-checking token counts; padding removed from packed token batches.

Related errors

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/a8249dac8351c0af. Report an issue: GitHub.