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

astrai/model/components/decoder_block.py

@@ -14,10 +14,18 @@ class DecoderBlock(nn.Module):
     def __init__(self, config, layer_id: int):
         super().__init__()
         cfg = asdict(config)
-        cfg["down_init_std"] = 0.02 / (2 * config.n_layers) ** 0.5
+        cfg.update(
+            dim=config.hidden_size,
+            dim_ffn=config.intermediate_size,
+            n_layers=config.num_hidden_layers,
+            n_heads=config.num_attention_heads,
+            n_kv_heads=config.num_key_value_heads,
+            norm_eps=config.rms_norm_eps,
+            down_init_std=0.02 / (2 * config.num_hidden_layers) ** 0.5,
+        )
         self.attention = AttnFactory.create(config.attn_type, **cfg, layer_id=layer_id)
-        self.input_norm = RMSNorm(config.dim, config.norm_eps)
-        self.post_attention_norm = RMSNorm(config.dim, config.norm_eps)
+        self.input_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.post_attention_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
         self.mlp = FFNFactory.create(config.ffn_type, **cfg)
 
     def forward(

astrai/model/encoder.py

@@ -18,20 +18,28 @@ class EmbeddingEncoder(AutoModel):
     def __init__(self, config: EncoderConfig):
         super().__init__(config)
         self.config = config
-        rope_dim = config.dim // config.n_heads
+        rope_dim = config.hidden_size // config.num_attention_heads
         rope_base = config.rope_theta if config.rope_theta is not None else 10000
         self.rotary_embedding = RotaryEmbedding(
-            rope_dim, config.max_len, rope_base, rope_scaling=config.rope_scaling
+            rope_dim,
+            config.max_position_embeddings,
+            rope_base,
+            rope_scaling=config.rope_scaling,
         )
         self.embed_tokens = Embedding(
-            config.vocab_size, config.dim, neftune_alpha=config.neftune_alpha
+            config.vocab_size,
+            config.hidden_size,
+            neftune_alpha=config.neftune_alpha,
         )
 
         self.layers = nn.ModuleList(
-            [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)]
+            [
+                DecoderBlock(config, layer_id)
+                for layer_id in range(config.num_hidden_layers)
+            ]
         )
 
-        self.norm = RMSNorm(config.dim, config.norm_eps)
+        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
 
         self.pooling_type = config.pooling_type or "mean"
         self.normalize_embeddings = config.normalize_embeddings or False

astrai/model/transformer.py

@@ -36,24 +36,32 @@ class AutoRegressiveLM(AutoModel):
         rope_dim = (
             config.qk_rope_head_dim
             if config.attn_type == "mla"
-            else config.dim // config.n_heads
+            else config.hidden_size // config.num_attention_heads
         )
         rope_base = config.rope_theta if config.rope_theta is not None else 10000
         self.rotary_embedding = RotaryEmbedding(
-            rope_dim, config.max_len, rope_base, rope_scaling=config.rope_scaling
+            rope_dim,
+            config.max_position_embeddings,
+            rope_base,
+            rope_scaling=config.rope_scaling,
         )
         self.embed_tokens = Embedding(
-            config.vocab_size, config.dim, neftune_alpha=config.neftune_alpha
+            config.vocab_size,
+            config.hidden_size,
+            neftune_alpha=config.neftune_alpha,
         )
 
         self.layers = nn.ModuleList(
-            [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)]
+            [
+                DecoderBlock(config, layer_id)
+                for layer_id in range(config.num_hidden_layers)
+            ]
         )
 
-        self.norm = RMSNorm(config.dim, config.norm_eps)
-        self.lm_head = Linear(config.dim, config.vocab_size)
+        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.lm_head = Linear(config.hidden_size, config.vocab_size)
 
-        if self.config.tie_weight is True:
+        if self.config.tie_word_embeddings is True:
             self.lm_head.weight = self.embed_tokens.weight
 
         self.apply(self._init_weights)
@@ -68,7 +76,7 @@ class AutoRegressiveLM(AutoModel):
 
         state_dict = dict(state_dict)
 
-        if self.config.tie_weight is True:
+        if self.config.tie_word_embeddings is True:
             # same tensor for embed and lm_head
             if embed_key in state_dict:
                 state_dict[lm_head_key] = state_dict[embed_key]
@@ -84,7 +92,7 @@ class AutoRegressiveLM(AutoModel):
             destination=destination, prefix=prefix, keep_vars=keep_vars
         )
 
-        if self.config.tie_weight is True:
+        if self.config.tie_word_embeddings is True:
             lm_head_key = prefix + "lm_head.weight"
             if lm_head_key in state_dict:
                 del state_dict[lm_head_key]