Switch tx model tests to AutoConfig + truncate layer_types under shrink

transformers 5.4 turned PreTrainedConfig into a @strict @DataClass with
class validators. Two patterns broke under transformers 5.8:

1. `PretrainedConfig.from_pretrained(model_name)` no longer round-trips
   model-specific config fields. With rope_parameters + a missing
   max_position_embeddings, validate fails. Switch every test caller to
   `AutoConfig.from_pretrained` (mirroring the production-side fix
   already adopted from PR #1561 in skyrl/backends/jax.py).

2. validate_layer_type asserts `len(layer_types) == num_hidden_layers`.
   tests/tx/utils/test_models.py:create_test_model shrinks
   num_hidden_layers to 1 to keep the test cheap, but layer_types is
   inherited from the real Qwen3-0.6B config (28 entries) and the
   wrapping Qwen3Config validator then raises. Truncate layer_types
   alongside the num_hidden_layers override.

Verified locally on the cpu jax suite (CI=true, CUDA hidden to match
the GitHub Actions cpu environment): all previously-failing tests
in test_deepseekv3.py, test_deepseekv3_lora_training.py,
test_llama3_lora_training.py, test_qwen3.py, test_qwen3_config.py,
and test_models.py now pass.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: erictang000

tests/tx/models/test_deepseekv3.py

@@ -7,7 +7,7 @@
 import pytest
 import torch
 from flax import nnx
-from transformers import AutoModelForCausalLM, AutoTokenizer, PretrainedConfig
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
 from transformers.models.deepseek_v3.modeling_deepseek_v3 import (
     DeepseekV3MoE as HFDeepseekV3MoE,
 )
@@ -40,7 +40,7 @@ def test_deepseekv3(tp: int):
     with tempfile.TemporaryDirectory() as tmp:
         hf_model.save_pretrained(tmp, safe_serialization=True)
 
-        base_config = PretrainedConfig.from_pretrained(model_name)
+        base_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
         config = DeepseekV3Config(base_config, max_lora_adapters=32, max_lora_rank=32, shard_attention_heads=True)
         # EP axis required for MoE expert sharding
         mesh = jax.make_mesh((1, 1, tp), ("fsdp", "ep", "tp"), axis_types=(jax.sharding.AxisType.Auto,) * 3)
@@ -87,7 +87,7 @@ def test_deepseekv3_moe_layer(ep: int, tp: int):
     hf_model = AutoModelForCausalLM.from_pretrained(
         model_name, attn_implementation="eager", use_safetensors=True, torch_dtype=torch.float32
     )
-    base_config = PretrainedConfig.from_pretrained(model_name)
+    base_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
     config = DeepseekV3Config(base_config, max_lora_adapters=0, max_lora_rank=0, shard_attention_heads=True)
 
     # Initial deepseek layers don't have MoE
@@ -136,7 +136,7 @@ def test_deepseekv3_moe_layer_lora(ep: int, tp: int):
     hf_model = AutoModelForCausalLM.from_pretrained(
         model_name, attn_implementation="eager", use_safetensors=True, torch_dtype=torch.float32
     )
-    base_config = PretrainedConfig.from_pretrained(model_name)
+    base_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
     config = DeepseekV3Config(base_config, max_lora_adapters=3, max_lora_rank=4, shard_attention_heads=True)
 
     hf_moe_layer = hf_model.model.layers[1].mlp
@@ -211,7 +211,7 @@ def test_deepseekv3_gradient_checkpointing():
     that gradient checkpointing works correctly with heterogeneous layer types.
     """
     model_name = "yujiepan/deepseek-v3-tiny-random"
-    base_config = PretrainedConfig.from_pretrained(model_name)
+    base_config = AutoConfig.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)

tests/tx/models/test_deepseekv3_lora_training.py

@@ -3,7 +3,7 @@
 import optax
 from flax import nnx
 from huggingface_hub import snapshot_download
-from transformers import PretrainedConfig
+from transformers import AutoConfig
 
 from skyrl.tinker.types import LoraConfig
 from skyrl.tx.layers.lora import init_lora_adapter
@@ -19,7 +19,7 @@
 
 def test_lora_training_moe_rank_normalized():
     base_model = "yujiepan/deepseek-v3-tiny-random"
-    base_config = PretrainedConfig.from_pretrained(base_model, trust_remote_code=True)
+    base_config = AutoConfig.from_pretrained(base_model, trust_remote_code=True)
     config = DeepseekV3Config(base_config, max_lora_adapters=5, max_lora_rank=32, shard_attention_heads=True)
 
     checkpoint_path = snapshot_download(base_model, allow_patterns=["*.safetensors"])
@@ -100,7 +100,7 @@ def loss_for_lora(lora_params):
 
 def test_lora_training_high_rank():
     base_model = "yujiepan/deepseek-v3-tiny-random"
-    base_config = PretrainedConfig.from_pretrained(base_model, trust_remote_code=True)
+    base_config = AutoConfig.from_pretrained(base_model, trust_remote_code=True)
     config = DeepseekV3Config(base_config, max_lora_adapters=5, max_lora_rank=32, shard_attention_heads=True)
 
     checkpoint_path = snapshot_download(base_model, allow_patterns=["*.safetensors"])

tests/tx/models/test_llama3_lora_training.py

@@ -3,7 +3,7 @@
 import optax
 from flax import nnx
 from huggingface_hub import snapshot_download
-from transformers import PretrainedConfig
+from transformers import AutoConfig
 
 from skyrl.tinker.types import LoraConfig
 from skyrl.tx.layers.lora import init_lora_adapter
@@ -19,7 +19,7 @@
 
 def test_lora_training():
     base_model = "unsloth/Llama-3.2-1B"
-    base_config = PretrainedConfig.from_pretrained(base_model)
+    base_config = AutoConfig.from_pretrained(base_model)
     config = Llama3Config(base_config, max_lora_adapters=5, max_lora_rank=32, shard_attention_heads=True)
 
     checkpoint_path = snapshot_download(base_model, allow_patterns=["*.safetensors"])

tests/tx/models/test_qwen3.py

@@ -7,7 +7,7 @@
 import torch
 from flax import nnx
 from peft import LoraConfig, get_peft_model
-from transformers import AutoModelForCausalLM, AutoTokenizer, PretrainedConfig
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
 from transformers.models.qwen3_moe.modeling_qwen3_moe import (
     Qwen3MoeSparseMoeBlock as HFQwen3MoeSparseMoeBlock,
 )
@@ -71,7 +71,7 @@ def test_qwen3_moe_layer(ep: int, tp: int):
     hf_model = AutoModelForCausalLM.from_pretrained(
         model_name, attn_implementation="eager", use_safetensors=True, torch_dtype=torch.float32
     )
-    base_config = PretrainedConfig.from_pretrained(model_name)
+    base_config = AutoConfig.from_pretrained(model_name)
     config = Qwen3Config(base_config, max_lora_adapters=0, max_lora_rank=0, shard_attention_heads=True)
 
     hf_moe_layer = hf_model.model.layers[0].mlp
@@ -126,7 +126,7 @@ def test_qwen3_moe_layer_lora(ep: int, tp: int):
     hf_model = AutoModelForCausalLM.from_pretrained(
         model_name, attn_implementation="eager", use_safetensors=True, torch_dtype=torch.float32
     )
-    base_config = PretrainedConfig.from_pretrained(model_name)
+    base_config = AutoConfig.from_pretrained(model_name)
     config = Qwen3Config(base_config, max_lora_adapters=3, max_lora_rank=4, shard_attention_heads=True)
 
     hf_moe_layer = hf_model.model.layers[0].mlp

tests/tx/models/test_qwen3_5_lora_training.py

@@ -3,7 +3,7 @@
 import optax
 from flax import nnx
 from huggingface_hub import snapshot_download
-from transformers import PretrainedConfig
+from transformers import AutoConfig
 
 from skyrl.tinker.types import LoraConfig
 from skyrl.tx.layers.lora import init_lora_adapter
@@ -19,7 +19,7 @@
 
 def test_lora_training():
     base_model = "Qwen/Qwen3.5-0.8B"
-    base_config = PretrainedConfig.from_pretrained(base_model)
+    base_config = AutoConfig.from_pretrained(base_model)
     config = Qwen3_5Config(base_config, max_lora_adapters=5, max_lora_rank=32, shard_attention_heads=True)
 
     checkpoint_path = snapshot_download(base_model, allow_patterns=["*.safetensors"])

tests/tx/models/test_qwen3_config.py

@@ -1,13 +1,13 @@
 """Tests for Qwen3Config."""
 
-from transformers import PretrainedConfig
+from transformers import AutoConfig
 
 from skyrl.tx.models.configs import Qwen3Config
 
 
 def test_config_wraps_pretrained_config():
     """Test that Qwen3Config wraps a PretrainedConfig and adds LoRA params."""
-    hf_config = PretrainedConfig.from_pretrained("Qwen/Qwen3-0.6B")
+    hf_config = AutoConfig.from_pretrained("Qwen/Qwen3-0.6B")
     config = Qwen3Config(hf_config, max_lora_adapters=8, max_lora_rank=16, shard_attention_heads=False)
 
     # Check LoRA params were set
@@ -23,7 +23,7 @@ def test_config_wraps_pretrained_config():
 
 def test_config_preserves_moe_config():
     """Test that MoE-specific configs are preserved."""
-    hf_config = PretrainedConfig.from_pretrained("trl-internal-testing/tiny-Qwen3MoeForCausalLM")
+    hf_config = AutoConfig.from_pretrained("trl-internal-testing/tiny-Qwen3MoeForCausalLM")
     config = Qwen3Config(hf_config, max_lora_adapters=3, max_lora_rank=4, shard_attention_heads=True)
 
     # Check that MoE-specific attributes are preserved

tests/tx/models/test_qwen3_generate.py

@@ -8,7 +8,7 @@
 import pytest
 import torch
 from flax import nnx
-from transformers import AutoModelForCausalLM, AutoTokenizer, PretrainedConfig
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
 
 from skyrl.tinker import types
 from skyrl.tx.models.configs import Qwen3Config
@@ -45,7 +45,7 @@ def test_qwen3_generate():
     # Generate with our implementation (batched with right-padding)
     with tempfile.TemporaryDirectory() as tmp:
         hf_model.save_pretrained(tmp, safe_serialization=True)
-        base_config = PretrainedConfig.from_pretrained(model_name)
+        base_config = AutoConfig.from_pretrained(model_name)
         config = Qwen3Config(base_config, max_lora_adapters=2, max_lora_rank=32, shard_attention_heads=True)
 
         mesh = jax.make_mesh((1, 1), ("fsdp", "tp"), axis_types=(jax.sharding.AxisType.Auto,) * 2)
@@ -131,7 +131,7 @@ def test_qwen3_generate_speed():
     hf_model = AutoModelForCausalLM.from_pretrained(
         model_name, attn_implementation="eager", use_safetensors=True, torch_dtype=torch.float32
     )
-    base_config = PretrainedConfig.from_pretrained(model_name)
+    base_config = AutoConfig.from_pretrained(model_name)
     config = Qwen3Config(base_config, max_lora_adapters=32, max_lora_rank=32, shard_attention_heads=True)
 
     inputs = [

tests/tx/utils/test_models.py

@@ -34,6 +34,13 @@ def create_test_model(base_model_name: str, rank: int, alpha: int, adapter_index
     base_config.intermediate_size = 128
     base_config.num_attention_heads = 2
     base_config.num_key_value_heads = 2
+    # transformers >=5.4 has a strict validator (validate_layer_type) that
+    # asserts len(layer_types) == num_hidden_layers when layer_types is set.
+    # When we shrink num_hidden_layers above, also truncate layer_types so
+    # validation still passes.
+    layer_types = getattr(base_config, "layer_types", None)
+    if layer_types is not None:
+        base_config.layer_types = list(layer_types[: base_config.num_hidden_layers])
 
     config = Qwen3Config(base_config, max_lora_adapters=5, max_lora_rank=32, shard_attention_heads=True)