xai-org/x-algorithm · critical · NotImplementedError

use_async_emb requires an embedding optimizer implementing A

Error message

use_async_emb requires an embedding optimizer implementing AsyncEmbOptimizer, and {type(self).__name__} has no fused table update

What it means

When use_async_emb is enabled, the trainer calls gradient_update_start on the embedding optimizer to perform a fused embedding-table update. The default implementation raises NotImplementedError, so any optimizer class that does not implement AsyncEmbOptimizer's fused update fails here, with the class name embedded in the message.

Source

Thrown at phoenix/xrex/optimizers/recsys/async_emb_gradient_update.py:32

class AsyncEmbGradientUpdate(NamedTuple):
    unique_tokens: jax.Array
    grads: jax.Array
    segment_ids: jax.Array
    pending: jax.Array


@runtime_checkable
class AsyncEmbOptimizer(Protocol):
    def gradient_update_start(
        self,
        context: async_emb.AsyncEmbContextHandle,
        update: AsyncEmbGradientUpdate,
        table: jax.Array,
        state: Any,
        gate: jax.Array,
    ) -> tuple[tuple[jax.Array, ...], jax.Array, Any, dict[str, jax.Array]]:
        raise NotImplementedError(
            f"use_async_emb requires an embedding optimizer implementing "
            f"AsyncEmbOptimizer, and {type(self).__name__} has no fused table update"
        )

    def gradient_update_done(
        self, context: async_emb.AsyncEmbContextHandle, gate: jax.Array
    ) -> tuple[jax.Array, jax.Array, jax.Array]:
        raise NotImplementedError(
            f"use_async_emb requires an embedding optimizer implementing "
            f"AsyncEmbOptimizer, and {type(self).__name__} has no fused table update"
        )

View on GitHub (pinned to 24c60942c5)

Solutions

1. Use an optimizer class that implements AsyncEmbOptimizer (fused table update) for the embeddings when use_async_emb=True
2. Subclass and implement gradient_update_start (and gradient_update_done) on your optimizer
3. Disable use_async_emb if fused async embedding updates are not needed

Example fix

# before
cfg = TrainConfig(use_async_emb=True)  # optimizer = default Adam
# after
cfg = TrainConfig(use_async_emb=False)
# or: optimizer = AsyncEmbAdam(...)  # implements AsyncEmbOptimizer

from phoenix.xrex.optimizers.recsys import async_emb
if cfg.use_async_emb:
    assert isinstance(emb_opt, async_emb.AsyncEmbOptimizer)

def supports_async_emb(opt) -> bool:
    return (hasattr(opt, 'gradient_update_start')
            and type(opt).gradient_update_start is not AsyncEmbGradientUpdate.gradient_update_start)

try:
    upd = emb_opt.gradient_update_start(ctx, update, table, state, gate)
except NotImplementedError:
    raise ConfigError('use_async_emb=True requires an AsyncEmbOptimizer') from None

When it happens

Trigger: Enabling use_async_emb in the recsys training config while using a plain optimizer (e.g. the standard Adam wrapper) for the embedding tables; instantiating AsyncEmbGradientUpdate and calling gradient_update_start on a subclass that only overrides gradient_update_done.

Common situations: Turning on the async embedding update feature of xrex without swapping the embedding optimizer for an AsyncEmbOptimizer implementation; upgrading where the async-emb API gained new required methods.

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 requires exactly one token shard per communicator
• async_emb tokens_per_batch={tokens_per_batch} does not shard
• async_emb emb_width={emb_width} does not shard evenly over t

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