[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

[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,...]

[raulchen]

good idea. code could be simpler

[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?

[raulchen]

this was added in the mid-ground state when both stacked and non-stacked are supported.
Should be non longer needed.

[raulchen]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a significant and well-executed refactoring to improve memory efficiency during training by implementing per-layer gradient checkpointing using jax.lax.scan and stacking decoder layer weights at initialization. While the architectural changes are impressive and include a unified code path for training and inference with comprehensive test coverage, a security audit identified several critical vulnerabilities. These include arbitrary method execution in the distributed backend, potential SSRF/path traversal in checkpoint loading, and missing ownership checks in model management, which must be addressed for system robustness and security, especially in multi-tenant or exposed environments. Furthermore, a correctness issue was found related to the collection of hidden states, affecting a corresponding test.

[gemini-code-assist]

When output_hidden_states is true, the list of hidden states is missing the output of the final decoder layer. The current implementation list(all_hs[:-1]) excludes the last layer's output from the scan operation.

To ensure all intermediate hidden states are collected, you should include all outputs from the scan. The list of hidden states returned by forward_layers should contain the initial embeddings plus the output of every decoder layer.

After this change, you will also need to update the assertion in test_hidden_states_length_matches to expect num_hidden_layers + 2 states (embeddings + N layer outputs + final normed output).

Suggested change

	all_hidden_states = [hidden_states] + list(all_hs[:-1]) if output_hidden_states else []
	all_hidden_states = [hidden_states] + list(all_hs) if output_hidden_states else []

[gemini-code-assist]

This assertion incorrectly expects num_hidden_layers + 1 hidden states. With the recommended fix in tx/models/utils.py to correctly collect all hidden states, the model will return num_hidden_layers + 2 states when output_hidden_states=True (initial embeddings, N layer outputs, and the final normalized output). Please update this assertion to reflect the correct number of hidden states.

Suggested change

	assert len(hidden_states_no_ckpt) == len(hidden_states_ckpt) == num_hidden_layers + 1
	assert len(hidden_states_no_ckpt) == len(hidden_states_ckpt) == num_hidden_layers + 2