ViperEkura
39985840c7
- BaseModelConfig 添加 neftune_alpha 字段 (默认 0.0) - Embedding.__init__ 接受 neftune_alpha 参数，不再外部 set - AutoRegressiveLM / EmbeddingEncoder 从 config 传入 neftune_alpha - train.py 将 CLI 参数注入 config 后再创建模型 - TrainContextBuilder 移除 neftune 设置（不再是其职责）
27 lines
843 B
Python
27 lines
843 B
Python
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torch import Tensor
|
|
|
|
|
|
class Embedding(nn.Module):
|
|
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 = 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)
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
out = F.embedding(x, self.weight)
|
|
if self.training and self.neftune_noise_alpha > 0.0:
|
|
eps = self.neftune_noise_alpha / math.sqrt(out.size(1))
|
|
out = out + eps * torch.randn_like(out)
|
|
return out
|