xai-org/x-algorithm · error · ValueError

async_emb requires token shards to vary across the communica

Error message

async_emb requires token shards to vary across the communicator: table_axis {missing_axes} missing from data_axis {data_axis}

What it means

make_context_handle requires every table_axis (embedding-sharding axis defining the communicator group) to also appear in data_axis (token-sharding axes). If tokens are not sharded along an axis on which tables are sharded, ranks in one communicator would hold embeddings for tokens they never see, which the async_emb collective protocol cannot handle.

Source

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


def make_context_handle(
    mesh: jax.sharding.Mesh,
    table_axis: tuple[str, ...],
    *,
    data_axis: tuple[str, ...],
    tokens_per_batch: int,
    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:

View on GitHub (pinned to 24c60942c5)

Solutions

Add every table_axis name into data_axis, e.g. table_axis=('table',), data_axis=('batch','table'), so token shards vary across the communicator.
If you did not intend to shard tables, reduce table_axis to a single axis that is already in data_axis.
Cross-check the axis names against mesh.shape (see also the adjacent mesh-axis validation errors).

Example fix

# before
handle = make_context_handle(
    mesh, table_axis=('table',), data_axis=('batch',), ...
)

# after
handle = make_context_handle(
    mesh, table_axis=('table',), data_axis=('batch', 'table'), ...
)

Defensive patterns

Strategy: validation

Validate before calling

def check_axis_containment(table_axis, data_axis):
    missing = [a for a in table_axis if a not in data_axis]
    if missing:
        raise ValueError(
            f"table_axis {missing} must also appear in data_axis {data_axis} "
            f"so token shards vary across the communicator"
        )

check_axis_containment(table_axis, data_axis)  # before make_context_handle

Type guard

def table_axes_in_data(table_axis, data_axis) -> bool:
    return set(table_axis) <= set(data_axis)

Try / catch

try:
    handle = make_context_handle(mesh, table_axis, data_axis, ...)
except ValueError as e:
    if "missing from data_axis" in str(e):
        data_axis = tuple(dict.fromkeys((*data_axis, *table_axis)))
        handle = make_context_handle(mesh, table_axis, data_axis, ...)
    else:
        raise

Prevention

Standardize on data_axis = table_axis + one batch axis for async_emb runs.
Add a config-lint step that reifies axis invariants before launching multi-host jobs.

When it happens

Trigger: Passing table_axis=('table','model') but data_axis=('batch',) — the 'table'/'model' axes are missing from data_axis — when calling make_context_handle / _create_async_emb_executables.

Common situations: Configuring data parallelism only ('batch') while enabling N-way model-parallel embedding sharding; reusing a single-tower mesh config for the two-tower async_emb path; copy-paste of axis tuples between experiments with different sharding layouts.

Related errors

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