[JAX SC] Implement pipelined SparseCore embedding layer in Flax #496
+725
−141
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
![copybara-service[bot]](https://avatars.githubusercontent.com/in/44061?s=60&v=4){kind=small}
copybara-service
bot
force-pushed
the
test_812524634
branch
from
f053f91 to
9c5f322
Compare
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
copybara-service
bot
force-pushed
the
test_812524634
branch
from
9c5f322 to
424271f
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
[JAX SC] Implement pipelined SparseCore embedding layer in Flax
This change introduces
PipelinedSparseCoreEmbedto 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 usingmutable=Trueinmodel.apply. A test target for the Shakespeare example is also added.With pipelining enabled, the model converges as shown below: