[JAX SC] Implement pipelined SparseCore embedding layer in Flax

This change introduces `PipelinedSparseCoreEmbed` to enable pipelining of SparseCore operations with TensorCore. The layer stores inputs and activations/gradients across steps to allow for overlapping computation. The Shakespeare example is updated to use this new pipelined layer, requiring changes to how variables are handled in the training loop, including using `mutable=True` in `model.apply`. A test target for the Shakespeare example is also added.

With pipelining enabled, the model converges as shown below:
```
I 2025-09-28T13:41:31.529508-07:00 4428  jax_sc_shakespeare_jit_flax.py:415] Step 9: Loss = 7.518053
I 2025-09-28T13:41:31.599092-07:00 4428  jax_sc_shakespeare_jit_flax.py:415] Step 19: Loss = 7.317211
I 2025-09-28T13:41:31.668971-07:00 4428  jax_sc_shakespeare_jit_flax.py:415] Step 29: Loss = 7.0158334
I 2025-09-28T13:41:31.738482-07:00 4428  jax_sc_shakespeare_jit_flax.py:415] Step 39: Loss = 6.609592
I 2025-09-28T13:41:31.807324-07:00 4428  jax_sc_shakespeare_jit_flax.py:415] Step 49: Loss = 6.140646
I 2025-09-28T13:41:31.876329-07:00 4428  jax_sc_shakespeare_jit_flax.py:415] Step 59: Loss = 5.7092185
I 2025-09-28T13:41:31.949826-07:00 4428  jax_sc_shakespeare_jit_flax.py:415] Step 69: Loss = 5.3921995
I 2025-09-28T13:41:32.018904-07:00 4428  jax_sc_shakespeare_jit_flax.py:415] Step 79: Loss = 5.183298
I 2025-09-28T13:41:32.088231-07:00 4428  jax_sc_shakespeare_jit_flax.py:415] Step 89: Loss = 5.0565457
I 2025-09-28T13:41:32.156878-07:00 4428  jax_sc_shakespeare_jit_flax.py:415] Step 99: Loss = 4.9826536
...
I 2025-09-28T13:41:38.309954-07:00 4428  jax_sc_shakespeare_jit_flax.py:415] Step 999: Loss = 0.31050137
```

PiperOrigin-RevId: 812524634

Author: adityagupta1089

jax_tpu_embedding/sparsecore/examples/models/shakespeare/flax_model.py

@@ -20,6 +20,7 @@
 from jax_tpu_embedding.sparsecore.lib.nn import embedding
 from jax_tpu_embedding.sparsecore.lib.nn import embedding_spec
 
+
 shard_map = jax.experimental.shard_map.shard_map
 Nested = embedding.Nested
 
@@ -38,6 +39,7 @@ class Model(nn.Module):
   feature_name: str = 'shakespeare_feature'
   mesh: jax.sharding.Mesh | None = None
   sharding_axis: str = 'sparsecore_sharding'
+  enable_pipelining: bool = False
 
   def add_sharding_constraint(self, x: jax.Array, names: tuple[str | None]):
     # Add a sharding constraint to the array.
@@ -66,6 +68,7 @@ def __call__(self, embedding_lookup_inputs: embedding.PreprocessedInput):
         feature_specs=self.feature_specs,
         mesh=self.mesh,
         sharding_axis=self.sharding_axis,
+        enable_pipelining=self.enable_pipelining,
     )(embedding_lookup_inputs)
 
     # Unpack the activations.

jax_tpu_embedding/sparsecore/lib/flax/embed.py

@@ -16,10 +16,12 @@
 import functools
 from typing import Any, Callable, Mapping, TypeVar
 
+import flax
 from flax import linen as nn
 from flax import typing
 import jax
 from jax.experimental import layout
+import jax.numpy as jnp
 from jax_tpu_embedding.sparsecore.lib.nn import embedding
 from jax_tpu_embedding.sparsecore.lib.nn import embedding_spec
 from jax_tpu_embedding.sparsecore.utils import utils
@@ -69,7 +71,32 @@ def wrapper(*args, **kwargs):
 
 
 class SparseCoreEmbed(nn.Module):
-  """SparseCore embedding layer."""
+  """SparseCore embedding layer.
+
+  ## Pipelining
+
+  This layer supports pipelining of SparseCore computations if
+  `enable_pipelining` is set to True. Pipelining decouples the SC computation
+  with TC computation by processing multiple batches concurrently stored in
+  internal state (variables). This allows for greater SC-TC overlap and
+  generally better performance at the cost of higher memory usage.
+  There's however a comparitively slower convergence which is tolerable in most
+  cases. See internal link:jax-sc-embedding-pipelining for more information.
+
+  When pipelining is enabled, it implements a two-stage pipeline: embedding
+  lookups for batch `i` run concurrently with TensorCore computations for batch
+  `i-1` and embedding gradient updates for batch `i-2`. This results in
+  activations being delayed by one step and gradient updates by two steps
+  relative to the inputs.
+
+  NOTE for pipelining:
+  * The first two steps return zero activations (warm-up), therefore user needs
+  to run two additional steps. The dense input for first(0) and last(N+1) could
+  be dummy input.
+  * If pipelining is enabled, user will have to pass
+  `mutable=['sparsecore_pipeline_state']` to `.apply()` to
+  update internal pipeline state.
+  """
 
   # A sequence of FeatureSpecs to specify the configurations for the
   # input feature.
@@ -81,6 +108,7 @@ class SparseCoreEmbed(nn.Module):
   # Sharding strategy for embedding tables.
   table_sharding_strategy: str = 'MOD'
   enable_minibatching: bool = False
+  enable_pipelining: bool = False
 
   num_sc_per_device: int = -1  # Initialized in __post_init__.
 
@@ -165,6 +193,7 @@ def preprocess_inputs(
         all_reduce_interface=all_reduce_interface,
     )[0]
 
+  @nn.compact
   def __call__(
       self, embedding_lookup_inputs: EmbeddingLookupInput
   ) -> embedding.Nested[jax.Array]:
@@ -176,12 +205,19 @@ def __call__(
     Returns:
       The activations structure with the same structure as feature_specs.
     """
+    if self.enable_pipelining:
+      return self._pipelined_call(embedding_lookup_inputs)
     return _emb_lookup(
         self,
         embedding_lookup_inputs,
         self.embedding_table,
     )
 
+  def _get_embedding_table_path(self) -> str:
+    """Returns the path to the embedding table within the module."""
+    return '/'.join(self.path + (EMBEDDING_PARAM_NAME,))
+
+  @nn.compact
   def apply_gradient(
       self,
       gradients: embedding.Nested[jax.Array],
@@ -196,13 +232,131 @@ def apply_gradient(
     Returns:
       The updated activation embedding tables.
     """
+    if self.enable_pipelining:
+      return self._pipelined_apply_gradient(gradients, embedding_lookup_inputs)
     _, embed_table = _emb_lookup_bwd(
         self,
         (embedding_lookup_inputs, self.embedding_table),
         gradients,
     )
-    path = '/'.join(self.path + (EMBEDDING_PARAM_NAME,))
-    return {path: embed_table}
+    return {self._get_embedding_table_path(): embed_table}
+
+  ##############################################################################
+  # Pipelining implementation
+  ##############################################################################
+  def _pipelined_call(
+      self, embedding_lookup_inputs: embedding.PreprocessedInput
+  ) -> embedding.Nested[jax.Array]:
+    """Pipelined version of __call__."""
+    step_im1_sparse_activations = self._get_step_im1_sparse_activations()
+    step_im2_sparse_inputs = self._get_step_im2_sparse_inputs(
+        embedding_lookup_inputs
+    )
+    step_im2_sparse_gradients = self._get_step_im2_sparse_gradients()
+
+    # Update embedding table using values from step i-2.
+    # Perform the update using a custom_vjp to avoid differentiating the
+    # optimizer step.
+    updated_table_val = _perform_update(
+        self,
+        step_im2_sparse_inputs.value,
+        step_im2_sparse_gradients.value,
+        self.embedding_table,
+    )
+    if self.is_mutable_collection('params'):
+      self.scope.variables()['params'][
+          EMBEDDING_PARAM_NAME
+      ] = self.scope.variables()['params'][EMBEDDING_PARAM_NAME].replace_boxed(
+          updated_table_val
+      )
+
+    updated_table = updated_table_val
+
+    # The activations for the current step's forward pass are from step i-1.
+    result_activations = step_im1_sparse_activations.value
+
+    # Now, perform the lookup for the current step (i) using the newly updated
+    # embedding table and store the inputs and resulting activations for future
+    # steps.
+    self._get_step_im1_sparse_inputs(embedding_lookup_inputs).value = (
+        embedding_lookup_inputs
+    )
+    step_im1_sparse_activations.value = _emb_lookup(
+        self, embedding_lookup_inputs, updated_table
+    )
+
+    return result_activations
+
+  def _pipelined_apply_gradient(
+      self,
+      gradients: embedding.Nested[jax.Array],
+      embedding_lookup_inputs: embedding.PreprocessedInput,
+  ) -> Mapping[str, Mapping[str, jax.Array]]:
+    """Pipelined version of apply_gradient."""
+    step_im1_sparse_inputs = self._get_step_im1_sparse_inputs(
+        embedding_lookup_inputs
+    )
+    step_im2_sparse_inputs = self._get_step_im2_sparse_inputs(
+        embedding_lookup_inputs
+    )
+    step_im2_sparse_gradients = self._get_step_im2_sparse_gradients()
+
+    # Store sparse inputs and gradients for use on next step.
+    step_im2_sparse_inputs.value = step_im1_sparse_inputs.value
+    step_im2_sparse_gradients.value = gradients
+
+    return {}
+
+  def _get_unfrozen_feature_specs(
+      self,
+  ) -> embedding.Nested[embedding_spec.FeatureSpec]:
+    return (
+        flax.core.unfreeze(self.feature_specs)
+        if isinstance(self.feature_specs, flax.core.FrozenDict)
+        else self.feature_specs
+    )
+
+  def _get_step_im1_sparse_inputs(
+      self, embedding_lookup_inputs: embedding.PreprocessedInput
+  ) -> flax.core.scope.Variable[embedding.PreprocessedInput]:
+    return self.variable(
+        SPARSECORE_PIPELINE_STATE_COLLECTION,
+        'step_im1_sparse_inputs',
+        lambda: jax.tree.map(jnp.zeros_like, embedding_lookup_inputs),
+    )
+
+  def _get_step_im1_sparse_activations(
+      self,
+  ) -> flax.core.scope.Variable[embedding.Nested[jax.Array]]:
+    return self.variable(
+        SPARSECORE_PIPELINE_STATE_COLLECTION,
+        'step_im1_sparse_activations',
+        lambda: jax.tree.map(
+            lambda f: jnp.zeros(f.output_shape, dtype=jnp.float32),
+            self._get_unfrozen_feature_specs(),
+        ),
+    )
+
+  def _get_step_im2_sparse_inputs(
+      self, embedding_lookup_inputs: embedding.PreprocessedInput
+  ) -> flax.core.scope.Variable[embedding.PreprocessedInput]:
+    return self.variable(
+        SPARSECORE_PIPELINE_STATE_COLLECTION,
+        'step_im2_sparse_inputs',
+        lambda: jax.tree.map(jnp.zeros_like, embedding_lookup_inputs),
+    )
+
+  def _get_step_im2_sparse_gradients(
+      self,
+  ) -> flax.core.scope.Variable[embedding.Nested[jax.Array]]:
+    return self.variable(
+        SPARSECORE_PIPELINE_STATE_COLLECTION,
+        'step_im2_sparse_gradients',
+        lambda: jax.tree.map(
+            lambda f: jnp.zeros(f.output_shape, dtype=jnp.float32),
+            self._get_unfrozen_feature_specs(),
+        ),
+    )
 
 
 ################################################################################
@@ -255,7 +409,7 @@ def _emb_lookup_bwd(embedding_layer, res, gradients):
 
   pt = embedding_layer.embedding_table_partition
   pd = embedding_layer.data_partition
-  emb_table_grads = shard_map(
+  updated_emb_table = shard_map(
       functools.partial(
           embedding.tpu_sparse_dense_matmul_grad,
           feature_specs=embedding_layer.feature_specs,
@@ -272,11 +426,12 @@ def _emb_lookup_bwd(embedding_layer, res, gradients):
       emb_table,
   )
 
+  # > Reinterpret emb table as grads (unused).
   # tpu_sparse_dense_matmul_grad returns a general Mapping (usually a dict).
   # It may not be the same type as the embedding table (e.g. FrozenDict).
   # Here we use flatten / unflatten to ensure the types are the same.
   emb_table_grads = jax.tree.unflatten(
-      jax.tree.structure(emb_table), jax.tree.leaves(emb_table_grads)
+      jax.tree.structure(emb_table), jax.tree.leaves(updated_emb_table)
   )
 
   return (
@@ -286,3 +441,68 @@ def _emb_lookup_bwd(embedding_layer, res, gradients):
 
 
 _emb_lookup.defvjp(_emb_lookup_fwd, _emb_lookup_bwd)
+
+
+SPARSECORE_PIPELINE_STATE_COLLECTION = 'sparsecore_pipeline_state'
+
+
+################################################################################
+# Define custom VJP for embedding update for pipelining.
+# This is used to prevent autodiff from differentiating through the optimizer
+# step.
+################################################################################
+@functools.partial(jax.custom_vjp, nondiff_argnums=(0,))
+def _perform_update(
+    module: 'SparseCoreEmbed',
+    im2_inputs: embedding.PreprocessedInput,
+    im2_grads: embedding.Nested[jax.Array],
+    emb_table: embedding.Nested[jax.Array],
+) -> embedding.Nested[jax.Array]:
+  """Performs the embedding update, but is opaque to autodiff."""
+  _, updated_table = _emb_lookup_bwd(  # pylint: disable=protected-access
+      module,
+      (im2_inputs, emb_table),
+      im2_grads,
+  )
+  return updated_table
+
+
+def _perform_update_fwd(
+    module: 'SparseCoreEmbed',
+    im2_inputs: embedding.PreprocessedInput,
+    im2_grads: embedding.Nested[jax.Array],
+    emb_table: embedding.Nested[jax.Array],
+):
+  """Forward pass for _perform_update."""
+  updated_table = _perform_update(module, im2_inputs, im2_grads, emb_table)
+  # Return inputs as residuals for backward pass.
+  return updated_table, (im2_inputs, im2_grads, emb_table)
+
+
+def _perform_update_bwd(
+    module: 'SparseCoreEmbed',
+    res: tuple[
+        embedding.PreprocessedInput,
+        embedding.Nested[jax.Array],
+        embedding.Nested[jax.Array],
+    ],
+    g: embedding.Nested[jax.Array],
+) -> tuple[
+    embedding.PreprocessedInput,
+    embedding.Nested[jax.Array],
+    embedding.Nested[jax.Array],
+]:
+  """Backward pass for _perform_update."""
+  # g is the gradient w.r.t. the output (updated_table).
+  # We want this to flow back to the original emb_table as if this function
+  # was an identity function.
+  im2_inputs, im2_grads, emb_table = res
+  del module, emb_table
+  return (
+      jax.tree.map(jnp.zeros_like, im2_inputs),
+      jax.tree.map(jnp.zeros_like, im2_grads),
+      g,  # Pass gradient through to the original embedding table.
+  )
+
+
+_perform_update.defvjp(_perform_update_fwd, _perform_update_bwd)