From 39985840c724884eb3b0994884f0c96063eb56d6 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Fri, 19 Jun 2026 14:23:27 +0800
Subject: [PATCH] =?UTF-8?q?refactor=20:=20neftune=5Falpha=20=E5=9C=A8=20Em?=
 =?UTF-8?q?bedding=20=E6=9E=84=E9=80=A0=E6=97=B6=E4=BC=A0=E5=85=A5?=
 =?UTF-8?q?=EF=BC=8C=E7=94=B1=E6=A8=A1=E5=9E=8B=E9=85=8D=E7=BD=AE=E9=93=BE?=
 =?UTF-8?q?=E8=B7=AF=E8=B4=9F=E8=B4=A3?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 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        | 1 +
 astrai/model/components/embedding.py | 7 +++++--
 astrai/model/encoder.py              | 4 +++-
 astrai/model/transformer.py          | 4 +++-
 astrai/trainer/train_context.py      | 1 -
 scripts/tools/train.py               | 1 +
 6 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/astrai/config/model_config.py b/astrai/config/model_config.py
index d0c1cf2..d9b8e5c 100644
--- a/astrai/config/model_config.py
+++ b/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
diff --git a/astrai/model/components/embedding.py b/astrai/model/components/embedding.py
index f8f4551..4d400fb 100644
--- a/astrai/model/components/embedding.py
+++ b/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)
diff --git a/astrai/model/encoder.py b/astrai/model/encoder.py
index 1ea1b2c..4b83652 100644
--- a/astrai/model/encoder.py
+++ b/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)]
diff --git a/astrai/model/transformer.py b/astrai/model/transformer.py
index 44655f8..9810545 100644
--- a/astrai/model/transformer.py
+++ b/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)]
diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py
index 031ab51..c7c7877 100644
--- a/astrai/trainer/train_context.py
+++ b/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:
diff --git a/scripts/tools/train.py b/scripts/tools/train.py
index e845596..999ff0a 100644
--- a/scripts/tools/train.py
+++ b/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