[tx] add per-layer gradient checkpointing with scan for memory-efficient training

[raulchen]

Summary

• Add per-layer gradient checkpointing using jax.lax.scan with jax.checkpoint
• Reduces peak memory by ~num_layers factor (only one layer's activations held during backward)
• Enable via gradient_checkpointing=True in model config

Implementation

Uses jax.lax.scan so XLA compiles ONE loop body and reuses buffers during backward recomputation. With a Python loop, XLA unrolls N separate checkpoint regions and can't optimize buffer reuse across them.

Tradeoff: requires stacking all layer weights once per forward pass. This is acceptable because checkpointing already trades
compute for memory.

Test plan

• Forward outputs match with/without checkpointing
• Hidden states length and values match
• is_training=False uses standard path with KV cache
• Gradients match between checkpointed and non-checkpointed paths

Benchmark results #891 (comment)

[gemini-code-assist]

Code Review

This pull request successfully integrates per-layer gradient checkpointing using jax.lax.scan and jax.checkpoint, which is a significant improvement for memory efficiency during training. The implementation correctly separates the checkpointed and non-checkpointed forward paths, and the new tests adequately cover the correctness of outputs, hidden states, and gradients. The is_training flag and gradient_checkpointing configuration are well-handled across the models and backend. There are a few minor areas for improvement related to test clarity and robustness.

[vercel]

@raulchen is attempting to deploy a commit to the Tyler's projects Team on Vercel.

A member of the Team first needs to authorize it.

[pcmoritz]

Can we unify these two methods into a single method forward_layers that will do the right thing depending on the parameters that are passed in? The reasoning is that model authors should not be concerned with the difference between checkpointing and non-checkpointing, they should just be able to write the model, and also we would want to avoid duplicating the code into all the model definitions.

Internally you can of course still have _forward_layers_checkpointed and _ forward_layers and implement the "public" forward_layers function in terms of those.

[pcmoritz]

An alternative way to do this would be to use https://flax.readthedocs.io/en/v0.8.3/api_reference/flax.experimental.nnx/module.html#flax.experimental.nnx.Module.train (and https://flax.readthedocs.io/en/v0.8.3/api_reference/flax.experimental.nnx/module.html#flax.experimental.nnx.Module.eval) to switch the module to train/eval mode. This would have some benefits: (a) it is the standard way to do something like this (similar to how pytorch does it, there you see this kind of pattern a lot), (b) if we are going to add dropout in the future, handling training / eval like this will be much easier and (c) we don't need to pass the flags around if we want to access similar attributes. It is less explicit though, and we would need to provide our own attribute to set. Curious about your thoughts :)