xai-org/x-algorithm · error · ValueError

fused lazy decay needs timestamped state (step/last_step)

Error message

fused lazy decay needs timestamped state (step/last_step)

What it means

Thrown by the fused rowwise Adagrad embedding optimizer when lazy decay (decaying stale embedding rows only when they are touched) is enabled but the optimizer state lacks step counters. The fused CUDA kernel applies per-row decay based on the difference between the global step and each row's last-updated step, so both state.step and state.last_step must be present.

Source

Thrown at phoenix/xrex/optimizers/recsys/rowwise_adagrad.py:103

        token_ids: jax.Array,
        state: Any,
    ) -> tuple[Any, Any]:
        return embeddings, None

    @property
    def decay_factor(self) -> float:
        return math.exp(-self._decay_rate) if self._decay_rate is not None else 1.0

    def gradient_update_start(
        self,
        context: async_emb.AsyncEmbContextHandle,
        update: AsyncEmbGradientUpdate,
        table: jax.Array,
        state: RecsysRowwiseAdagradState,
        gate: jax.Array,
    ) -> tuple[tuple[jax.Array, ...], jax.Array, RecsysRowwiseAdagradState]:
        if self._lazy_decay and (state.step is None or state.last_step is None):
            raise ValueError("fused lazy decay needs timestamped state (step/last_step)")

        if 32 % context.shard_width != 0:
            raise ValueError(
                f"the fused rowwise Adagrad update needs a row shard that divides a warp, "
                f"got shard_width={context.shard_width} (emb_width={context.emb_width})"
            )

        from xrex.cuda.async_emb import async_emb

        metrics: dict[str, jax.Array] = {}

        if self._lazy_decay:

            @shard_map(
                mesh=context.mesh,
                in_specs=(
                    P(context.data_axis, None),
                    P(context.data_axis),

View on GitHub (pinned to 24c60942c5)

Solutions

1. Initialize the optimizer state with step counters (e.g. zeros for last_step and a step counter array) before enabling _lazy_decay
2. If timestamped state is unavailable, disable _lazy_decay and use dense decay instead
3. Migrate/regenerate the checkpoint so it includes step and last_step fields

Example fix

# before
opt = RecsysRowwiseAdagrad(lazy_decay=True)
# state = opt.init_state(...)  # no step/last_step

# after
opt = RecsysRowwiseAdagrad(lazy_decay=True)
state = opt.init_state(..., with_step_counters=True)  # state.step / state.last_step populated

if opt._lazy_decay and (state.step is None or state.last_step is None):
    raise ValueError("initialize step/last_step before enabling lazy decay")

def has_timestamps(s: RecsysRowwiseAdagradState) -> bool:
    return s.step is not None and s.last_step is not None

When it happens

Trigger: Calling gradient_update_start with _lazy_decay=True while the RecsysRowwiseAdagradState was constructed without step/last_step fields (e.g. a freshly initialized or deserialized state that omits timestamps).

Common situations: Enabling lazy decay in a recsys training config after previously running without it, reusing old checkpoints whose optimizer state predates timestamping, or building the state manually in tests.

Related errors

• use_async_emb requires an embedding optimizer implementing A
• the fused rowwise Adagrad update needs a row shard that divi
• type checking expression %s failed: invalid argument type: e
• Expected 2D emb_table, got shape={embeddings.shape}
• Unsupported emb_table dtype: {embeddings.dtype}

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