refactor : align config field names with Hugging Face

- dim -> hidden_size, n_layers -> num_hidden_layers
- dim_ffn -> intermediate_size, n_heads -> num_attention_heads
- n_kv_heads -> num_key_value_heads, max_len -> max_position_embeddings
- norm_eps -> rms_norm_eps, tie_weight -> tie_word_embeddings
- update model, inference, training, scripts, tests, docs

tests/conftest.py

@@ -107,13 +107,13 @@ def test_model():
     """Session-scoped small AutoRegressiveLM model, created once."""
     config = AutoRegressiveLMConfig(
         vocab_size=1000,
-        dim=8,
-        n_heads=2,
-        n_kv_heads=1,
-        dim_ffn=16,
-        max_len=64,
-        n_layers=2,
-        norm_eps=1e-5,
+        hidden_size=8,
+        num_attention_heads=2,
+        num_key_value_heads=1,
+        intermediate_size=16,
+        max_position_embeddings=64,
+        num_hidden_layers=2,
+        rms_norm_eps=1e-5,
     )
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model = AutoRegressiveLM(config).to(device=device)
@@ -137,13 +137,13 @@ def base_test_env(test_model, test_tokenizer):
         json.dump(
             {
                 "vocab_size": 1000,
-                "dim": 8,
-                "n_heads": 2,
-                "n_kv_heads": 1,
-                "dim_ffn": 16,
-                "max_len": 64,
-                "n_layers": 2,
-                "norm_eps": 1e-5,
+                "hidden_size": 8,
+                "num_attention_heads": 2,
+                "num_key_value_heads": 1,
+                "intermediate_size": 16,
+                "max_position_embeddings": 64,
+                "num_hidden_layers": 2,
+                "rms_norm_eps": 1e-5,
             },
             f,
         )

tests/inference/test_scheduler.py

@@ -14,11 +14,11 @@ def mock_model_and_tokenizer():
     """Create mock model and tokenizer."""
     mock_model = MagicMock()
     mock_model.config = MagicMock()
-    mock_model.config.n_kv_heads = 8
-    mock_model.config.n_heads = 8
-    mock_model.config.dim = 128
-    mock_model.config.n_layers = 2
-    mock_model.config.max_len = 100
+    mock_model.config.num_key_value_heads = 8
+    mock_model.config.num_attention_heads = 8
+    mock_model.config.hidden_size = 128
+    mock_model.config.num_hidden_layers = 2
+    mock_model.config.max_position_embeddings = 100
     mock_model.parameters.return_value = iter(
         [MagicMock(dtype=torch.float32, device=torch.device("cpu"))]
     )
@@ -213,13 +213,13 @@ def _make_real_scheduler(device):
 
     cfg = AutoRegressiveLMConfig(
         vocab_size=200,
-        dim=16,
-        n_heads=2,
-        n_kv_heads=1,
-        dim_ffn=32,
-        max_len=64,
-        n_layers=2,
-        norm_eps=1e-5,
+        hidden_size=16,
+        num_attention_heads=2,
+        num_key_value_heads=1,
+        intermediate_size=32,
+        max_position_embeddings=64,
+        num_hidden_layers=2,
+        rms_norm_eps=1e-5,
     )
     model = AutoRegressiveLM(cfg).to(device=device).eval()
     tokenizer = _Tok()

tests/module/test_encoder.py

@@ -12,13 +12,13 @@ from astrai.model.encoder import EmbeddingEncoder
 
 TINY_CONFIG = dict(
     vocab_size=128,
-    dim=8,
-    n_heads=2,
-    n_kv_heads=1,
-    dim_ffn=16,
-    max_len=64,
-    n_layers=2,
-    norm_eps=1e-5,
+    hidden_size=8,
+    num_attention_heads=2,
+    num_key_value_heads=1,
+    intermediate_size=16,
+    max_position_embeddings=64,
+    num_hidden_layers=2,
+    rms_norm_eps=1e-5,
 )
 
 _device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -42,7 +42,7 @@ def test_encoder_forward_pooling(pooling_type):
     with torch.no_grad():
         output = model(input_ids)
 
-    assert output.shape == (batch_size, TINY_CONFIG["dim"])
+    assert output.shape == (batch_size, TINY_CONFIG["hidden_size"])
     assert not torch.isnan(output).any()
 
 
@@ -60,7 +60,7 @@ def test_encoder_forward_with_padding():
     with torch.no_grad():
         output = model(input_ids, input_mask=input_mask)
 
-    assert output.shape == (batch_size, TINY_CONFIG["dim"])
+    assert output.shape == (batch_size, TINY_CONFIG["hidden_size"])
     assert not torch.isnan(output).any()
 
 
@@ -90,7 +90,7 @@ def test_encoder_from_transformer_checkpoint():
     model = _make_model()
     state_dict = model.state_dict()
     state_dict["lm_head.weight"] = torch.randn(
-        TINY_CONFIG["vocab_size"], TINY_CONFIG["dim"], device=_device
+        TINY_CONFIG["vocab_size"], TINY_CONFIG["hidden_size"], device=_device
     )
 
     new_model = _make_model()

tests/module/test_forward_configs.py

@@ -6,13 +6,13 @@ from astrai.model.transformer import AutoRegressiveLM
 
 TINY_CONFIG = dict(
     vocab_size=128,
-    dim=8,
-    n_heads=2,
-    n_kv_heads=1,
-    dim_ffn=16,
-    max_len=64,
-    n_layers=2,
-    norm_eps=1e-5,
+    hidden_size=8,
+    num_attention_heads=2,
+    num_key_value_heads=1,
+    intermediate_size=16,
+    max_position_embeddings=64,
+    num_hidden_layers=2,
+    rms_norm_eps=1e-5,
 )
 
 
@@ -58,8 +58,13 @@ CONFIGS = [
         id="gqa_qk_norm",
     ),
     pytest.param(
-        {**TINY_CONFIG, "attn_type": "gqa", "ffn_type": "mlp", "tie_weight": True},
-        id="gqa_tie_weight",
+        {
+            **TINY_CONFIG,
+            "attn_type": "gqa",
+            "ffn_type": "mlp",
+            "tie_word_embeddings": True,
+        },
+        id="gqa_tie_word_embeddings",
     ),
 ]
 
@@ -82,7 +87,11 @@ def test_model_forward(config_kwargs):
     assert "logits" in output
     assert "hidden_states" in output
     assert output["logits"].shape == (batch_size, seq_len, config.vocab_size)
-    assert output["hidden_states"].shape == (batch_size, seq_len, config.dim)
+    assert output["hidden_states"].shape == (
+        batch_size,
+        seq_len,
+        config.hidden_size,
+    )
     assert not torch.isnan(output["logits"]).any()
     assert not torch.isnan(output["hidden_states"]).any()
 

tests/module/test_lora.py

@@ -19,13 +19,13 @@ from astrai.model.components.lora import (
 
 MODEL_KWARGS = dict(
     vocab_size=1000,
-    dim=64,
-    n_heads=4,
-    n_kv_heads=2,
-    dim_ffn=128,
-    n_layers=2,
-    max_len=32,
-    norm_eps=1e-5,
+    hidden_size=64,
+    num_attention_heads=4,
+    num_key_value_heads=2,
+    intermediate_size=128,
+    num_hidden_layers=2,
+    max_position_embeddings=32,
+    rms_norm_eps=1e-5,
 )
 
 
@@ -192,7 +192,7 @@ def test_inject_lora_on_moe_model():
         n_routed_experts=4,
         n_shared_experts=1,
         n_activated_experts=2,
-        dim_ffn=32,
+        intermediate_size=32,
     )
     inject_lora(model, r=4, alpha=8, target_modules={"up", "gate", "down"})
     assert _get_lora_count(model) > 0

tests/module/test_tie_weight.py

@@ -17,13 +17,13 @@ def transformer_test_env():
 
     config = {
         "vocab_size": 1000,
-        "dim": 8,
-        "n_heads": 2,
-        "n_kv_heads": 1,
-        "dim_ffn": 16,
-        "max_len": 64,
-        "n_layers": 2,
-        "norm_eps": 1e-5,
+        "hidden_size": 8,
+        "num_attention_heads": 2,
+        "num_key_value_heads": 1,
+        "intermediate_size": 16,
+        "max_position_embeddings": 64,
+        "num_hidden_layers": 2,
+        "rms_norm_eps": 1e-5,
     }
 
     with open(config_path, "w") as f:
@@ -45,7 +45,7 @@ def test_tie_weight_init(transformer_test_env):
     config_data = transformer_test_env["config"].copy()
 
     # case 1: tie weight
-    config_data["tie_weight"] = True
+    config_data["tie_word_embeddings"] = True
 
     with open(config_path, "w") as f:
         json.dump(config_data, f)
@@ -63,7 +63,7 @@ def test_tie_weight_init(transformer_test_env):
     assert not torch.equal(model.lm_head.weight, original_weight)
 
     # case 2: not tie weight
-    config_data["tie_weight"] = False
+    config_data["tie_word_embeddings"] = False
 
     with open(config_path, "w") as f:
         json.dump(config_data, f)
@@ -88,7 +88,7 @@ def test_model_save_load_with_tie_weight(transformer_test_env):
     config_data = transformer_test_env["config"].copy()
 
     # case 1: tie weight
-    config_data["tie_weight"] = True
+    config_data["tie_word_embeddings"] = True
     config_path = os.path.join(test_dir, "config.json")
 
     with open(config_path, "w") as f:
@@ -108,7 +108,7 @@ def test_model_save_load_with_tie_weight(transformer_test_env):
     assert "lm_head.weight" not in model.state_dict()
 
     # case 2: not tie weight (form tie-weight state dict load)
-    config_data["tie_weight"] = False
+    config_data["tie_word_embeddings"] = False
     with open(config_path, "w") as f:
         json.dump(config_data, f)
 

tests/trainer/test_grpo_strategy.py

@@ -13,16 +13,16 @@ class _FakeExecutor:
         return model.state_dict()
 
 
-def _make_config(vocab_size=200, max_len=64):
+def _make_config(vocab_size=200, max_position_embeddings=64):
     return AutoRegressiveLMConfig(
         vocab_size=vocab_size,
-        dim=16,
-        n_heads=2,
-        n_kv_heads=1,
-        dim_ffn=32,
-        max_len=max_len,
-        n_layers=2,
-        norm_eps=1e-5,
+        hidden_size=16,
+        num_attention_heads=2,
+        num_key_value_heads=1,
+        intermediate_size=32,
+        max_position_embeddings=max_position_embeddings,
+        num_hidden_layers=2,
+        rms_norm_eps=1e-5,
     )
 
 

tests/trainer/test_online_strategy.py

@@ -33,16 +33,16 @@ class _FakeExecutor:
         return model.state_dict()
 
 
-def _make_config(vocab_size=200, max_len=64):
+def _make_config(vocab_size=200, max_position_embeddings=64):
     return AutoRegressiveLMConfig(
         vocab_size=vocab_size,
-        dim=16,
-        n_heads=2,
-        n_kv_heads=1,
-        dim_ffn=32,
-        max_len=max_len,
-        n_layers=2,
-        norm_eps=1e-5,
+        hidden_size=16,
+        num_attention_heads=2,
+        num_key_value_heads=1,
+        intermediate_size=32,
+        max_position_embeddings=max_position_embeddings,
+        num_hidden_layers=2,
+        rms_norm_eps=1e-5,
     )
 
 

tests/trainer/test_rollout.py

@@ -71,16 +71,16 @@ class ConstantRewardModel(BaseRewardModel):
         return torch.full((B, G), float(self.value))
 
 
-def _make_config(vocab_size=200, max_len=128):
+def _make_config(vocab_size=200, max_position_embeddings=128):
     return AutoRegressiveLMConfig(
         vocab_size=vocab_size,
-        dim=16,
-        n_heads=2,
-        n_kv_heads=1,
-        dim_ffn=32,
-        max_len=max_len,
-        n_layers=2,
-        norm_eps=1e-5,
+        hidden_size=16,
+        num_attention_heads=2,
+        num_key_value_heads=1,
+        intermediate_size=32,
+        max_position_embeddings=max_position_embeddings,
+        num_hidden_layers=2,
+        rms_norm_eps=1e-5,
     )
 
 
@@ -158,7 +158,7 @@ def _make_generator(device, **kw):
         model,
         tokenizer,
         max_batch_size=kw.get("max_batch_size", 8),
-        max_len=kw.get("max_len", 128),
+        max_len=kw.get("max_position_embeddings", 128),
     )
     generator = RolloutGenerator(
         scheduler=scheduler,
@@ -254,7 +254,7 @@ def _make_runner(device, **kw):
         group_size=kw.get("group_size", 2),
         max_tokens=kw.get("max_tokens", 8),
         max_batch_size=kw.get("max_batch_size", 8),
-        max_len=kw.get("max_len", 128),
+        max_len=kw.get("max_position_embeddings", 128),
     )
     rm = ConstantRewardModel(1.0)
     return (