feat: 新增NTK-Aware RoPE缩放支持

- RotaryEmbedding接受rope_scaling配置,自动计算scaled base - AutoRegressiveLMConfig和EncoderConfig新增rope_scaling字段
2026-05-25 21:22:07 +08:00
parent a4688021bf
commit 737585a32a
4 changed files with 28 additions and 4 deletions
@@ -59,7 +59,9 @@ class AutoRegressiveLM(AutoModel):
            else config.dim // config.n_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)
+        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.layers = nn.ModuleList(