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

skyrl-tx/tests/models/lora_test_utils.py

@@ -0,0 +1,81 @@
+"""Shared test utilities for LoRA training tests."""
+
+import jax
+import jax.numpy as jnp
+
+from tx.utils.models import get_adapter_idx
+
+
+def get_adapter_params(params, adapter_idx: int):
+    """Extract adapter params at a specific index.
+
+    Decoder layer LoRA params have shape (num_layers, num_adapters, ...).
+    Embed tokens LoRA params have shape (num_adapters, ...).
+    """
+
+    def extract(path, p):
+        idx = get_adapter_idx(path, adapter_idx)
+        return p[idx].copy()
+
+    return jax.tree.map_with_path(extract, params)
+
+
+def _slice_out_of_rank(params, adapter_idx: int, get_rank):
+    """Extract out-of-rank params using a rank function.
+
+    Args:
+        params: LoRA parameters tree.
+        adapter_idx: Adapter index to extract.
+        get_rank: Function (path) -> int returning effective rank for that path.
+    """
+
+    def slice_param(path, p):
+        path_str = str(path)
+        if "lora_A" not in path_str and "lora_B" not in path_str:
+            return p
+        rank = get_rank(path)
+        idx = get_adapter_idx(path, adapter_idx)
+        if "lora_A" in path_str:
+            return p[idx + (..., slice(rank, None))].copy()
+        return p[idx + (..., slice(rank, None), slice(None))].copy()
+
+    return jax.tree.map_with_path(slice_param, params)
+
+
+def get_out_of_rank_params(params, adapter_idx: int, rank: int):
+    """Extract out-of-rank params for an adapter."""
+    return _slice_out_of_rank(params, adapter_idx, lambda _: rank)
+
+
+def verify_params_unchanged(initial_params, final_params, error_msg_prefix: str):
+    """Verify that params haven't changed between initial and final state."""
+    for (path, initial), (_, final) in zip(
+        jax.tree.leaves_with_path(initial_params), jax.tree.leaves_with_path(final_params)
+    ):
+        assert jnp.allclose(initial, final), f"{error_msg_prefix} for {path}"
+
+
+def _is_routed_expert_path(path) -> bool:
+    """Check if path is for routed experts (not shared_experts)."""
+    keys = []
+    for p in path:
+        if hasattr(p, "key"):
+            keys.append(str(p.key))
+        elif hasattr(p, "name"):
+            keys.append(str(p.name))
+    for i, key in enumerate(keys):
+        if key == "experts" and i > 0 and keys[i - 1] == "mlp":
+            return True
+    return False
+
+
+def get_moe_out_of_rank_params(params, adapter_idx: int, rank: int, num_experts: int):
+    """Extract out-of-rank params for MoE models.
+
+    For routed experts, uses effective rank = max(1, rank // num_experts).
+    """
+
+    def get_rank(path):
+        return max(1, rank // num_experts) if _is_routed_expert_path(path) else rank
+
+    return _slice_out_of_rank(params, adapter_idx, get_rank)

skyrl-tx/tests/models/test_deepseekv3.py

@@ -186,3 +186,51 @@ def test_deepseekv3_moe_layer_lora(ep: int, tp: int):
             output_merged = moe_layer_merged(x_sample)
 
             assert np.allclose(output_with_lora[sample_idx : sample_idx + 1], output_merged, rtol=1e-3, atol=1e-3)
+
+
+def test_deepseekv3_gradient_checkpointing():
+    """Test that gradient checkpointing produces identical outputs for DeepSeekV3.
+
+    DeepSeekV3 has split stacking (dense_layers + moe_layers), so this tests
+    that gradient checkpointing works correctly with heterogeneous layer types.
+    """
+    model_name = "yujiepan/deepseek-v3-tiny-random"
+    base_config = PretrainedConfig.from_pretrained(model_name, trust_remote_code=True)
+
+    batch_size, seq_len = 2, 8
+    mesh = jax.make_mesh((1, 1, 1), ("fsdp", "ep", "tp"), axis_types=(jax.sharding.AxisType.Auto,) * 3)
+
+    results = {}
+    for use_checkpointing in [False, True]:
+        config = DeepseekV3Config(
+            base_config,
+            max_lora_adapters=1,
+            max_lora_rank=1,
+            shard_attention_heads=True,
+            gradient_checkpointing=use_checkpointing,
+        )
+        with jax.set_mesh(mesh):
+            model = DeepseekV3ForCausalLM(config, dtype=jnp.float32, rngs=nnx.Rngs(0))
+
+            input_ids = jax.random.randint(jax.random.key(42), (batch_size, seq_len), 0, config.vocab_size)
+            attention_mask = jnp.ones((batch_size, seq_len), dtype=jnp.int32)
+
+            out = model(input_ids, attention_mask=attention_mask, output_hidden_states=True)
+            logits = model.compute_logits(out.last_hidden_state)
+
+            results[use_checkpointing] = {
+                "logits": np.array(logits),
+                "hidden_states": [np.array(hs) for hs in out.hidden_states],
+                "kv_cache_shape": out.kv_cache.keys.shape,
+            }
+
+    # Verify outputs match
+    np.testing.assert_allclose(results[False]["logits"], results[True]["logits"], rtol=1e-4, atol=1e-6)
+
+    # Verify hidden states match
+    assert len(results[False]["hidden_states"]) == len(results[True]["hidden_states"])
+    for i, (hs_no_ckpt, hs_ckpt) in enumerate(zip(results[False]["hidden_states"], results[True]["hidden_states"])):
+        np.testing.assert_allclose(hs_no_ckpt, hs_ckpt, rtol=1e-4, atol=1e-6, err_msg=f"Mismatch at hidden state {i}")
+
+    # Verify KV cache shape is correct (num_layers, batch, seq, heads, dim)
+    assert results[True]["kv_cache_shape"][0] == config.num_hidden_layers

skyrl-tx/tests/models/test_deepseekv3_lora_training.py

@@ -11,42 +11,11 @@
 from tx.layers.lora import init_lora_adapter
 from tx.tinker.types import LoraConfig
 
-
-def _is_routed_expert_path(path) -> bool:
-    """Disambiguate shared_experts and experts"""
-    keys = []
-    for p in path:
-        if hasattr(p, "key"):
-            keys.append(str(p.key))
-        elif hasattr(p, "name"):
-            keys.append(str(p.name))
-
-    for i, key in enumerate(keys):
-        if key == "experts" and i > 0 and keys[i - 1] == "mlp":
-            return True
-    return False
-
-
-def _get_out_of_rank_params(params, adapter_idx: int, rank: int, num_experts: int):
-    """Extract out-of-rank params, using effective rank for routed expert layers."""
-
-    def slice_param(path, p):
-        path_str = str(path)
-
-        if _is_routed_expert_path(path):
-            effective_rank = max(1, rank // num_experts)
-        else:
-            effective_rank = rank
-
-        if "lora_A" in path_str:
-            # lora_A shape: [adapters, ..., max_rank] - slice last dim
-            return p[adapter_idx, ..., effective_rank:].copy()
-        elif "lora_B" in path_str:
-            # lora_B shape: [adapters, ..., max_rank, out] - slice second-to-last dim
-            return p[adapter_idx, ..., effective_rank:, :].copy()
-        return p
-
-    return jax.tree.map_with_path(slice_param, params)
+from tests.models.lora_test_utils import (
+    get_adapter_params,
+    get_moe_out_of_rank_params,
+    verify_params_unchanged,
+)
 
 
 def test_lora_training_moe_rank_normalized():
@@ -85,15 +54,12 @@ def loss_fn(model, input_ids, target_ids, attention_mask):
 
         graphdef, lora_params, non_lora_params = nnx.split(model, model.is_lora_param, ...)
 
-        def get_adapter_params(params, adapter_idx):
-            return jax.tree.map(lambda p: p[adapter_idx].copy(), params)
-
         num_experts = config.n_routed_experts
 
         # Save initial states
         initial_adapter_2_params = get_adapter_params(lora_params, 2)
-        initial_adapter_0_out_of_rank = _get_out_of_rank_params(lora_params, 0, 16, num_experts)
-        initial_adapter_1_out_of_rank = _get_out_of_rank_params(lora_params, 1, 8, num_experts)
+        initial_adapter_0_out_of_rank = get_moe_out_of_rank_params(lora_params, 0, 16, num_experts)
+        initial_adapter_1_out_of_rank = get_moe_out_of_rank_params(lora_params, 1, 8, num_experts)
 
         initial_loss = None
 
@@ -116,24 +82,18 @@ def loss_for_lora(lora_params):
 
         final_loss = float(loss)
 
-        def verify_params_unchanged(initial_params, final_params, error_msg_prefix):
-            for (path, initial), (_, final) in zip(
-                jax.tree.leaves_with_path(initial_params), jax.tree.leaves_with_path(final_params)
-            ):
-                assert jnp.allclose(initial, final), f"{error_msg_prefix} for {path}"
-
         assert final_loss < initial_loss, f"Loss did not decrease: {initial_loss} -> {final_loss}"
 
         # Verify unused adapter was not modified
         final_adapter_2_params = get_adapter_params(lora_params, 2)
         verify_params_unchanged(initial_adapter_2_params, final_adapter_2_params, "Adapter 2 was modified")
 
         # Verify out-of-rank params were not modified
-        final_adapter_0_out_of_rank = _get_out_of_rank_params(lora_params, 0, 16, num_experts)
+        final_adapter_0_out_of_rank = get_moe_out_of_rank_params(lora_params, 0, 16, num_experts)
         verify_params_unchanged(
             initial_adapter_0_out_of_rank, final_adapter_0_out_of_rank, "Adapter 0 out-of-rank params modified"
         )
-        final_adapter_1_out_of_rank = _get_out_of_rank_params(lora_params, 1, 8, num_experts)
+        final_adapter_1_out_of_rank = get_moe_out_of_rank_params(lora_params, 1, 8, num_experts)
         verify_params_unchanged(
             initial_adapter_1_out_of_rank, final_adapter_1_out_of_rank, "Adapter 1 out-of-rank params modified"
         )
@@ -172,9 +132,6 @@ def loss_fn(model, input_ids, target_ids, attention_mask):
 
         graphdef, lora_params, non_lora_params = nnx.split(model, model.is_lora_param, ...)
 
-        def get_adapter_params(params, adapter_idx):
-            return jax.tree.map(lambda p: p[adapter_idx].copy(), params)
-
         num_experts = config.n_routed_experts
 
         # Save initial states for all unused adapters
@@ -183,8 +140,8 @@ def get_adapter_params(params, adapter_idx):
         initial_adapter_4_params = get_adapter_params(lora_params, 4)
 
         # Save out-of-rank params for adapters 0 and 1
-        initial_adapter_0_out_of_rank = _get_out_of_rank_params(lora_params, 0, 16, num_experts)
-        initial_adapter_1_out_of_rank = _get_out_of_rank_params(lora_params, 1, 8, num_experts)
+        initial_adapter_0_out_of_rank = get_moe_out_of_rank_params(lora_params, 0, 16, num_experts)
+        initial_adapter_1_out_of_rank = get_moe_out_of_rank_params(lora_params, 1, 8, num_experts)
 
         # Training loop
         for step in range(10):
@@ -200,12 +157,6 @@ def loss_for_lora(lora_params):
 
             print(f"Step {step}: loss = {float(loss):.4f}")
 
-        def verify_params_unchanged(initial_params, final_params, error_msg_prefix):
-            for (path, initial), (_, final) in zip(
-                jax.tree.leaves_with_path(initial_params), jax.tree.leaves_with_path(final_params)
-            ):
-                assert jnp.allclose(initial, final), f"{error_msg_prefix} for {path}"
-
         # Verify unused adapters (2, 3, 4) were not modified
         final_adapter_2_params = get_adapter_params(lora_params, 2)
         verify_params_unchanged(initial_adapter_2_params, final_adapter_2_params, "Adapter 2 was modified")
@@ -217,11 +168,11 @@ def verify_params_unchanged(initial_params, final_params, error_msg_prefix):
         verify_params_unchanged(initial_adapter_4_params, final_adapter_4_params, "Adapter 4 was modified")
 
         # Verify out-of-rank params were not modified
-        final_adapter_0_out_of_rank = _get_out_of_rank_params(lora_params, 0, 16, num_experts)
+        final_adapter_0_out_of_rank = get_moe_out_of_rank_params(lora_params, 0, 16, num_experts)
         verify_params_unchanged(
             initial_adapter_0_out_of_rank, final_adapter_0_out_of_rank, "Adapter 0 out-of-rank params modified"
         )
-        final_adapter_1_out_of_rank = _get_out_of_rank_params(lora_params, 1, 8, num_experts)
+        final_adapter_1_out_of_rank = get_moe_out_of_rank_params(lora_params, 1, 8, num_experts)
         verify_params_unchanged(
             initial_adapter_1_out_of_rank, final_adapter_1_out_of_rank, "Adapter 1 out-of-rank params modified"
         )

skyrl-tx/tests/models/test_llama3_lora_training.py

@@ -11,6 +11,8 @@
 from tx.layers.lora import init_lora_adapter
 from tx.tinker.types import LoraConfig
 
+from tests.models.lora_test_utils import get_adapter_params, get_out_of_rank_params, verify_params_unchanged
+
 
 def test_lora_training():
     base_model = "unsloth/Llama-3.2-1B"
@@ -45,21 +47,6 @@ def loss_fn(model, input_ids, target_ids, attention_mask):
         # that we want to compute gradients for
         graphdef, lora_params, non_lora_params = nnx.split(model, model.is_lora_param, ...)
 
-        # Helper to extract adapter params at specific index
-        def get_adapter_params(params, adapter_idx):
-            return jax.tree.map(lambda p: p[adapter_idx].copy(), params)
-
-        # Helper to extract out-of-rank params for an adapter
-        def get_out_of_rank_params(params, adapter_idx, rank):
-            def slice_param(path, p):
-                if "lora_A" in str(path):
-                    return p[adapter_idx, :, rank:].copy()
-                elif "lora_B" in str(path):
-                    return p[adapter_idx, rank:, :].copy()
-                return p
-
-            return jax.tree.map_with_path(slice_param, params)
-
         # Save initial states
         initial_adapter_2_params = get_adapter_params(lora_params, 2)
         initial_adapter_0_out_of_rank = get_out_of_rank_params(lora_params, 0, 16)
@@ -79,12 +66,6 @@ def loss_for_lora(lora_params):
 
             print(f"Step {step}: loss = {float(loss):.4f}")
 
-        def verify_params_unchanged(initial_params, final_params, error_msg_prefix):
-            for (path, initial), (_, final) in zip(
-                jax.tree.leaves_with_path(initial_params), jax.tree.leaves_with_path(final_params)
-            ):
-                assert jnp.allclose(initial, final), f"{error_msg_prefix} for {path}"
-
         # Verify adapter 2 (unused) was not modified
         final_adapter_2_params = get_adapter_params(lora_params, 2)
         verify_params_unchanged(initial_adapter_2_params, final_adapter_2_params, "Adapter 2 was modified")