[tx] Per-layer gradient checkpointing with stacked decoder layers

[raulchen]

Summary

Implement per-layer gradient checkpointing using jax.lax.scan with permanently stacked decoder layer weights via
nnx.vmap. This reduces peak memory by ~num_layers factor during training while maintaining a unified code path for training
and inference.

Key Changes

1. Per-layer Gradient Checkpointing

• Use jax.lax.scan with jax.checkpoint to recompute activations during backward pass
• XLA compiles ONE loop body and reuses buffers, unlike Python loops which unroll N separate checkpoint regions
• Enable via gradient_checkpointing=True in model config

2. Stacked Layer Weights

• Stack decoder layer weights at initialization using nnx.vmap → shape (num_layers, ...)
• Eliminates runtime stacking overhead (weights already in stacked format)
• Single forward_layers() function for both training and inference
• KV cache uses stacked format (num_layers, batch, seq, heads, dim)

3. DeepSeekV3 Split Stacking

• Handles heterogeneous layers (dense MLP vs MoE) with separate stacks
• dense_layers for initial layers, moe_layers for MoE layers
• KV caches merged after forward pass

Files Changed

• tx/models/utils.py - New: create_stacked_layers(), forward_layers()
• tx/models/{llama3,qwen3,deepseekv3}.py - Use stacked layers
• tx/layers/lora.py - Stacked LoRA indexing
• tx/utils/models.py - Stack/unstack for HF checkpoint compatibility
• tx/utils/generator.py - Stacked KV cache
• tx/tinker/backends/jax.py - Fix gradient accumulation for stacked params

Test plan

• Forward outputs match with/without checkpointing
• Gradients match with/without checkpointing
• All model tests pass (37)
• All tinker tests pass (19)
• DeepSeekV3 EP=2 tests pass

[gemini-code-assist]

Code Review

This pull request introduces a major refactoring to use stacked weights for model layers, which is a significant architectural improvement for performance in JAX. The changes are extensive, replacing nnx.List of layers with nnx.vmap-created stacked modules and using jax.lax.scan for an efficient forward pass. Key logic is well-encapsulated in the new tx/models/utils.py file. The refactoring is consistently applied across different model architectures (Llama3, Qwen3, DeepseekV3), and the test suite has been updated to reflect these changes, including the addition of a shared test utility module and new tests for gradient checkpointing. The code quality is high and the changes are well-executed.

[raulchen]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a significant and well-executed refactoring to implement per-layer gradient checkpointing and stacked decoder layers, aiming to reduce memory usage during training. The core logic is cleanly encapsulated in the new tx/models/utils.py, which provides a unified forward_layers function for both training and inference. The changes are consistently propagated across all models, with careful handling of the heterogeneous layers in DeepSeekV3. The utilities for LoRA parameter management and weight loading/saving are correctly generalized for the new stacked architecture. The test suite has been impressively updated to validate the correctness of this new approach, including checks for gradient matching. I have one minor suggestion to improve type safety. Overall, this is a high-quality and impactful contribution.

[gemini-code-assist]

The __call__ method of this base class uses self.mlp, but mlp is not defined or hinted in the base class itself. This can be confusing for readers and may be flagged by static type checkers. Adding a class-level type hint for mlp would improve readability and maintainability.

For example:

class DeepseekV3DecoderLayer(nnx.Module):
    mlp: nnx.Module
    ...

[pcmoritz]

It will be good to get rid of this function (so we also don't depend on the jax.core.Tracer class, which looks like more of an internal class).

The best way to handle this I think is to introduce a required sharding argument in the constructors of the classes LoRA{Layer} and then call nnx.with_partitioning(...) inside the classes. This way the sharding is always required and we can even get rid of the assert.

I'm happy to make a separate PR with the refactor (do let me know if you see a problems with this).

[pcmoritz]

Can we simplify this a little and get rid of these classes? What I have in mind is to just have the class DeepseekV3DecoderLayer from above and do

class DeepseekV3DecoderLayer(nnx.Module):
    def __init__(self, mlp_cls: Callable[[ModelConfig, ...], nnx.Module], config: DeepseekV3Config, *, dtype: jnp.dtype, rngs: nnx.Rngs) -> None:
        # ...
        self.mlp = mlp_cls(config, dtype=dtype, rngs=rngs)

and then down below where the stacked layers are created, do

def create_dense_layer(rngs: nnx.Rngs) -> DeepseekV3DecoderLayer:
    return DeepseekV3DecoderLayer(DeepseekV3MLP, config, dtype=dtype, rngs=rngs)

[pcmoritz]

This merits a KVCache.concatenate() function I think, it could even be implemented generically like

@staticmethod
def concatenate(*kv_caches: KVCache) -> KVCache:
    return KVCache(
        keys=jnp.concatenate(kv_cache.key for kv_cache in kv_caches),
        values=jnp.concatenate(kv_cache.value for kv_cache in kv_caches),
        cache_position=kv_caches[0].cache_position,
    )

The reasoning is that we try to keep this kind of complexity out of the modeling files as much as possible.

[pcmoritz]

Similar to the below, it would be nice to write this as something like

dense_kv_cache, moe_kv_cache = kv_cache.split(self.num_dense_layers, self.num_moe_layers)

If you wanted, you could probably implement this genericlly by lifting https://docs.jax.dev/en/latest/_autosummary/jax.numpy.split.html, something like

def split(kv_cache: KVCache, indices_or_sections):
    keys_list = jnp.split(kv_cache.keys, indices_or_sections, axis=0)
    values_list =  jnp.split(kv_cache.keys, indices_or_sections, axis=0)
    return [KVCache(keys=keys, values=values, cache_position=kv_cache.cache_position) for (keys, values) in zip(keys_list, values_list))

but not sure if that's actually better. Otherwise just supporting the case of two splits is totally fine.

[pcmoritz]

I think we should change this function to

def get_adapter_idx(path: str, adapter_index: int):
    is_stacked = _is_stacked(path)
    return (slice(None), adapter_index) if is_stacked else (adapter_index,)

and then use it everywhere in the PR. I believe all the call sited of _adapter_index are already of this form, and there are lots more call sites where we can get rid of a pattern like

if is_stacked:
     # Process stacked weights
else:
     # Process unstacked weights

This can always be done like

idx = get_adapter_idx(path, adapter_index)
# Process weights with weights[idx,...]

[pcmoritz]

Do we actually need this function? Can't we just do things like

idx = get_adapter_idx(path, adapter_index)
p.at[idx, ..., :, :rank]
p.at[idx, ..., :rank, :]

[pcmoritz]

Do you know why this case is needed? If it isn't, we only need one method here and can always use the logic of is_stacked_lora_path.

My rationale for this question is, if we always stack the layers that are stackable, a case like ('model', 'layers', '0', 'self_attn', ...) should never happen, right?