{"record":{"id":"a8249dac8351c0af","repo":"xai-org/x-algorithm","slug":"async-emb-tokens-per-batch-tokens-per-batch-does","errorCode":null,"errorMessage":"async_emb tokens_per_batch={tokens_per_batch} does not shard evenly over the {group_size}-rank communicator","messagePattern":"async_emb tokens_per_batch=(.+?) does not shard evenly over the (.+?)-rank communicator","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/cuda/async_emb/async_emb.py","lineNumber":121,"sourceCode":"            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(\n        group_key,\n        (\n            tokens_per_batch // group_size,\n            emb_width // group_size,\n            emb_width,\n            num_unique,","sourceCodeStart":103,"sourceCodeEnd":139,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/cuda/async_emb/async_emb.py#L103-L139","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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."],"exampleFix":"# before\ntokens_per_batch = 1_000_003  # odd count\nhandle = make_context_handle(mesh, ..., tokens_per_batch=tokens_per_batch, ...)\n\n# after: pad to a multiple of group_size\ngroup_size = math.prod(mesh.shape[a] for a in table_axis)\npad = (-tokens_per_batch) % group_size\ntokens_per_batch += pad  # padded tokens, masked out in the loss\nhandle = make_context_handle(mesh, ..., tokens_per_batch=tokens_per_batch, ...)","handlingStrategy":"validation","validationCode":"import math\n\ndef group_size_for(mesh, table_axis) -> int:\n    return math.prod(mesh.shape[a] for a in table_axis)\n\ngs = group_size_for(mesh, table_axis)\nassert tokens_per_batch % gs == 0, f\"tokens_per_batch {tokens_per_batch} not divisible by group_size {gs}\"\n# or pad: tokens_per_batch += (-tokens_per_batch) % gs","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Choose token batch sizes that are multiples of the table-axis parallel degree.","Re-assert divisibility whenever the mesh layout or packing strategy changes."],"tags":["jax","sharding","batch-size","distributed","async-embedding"],"backgroundTag":"batch-not-divisible-by-parallel-degree","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}