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