xai-org/x-algorithm · error · ValueError

async_emb emb_width={emb_width} does not shard evenly over t

Error message

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

What it means

make_context_handle requires emb_width (the embedding feature dimension) to be divisible by group_size, the communicator rank count. The embedding rows are split width-wise across the sharding group, and a non-divisible width cannot be partitioned evenly across ranks.

Source

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

            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,
            num_devices_per_node,
        ),
    )
    return AsyncEmbContextHandle(
        context_id=context_id,

View on GitHub (pinned to 24c60942c5)

Solutions

Pad the embedding table width up to a multiple of group_size and slice after the collective.
Or reduce/reshape the table_axis parallel degree so it divides emb_width.
Add a startup assert: assert emb_width % group_size == 0 with both values in the message.

Example fix

# before
handle = make_context_handle(mesh, ..., emb_width=100, ...)  # 100 % 8 != 0

# after: pad width to 104 (multiple of 8), slice after lookup
emb_width = math.ceil(emb_width / group_size) * group_size  # 104
handle = make_context_handle(mesh, ..., emb_width=emb_width, ...)
# after lookup: emb = emb[..., :100]

Defensive patterns

Strategy: validation

Validate before calling

import math

gs = math.prod(mesh.shape[a] for a in table_axis)
assert emb_width % gs == 0, f"emb_width {emb_width} not divisible by group_size {gs}; pad the table"
# or: emb_width = math.ceil(emb_width / gs) * gs  and slice after lookup

Prevention

Keep embedding widths aligned to powers of two / multiples of the max sharding degree you plan to run.
Validate all table widths against the mesh before launching multi-host training.

When it happens

Trigger: Calling make_context_handle with emb_width=100 and group_size=8 (100 % 8 != 0). group_size is the product of table_axis mesh sizes from get_flatten_replica_groups(mesh, table_axis).

Common situations: Swapping an embedding table with width not divisible by the table-axis parallel degree (e.g. width 96 with 8-way sharding is fine, width 100 is not); increasing model parallelism without re-checking embedding sizes; mixed-dimension feature groups sharing one async_emb context.

Related errors

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