feat(model): 添加QK归一化和门控注意力支持

tests/conftest.py

@@ -83,19 +83,19 @@ def base_test_env(request: pytest.FixtureRequest):
     n_dim_choices = [8, 16, 32]
     n_head_choices = [2, 4]
     
-    n_dim = int(np.random.choice(n_dim_choices))
-    n_head = int(np.random.choice(n_head_choices))
-    n_kvhead = n_head // 2
-    d_ffn = n_dim * 2
+    dim = int(np.random.choice(n_dim_choices))
+    n_heads = int(np.random.choice(n_head_choices))
+    n_kv_heads = n_heads // 2
+    dim_ffn = dim * 2
 
     config = {
         "vocab_size": 1000,
-        "n_dim": n_dim,
-        "n_head": n_head,
-        "n_kvhead": n_kvhead,
-        "d_ffn": d_ffn,
-        "m_len": 1024,
-        "n_layer": 4,
+        "dim": dim,
+        "n_heads": n_heads,
+        "n_kv_heads": n_kv_heads,
+        "dim_ffn": dim_ffn,
+        "max_len": 1024,
+        "n_layers": 4,
         "norm_eps": 1e-5
     }
     

tests/test_module.py

@@ -22,12 +22,12 @@ def test_env(request: pytest.FixtureRequest):
     
     config = {
         "vocab_size": 1000,
-        "n_dim": 128,
-        "n_head": 4,
-        "n_kvhead": 2,
-        "d_ffn": 256,
-        "m_len": 64,
-        "n_layer": 2,
+        "dim": 128,
+        "n_heads": 4,
+        "n_kv_heads": 2,
+        "dim_ffn": 256,
+        "max_len": 64,
+        "n_layers": 2,
         "norm_eps": 1e-5
     }
     with open(config_path, 'w') as f:
@@ -64,9 +64,9 @@ def test_model_parameter(test_env):
 def test_transformer(test_env):
     model = test_env["model"]
     input_ids = torch.randint(0, test_env["transformer_config"].vocab_size, 
-                              (4, test_env["transformer_config"].m_len))
+                              (4, test_env["transformer_config"].max_len))
     output_logits = model(input_ids)["logits"]
-    target_shape = (4, test_env["transformer_config"].m_len, test_env["transformer_config"].vocab_size)
+    target_shape = (4, test_env["transformer_config"].max_len, test_env["transformer_config"].vocab_size)
     assert output_logits.shape == target_shape
     
 # generator
@@ -80,7 +80,7 @@ def test_embedding_encoder_core(test_env):
     
     single_emb = encoder.encode("测试文本")
     assert isinstance(single_emb, torch.Tensor)
-    assert single_emb.shape[-1] == test_env["transformer_config"].n_dim
+    assert single_emb.shape[-1] == test_env["transformer_config"].dim
     
 
     batch_emb = encoder.encode(["测试1", "测试2"])

tests/test_tie_weight.py

@@ -16,12 +16,12 @@ def transformer_test_env():
     
     config = {
         "vocab_size": 1000,
-        "n_dim": 128,
-        "n_head": 4,
-        "n_kvhead": 2,
-        "d_ffn": 256,
-        "m_len": 64,
-        "n_layer": 2,
+        "dim": 128,
+        "n_heads": 4,
+        "n_kv_heads": 2,
+        "dim_ffn": 256,
+        "max_len": 64,
+        "n_layers": 2,
         "norm_eps": 1e-5
     }