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/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)