{"record":{"id":"28ffa0b1ca143068","repo":"xai-org/x-algorithm","slug":"async-emb-emb-width-emb-width-does-not-shard-eve","errorCode":null,"errorMessage":"async_emb emb_width={emb_width} does not shard evenly over the {group_size}-rank communicator","messagePattern":"async_emb emb_width=(.+?) 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":126,"sourceCode":"            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,\n            num_devices_per_node,\n        ),\n    )\n    return AsyncEmbContextHandle(\n        context_id=context_id,","sourceCodeStart":108,"sourceCodeEnd":144,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/cuda/async_emb/async_emb.py#L108-L144","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","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."],"exampleFix":"# before\nhandle = make_context_handle(mesh, ..., emb_width=100, ...)  # 100 % 8 != 0\n\n# after: pad width to 104 (multiple of 8), slice after lookup\nemb_width = math.ceil(emb_width / group_size) * group_size  # 104\nhandle = make_context_handle(mesh, ..., emb_width=emb_width, ...)\n# after lookup: emb = emb[..., :100]","handlingStrategy":"validation","validationCode":"import math\n\ngs = math.prod(mesh.shape[a] for a in table_axis)\nassert emb_width % gs == 0, f\"emb_width {emb_width} not divisible by group_size {gs}; pad the table\"\n# or: emb_width = math.ceil(emb_width / gs) * gs  and slice after lookup","typeGuard":null,"tryCatchPattern":null,"preventionTips":["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."],"tags":["jax","sharding","embedding-dimension","distributed","async-embedding"],"backgroundTag":"dimension-not-divisible-by-parallel-degree","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}