refactor : neftune_alpha 在 Embedding 构造时传入，由模型配置链路负责

- BaseModelConfig 添加 neftune_alpha 字段 (默认 0.0)
- Embedding.__init__ 接受 neftune_alpha 参数，不再外部 set
- AutoRegressiveLM / EmbeddingEncoder 从 config 传入 neftune_alpha
- train.py 将 CLI 参数注入 config 后再创建模型
- TrainContextBuilder 移除 neftune 设置（不再是其职责）

astrai/config/model_config.py

@@ -20,6 +20,7 @@ class BaseModelConfig(BaseConfig):
     """Base config with ``model_type`` dispatch and file I/O."""
 
     model_type: Optional[str] = None
+    neftune_alpha: float = 0.0
 
 
 @dataclass

astrai/model/components/embedding.py

@@ -7,10 +7,13 @@ from torch import Tensor
 
 
 class Embedding(nn.Module):
-    def __init__(self, vocab_size: int, embedding_dim: int):
+    def __init__(self, vocab_size: int, embedding_dim: int, neftune_alpha: float = 0.0):
         super().__init__()
         self.weight = nn.Parameter(torch.empty((vocab_size, embedding_dim)))
-        self.neftune_noise_alpha = 0.0
+        self.neftune_noise_alpha = neftune_alpha
+
+    def set_neftune_alpha(self, alpha: float):
+        self.neftune_noise_alpha = alpha
 
     def reset_parameters(self):
         nn.init.normal_(self.weight, mean=0.0, std=0.02)

astrai/model/encoder.py

@@ -23,7 +23,9 @@ class EmbeddingEncoder(AutoModel):
         self.rotary_embedding = RotaryEmbedding(
             rope_dim, config.max_len, rope_base, rope_scaling=config.rope_scaling
         )
-        self.embed_tokens = Embedding(config.vocab_size, config.dim)
+        self.embed_tokens = Embedding(
+            config.vocab_size, config.dim, neftune_alpha=config.neftune_alpha
+        )
 
         self.layers = nn.ModuleList(
             [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)]

astrai/model/transformer.py

@@ -59,7 +59,9 @@ class AutoRegressiveLM(AutoModel):
         self.rotary_embedding = RotaryEmbedding(
             rope_dim, config.max_len, rope_base, rope_scaling=config.rope_scaling
         )
-        self.embed_tokens = Embedding(config.vocab_size, config.dim)
+        self.embed_tokens = Embedding(
+            config.vocab_size, config.dim, neftune_alpha=config.neftune_alpha
+        )
 
         self.layers = nn.ModuleList(
             [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)]

astrai/trainer/train_context.py

@@ -63,7 +63,6 @@ class TrainContextBuilder:
 
         model = cfg.model_fn()
         model = model.to(device=device)
-        model.embed_tokens.neftune_noise_alpha = cfg.neftune_alpha
 
         model_config = {}
         if self._resume_dir:

scripts/tools/train.py

@@ -309,6 +309,7 @@ def train(
     # Load config
     config_path = os.path.join(param_path, "config.json")
     config = AutoRegressiveLMConfig.from_file(config_path)
+    config.neftune_alpha = neftune_alpha
 
     if window_size is None:
         window_size = config.max_len