xai-org/x-algorithm · error · ValueError

async_emb requires exactly one token shard per communicator

Error message

async_emb requires exactly one token shard per communicator rank: data_axis {data_axis} shards tokens over {off_communicator_shards} positions outside table_axis {table_axis}

What it means

After validating axis containment, make_context_handle computes off_communicator_shards = product of mesh.shape[axis] for data_axis axes NOT in table_axis, and requires it to equal 1. async_emb assumes each communicator rank owns exactly one token shard; extra outer data-parallel replication (e.g. a 'replica' axis outside the communicator) breaks the one-shard-per-rank invariant.

Source

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

    emb_width: int,
    num_unique: int,
    num_devices_per_node: int,
) -> AsyncEmbContextHandle:
    for axis in dict.fromkeys((*table_axis, *data_axis)):
        if axis not in mesh.shape:
            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(

View on GitHub (pinned to 24c60942c5)

Solutions

1. Add the extra data axis to table_axis too (so it participates in the communicator), e.g. table_axis=('batch','replica').
2. Or drop the extra axis from data_axis / from the mesh so tokens are sharded exactly once outside the communicator.
3. If outer replication is required, run separate processes/jobs per replica instead of a single mesh with an extra data axis.

Example fix

# before: 2x outer data-parallel replicas -> off_communicator_shards=2
mesh = Mesh(devices.reshape(2, 8), ('replica', 'batch'))
make_context_handle(mesh, table_axis=('batch',), data_axis=('batch', 'replica'), ...)

# after: replicate along table axis as well
mesh = Mesh(devices.reshape(2, 8), ('replica', 'batch'))
make_context_handle(mesh, table_axis=('batch', 'replica'), data_axis=('batch', 'replica'), ...)

import math

def check_single_shard_per_rank(mesh, table_axis, data_axis):
    off = math.prod(mesh.shape[a] for a in data_axis if a not in table_axis)
    if off != 1:
        raise ValueError(
            f"{off} token shards outside the communicator; add the extra axes "
            f"to table_axis or drop them from data_axis"
        )

check_single_shard_per_rank(mesh, table_axis, data_axis)  # before make_context_handle

def single_shard_per_rank(mesh, table_axis, data_axis) -> bool:
    return math.prod(mesh.shape[a] for a in data_axis if a not in table_axis) == 1

When it happens

Trigger: Having a mesh axis in data_axis but not in table_axis with size > 1, e.g. mesh axes {'batch': 8, 'replica': 2} with table_axis=('batch',), data_axis=('batch','replica') → off_communicator_shards=2.

Common situations: Adding a data-parallel replica axis for gradient accumulation/multi-host replication without replicating the embedding table axis as well; scaling from single-host to multi-host layouts; mixing model-parallel and data-parallel axes in one mesh for async_emb training.

Related errors

• async_emb axis {axis!r} is not a mesh axis of {mesh}
• async_emb requires token shards to vary across the communica
• async_emb tokens_per_batch={tokens_per_batch} does not shard
• async_emb emb_width={emb_width} does not shard evenly over t
• Unable to create named shape with unnamed dimensions (shape:

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