{"record":{"id":"0f4e77c35542039b","repo":"xai-org/x-algorithm","slug":"async-emb-requires-exactly-one-token-shard-per-com","errorCode":null,"errorMessage":"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}","messagePattern":"async_emb requires exactly one token shard per communicator rank: data_axis (.+?) shards tokens over (.+?) positions outside table_axis (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/cuda/async_emb/async_emb.py","lineNumber":115,"sourceCode":"    emb_width: int,\n    num_unique: int,\n    num_devices_per_node: int,\n) -> AsyncEmbContextHandle:\n    for axis in dict.fromkeys((*table_axis, *data_axis)):\n        if axis not in mesh.shape:\n            raise ValueError(f\"async_emb axis {axis!r} is not a mesh axis of {mesh}\")\n    group_size, flatten_replicas = get_flatten_replica_groups(mesh, table_axis)\n    missing_axes = [axis for axis in table_axis if axis not in data_axis]\n    if missing_axes:\n        raise ValueError(\n            f\"async_emb requires token shards to vary across the communicator: \"\n            f\"table_axis {missing_axes} missing from data_axis {data_axis}\"\n        )\n    off_communicator_shards = math.prod(\n        mesh.shape[axis] for axis in data_axis if axis not in table_axis\n    )\n    if off_communicator_shards != 1:\n        raise ValueError(\n            f\"async_emb requires exactly one token shard per communicator rank: \"\n            f\"data_axis {data_axis} shards tokens over {off_communicator_shards} \"\n            f\"positions outside table_axis {table_axis}\"\n        )\n    if tokens_per_batch % group_size != 0:\n        raise ValueError(\n            f\"async_emb tokens_per_batch={tokens_per_batch} does not shard evenly \"\n            f\"over the {group_size}-rank communicator\"\n        )\n    if emb_width % group_size != 0:\n        raise ValueError(\n            f\"async_emb emb_width={emb_width} does not shard evenly over the \"\n            f\"{group_size}-rank communicator\"\n        )\n    device_ids = [d.id for d in mesh.devices.flatten()]\n    flatten_replicas = tuple(device_ids[pos] for pos in flatten_replicas)\n    group_key = get_context_id(group_size, flatten_replicas)\n    context_id = get_context_id(","sourceCodeStart":97,"sourceCodeEnd":133,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/cuda/async_emb/async_emb.py#L97-L133","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Add the extra data axis to table_axis too (so it participates in the communicator), e.g. table_axis=('batch','replica').","Or drop the extra axis from data_axis / from the mesh so tokens are sharded exactly once outside the communicator.","If outer replication is required, run separate processes/jobs per replica instead of a single mesh with an extra data axis."],"exampleFix":"# before: 2x outer data-parallel replicas -> off_communicator_shards=2\nmesh = Mesh(devices.reshape(2, 8), ('replica', 'batch'))\nmake_context_handle(mesh, table_axis=('batch',), data_axis=('batch', 'replica'), ...)\n\n# after: replicate along table axis as well\nmesh = Mesh(devices.reshape(2, 8), ('replica', 'batch'))\nmake_context_handle(mesh, table_axis=('batch', 'replica'), data_axis=('batch', 'replica'), ...)","handlingStrategy":"validation","validationCode":"import math\n\ndef check_single_shard_per_rank(mesh, table_axis, data_axis):\n    off = math.prod(mesh.shape[a] for a in data_axis if a not in table_axis)\n    if off != 1:\n        raise ValueError(\n            f\"{off} token shards outside the communicator; add the extra axes \"\n            f\"to table_axis or drop them from data_axis\"\n        )\n\ncheck_single_shard_per_rank(mesh, table_axis, data_axis)  # before make_context_handle","typeGuard":"def single_shard_per_rank(mesh, table_axis, data_axis) -> bool:\n    return math.prod(mesh.shape[a] for a in data_axis if a not in table_axis) == 1","tryCatchPattern":null,"preventionTips":["When adding a replica/DP axis, mirror it into table_axis so the communicator grows with it.","Prefer separate jobs per replica over extra data axes when using async_emb."],"tags":["jax","sharding","data-parallelism","distributed","async-embedding"],"backgroundTag":"jax-sharding-axis-mismatch","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}