LLM 上下文长度扩展：从 YARN 到 NTK-aware 插值

训练长度 4K 的模型在不同上下文长度下的表现:

┌──────────────┬──────────┬─────────────────────────┐
│ 上下文长度   │ 困惑度   │ Passkey Retrieval 准确率 │
├──────────────┼──────────┼─────────────────────────┤
│ 4K (训练内)  │ 5.82     │ 100%                    │
│ 8K (2x)      │ 7.31     │ 12%                     │
│ 16K (4x)     │ 12.44    │ 0%                      │
│ 32K (8x)     │ 28.71    │ 0%                      │
└──────────────┴──────────┴─────────────────────────┘

问题根源: RoPE 在超出训练范围后，旋转角度超出模型见过的范围

class PositionInterpolatedRoPE(nn.Module):
    """位置插值 RoPE"""
    def __init__(self, dim, max_train_len=4096, base=10000):
        super().__init__()
        self.dim = dim
        self.max_train_len = max_train_len
        self.base = base
        
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)
    
    def forward(self, x, seq_len, scaling_factor=1.0):
        """
        scaling_factor = L_train / L_new
        当 L_new > L_train 时, scaling_factor < 1
        """
        # 关键: 将位置索引缩放到训练范围内
        positions = torch.arange(seq_len, device=x.device, dtype=torch.float32)
        positions = positions / scaling_factor  # 压缩位置
        
        # 计算频率
        freqs = torch.outer(positions, self.inv_freq)  # (seq_len, dim/2)
        emb = torch.cat([freqs, freqs], dim=-1)
        
        cos = emb.cos().unsqueeze(0).unsqueeze(0)
        sin = emb.sin().unsqueeze(0).unsqueeze(0)
        
        return self.apply_rotary(x, cos, sin)
    
    def apply_rotary(self, x, cos, sin):
        d = x.shape[-1]
        x1, x2 = x[..., :d//2], x[..., d//2:]
        cos, sin = cos[..., :d//2], sin[..., :d//2]
        return torch.cat([x1*cos - x2*sin, x1*sin + x2*cos], dim=-1)

# PI 的微调流程
def pi_finetuning_pipeline(model, train_len=4096, target_len=16384):
    """
    位置插值微调流程
    """
    scaling_factor = train_len / target_len  # 0.25
    
    # 1. 修改模型的位置编码
    for layer in model.layers:
        layer.attention.rope.scaling_factor = scaling_factor
    
    # 2. 使用长文本数据微调
    # 关键: 数据应包含长文档，但不需要满长度
    training_config = {
        "max_seq_len": target_len,
        "batch_size": 1,  # 长序列需要小batch
        "gradient_accumulation": 32,
        "learning_rate": 2e-5,  # 比预训练低
        "steps": 1000,  # 通常 500-2000 步足够
        "data": "long_documents.jsonl",  # 长文档数据
        # 混合长短序列
        "mix_short_long": True,
        "short_ratio": 0.3,  # 30% 短序列保持原始能力
    }
    
    return training_config

class NTKAwareRoPE(nn.Module):
    """NTK-aware RoPE: 通过调整 base 实现高频不变"""
    def __init__(self, dim, base=10000, train_len=4096, target_len=16384):
        super().__init__()
        self.dim = dim
        
        # NTK-aware base 调整
        scale = target_len / train_len
        # 等价于: base_new = base * scale^(d/(d-2))
        self.base_new = base * (scale ** (dim / (dim - 2)))
        
        inv_freq = 1.0 / (self.base_new ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)
    
    def forward(self, x, seq_len):
        positions = torch.arange(seq_len, device=x.device, dtype=torch.float32)
        freqs = torch.outer(positions, self.inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        
        cos = emb.cos().unsqueeze(0).unsqueeze(0)
        sin = emb.sin().unsqueeze(0).unsqueeze(0)
        
        return self.apply_rotary(x, cos, sin)
    
    def apply_rotary(self, x, cos, sin):
        d = x.shape[-1]
        x1, x2 = x[..., :d//2], x[..., d//2:]
        cos, sin = cos[..., :d//2], sin[..., :d//2]
        return torch.cat([x1*cos - x2*sin, x1*sin + x2*cos], dim=-1)

def compare_frequencies(dim=128, base_old=10000, base_new=500000, train_len=4096, target_len=32768):
    """对比 PI 和 NTK 在不同频率上的行为"""
    
    # 原始频率
    inv_freq_old = 1.0 / (base_old ** (torch.arange(0, dim, 2).float() / dim))
    
    # PI: 统一缩放
    pi_scale = train_len / target_len
    inv_freq_pi = inv_freq_old * pi_scale  # 所有频率乘以相同因子
    
    # NTK: 调整 base
    ntk_scale = (target_len / train_len) ** (dim / (dim - 2))
    inv_freq_ntk = 1.0 / ((base_old * ntk_scale) ** (torch.arange(0, dim, 2).float() / dim))
    
    print("维度索引 | 原始频率     | PI频率       | NTK频率")
    print("-" * 60)
    for i in [0, 2, 8, 16, 32, 63]:  # 采样不同维度
        print(f"  {i:4d}  | {inv_freq_old[i]:.6f} | {inv_freq_pi[i]:.6f} | {inv_freq_ntk[i]:.6f}")
    
    # 高频(小索引): NTK ≈ 原始, PI 被缩放
    # 低频(大索引): NTK 被拉伸, PI 被缩放
    
# 输出示例:
# 维度索引 | 原始频率     | PI频率       | NTK频率
#   0    | 1.000000  | 0.125000   | 0.998754    ← 高频几乎不变!
#  16   | 0.021544  | 0.002693   | 0.021517    ← 中频轻微变化
#  63   | 0.000100  | 0.0000125  | 0.000794    ← 低频被拉伸!

频率维度划分 (从高频到低频):

┌─────────────┬─────────────────┬──────────────────┐
│   高频段     │    中频过渡段    │     低频段        │
│ (不变)       │  (平滑插值)      │   (完全外推)      │
├─────────────┼─────────────────┼──────────────────┤
│ dim 0 ~ r1  │  r1 ~ r2        │  r2 ~ dim/2      │
│ 保持原始     │ 线性插值          │  位置插值          │
└─────────────┴─────────────────┴──────────────────┘

r1 = floor(d * log(s/max_period_fast) / (2*log(base)))
r2 = floor(d * log(s/max_period_slow) / (2*log(base)))

class YaRNRoPE(nn.Module):
    """YaRN: 分频段处理的 RoPE 扩展"""
    def __init__(self, dim, base=10000, train_len=4096, target_len=32768,
                 beta_fast=32, beta_slow=1):
        super().__init__()
        self.dim = dim
        self.base = base
        self.train_len = train_len
        self.target_len = target_len
        self.scale = target_len / train_len
        self.beta_fast = beta_fast
        self.beta_slow = beta_slow
        
        # 计算分段边界
        self.r1 = self._find_correction_dim(beta_fast)
        self.r2 = self._find_correction_dim(beta_slow)
        
        # 计算 YaRN 的 inv_freq
        inv_freq = self._compute_yarn_inv_freq()
        self.register_buffer('inv_freq', inv_freq)
    
    def _find_correction_dim(self, num_rotations):
        """计算给定旋转次数对应的维度边界"""
        return (self.dim * math.log(self.train_len / (num_rotations * 2 * math.pi))
                / (2 * math.log(self.base)))
    
    def _compute_yarn_inv_freq(self):
        """计算 YaRN 的频率向量"""
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim))
        
        # 对每个频率维度应用不同的缩放
        for i in range(self.dim // 2):
            if i < self.r1:
                # 高频段: 保持不变
                pass
            elif i > self.r2:
                # 低频段: 完全外推 (相当于 PI)
                inv_freq[i] /= self.scale
            else:
                # 中频段: 线性插值
                t = (i - self.r1) / max(self.r2 - self.r1, 1)
                # 从 1.0 平滑过渡到 1/scale
                scale_factor = (1 - t) + t / self.scale
                inv_freq[i] /= scale_factor
        
        return inv_freq
    
    def forward(self, x, seq_len):
        positions = torch.arange(seq_len, device=x.device, dtype=torch.float32)
        freqs = torch.outer(positions, self.inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        
        cos = emb.cos().unsqueeze(0).unsqueeze(0)
        sin = emb.sin().unsqueeze(0).unsqueeze(0)
        
        return self.apply_rotary(x, cos, sin)
    
    def apply_rotary(self, x, cos, sin):
        d = x.shape[-1]
        x1, x2 = x[..., :d//2], x[..., d//2:]
        cos, sin = cos[..., :d//2], sin[..., :d//2]
        return torch.cat([x1*cos - x2*sin, x1*sin + x2*cos], dim=-1)

def yarn_attention_scale(scale, beta_fast=32, beta_slow=1):
    """
    YaRN 的注意力温度缩放
    当 scale > 1 时, log-scale 的注意力需要降温
    """
    if scale <= 1:
        return 1.0
    
    # 基于经验公式的缩放因子
    import math
    log_scale = math.log(scale) / math.log(2)  # log2(scale)
    
    # 0.1 是经验值, 控制缩放强度
    attention_scale = 0.1 * log_scale + 1.0
    
    return attention_scale

# 在注意力计算中应用:
# scores = (q @ k^T) / (sqrt(d_k) * attention_scale)
# 而非标准的 scores = (q @ k^T) / sqrt(d_k)

class LongRoPE:
    """LongRoPE: 搜索式非均匀 RoPE 缩放"""
    def __init__(self, dim, base=10000, train_len=4096, target_len=32768):
        self.dim = dim
        self.base = base
        self.train_len = train_len
        self.target_len = target_len
        self.scale = target_len / train_len
    
    def search_optimal_scaling(self, eval_fn, population=20, generations=100):
        """进化搜索最优的维度缩放向量"""
        import random
        
        # 初始化种群: 每个个体是一个 dim/2 维的缩放向量
        pop = []
        for _ in range(population):
            # 初始个体: 在 [1, scale] 范围内随机
            individual = [random.uniform(1.0, self.scale) 
                         for _ in range(self.dim // 2)]
            pop.append(individual)
        
        for gen in range(generations):
            # 评估适应度
            scored = [(ind, eval_fn(self._apply_scaling(ind))) for ind in pop]
            scored.sort(key=lambda x: x[1], reverse=True)
            
            # 保留 top 50%
            survivors = [ind for ind, _ in scored[:population // 2]]
            
            # 交叉 + 变异
            children = []
            for _ in range(population - len(survivors)):
                p1, p2 = random.sample(survivors, 2)
                # 均匀交叉
                child = [p1[i] if random.random() < 0.5 else p2[i] 
                        for i in range(self.dim // 2)]
                # 变异
                for i in range(len(child)):
                    if random.random() < 0.1:  # 10% 变异率
                        child[i] = random.uniform(1.0, self.scale)
                children.append(child)
            
            pop = survivors + children
        
        # 返回最优个体
        best = max(scored, key=lambda x: x[1])[0]
        return best
    
    def _apply_scaling(self, scaling_vec):
        """将缩放向量应用于 inv_freq"""
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim))
        inv_freq_scaled = inv_freq / torch.tensor(scaling_vec)
        return inv_freq_scaled

Stage 1: 搜索 2x 外推的缩放向量
  → 微调 400 步
  → 评估困惑度

Stage 2: 以 Stage 1 结果为基础, 搜索 4x 外推
  → 微调 400 步
  → 评估困惑度

Stage 3: (可选) 继续扩展到 8x, 16x

# HuggingFace Transformers 配置
from transformers import AutoConfig

# YaRN 配置 (Llama-3 风格)
config = AutoConfig.from_pretrained("meta-llama/Llama-3-8B-Instruct")
config.rope_scaling = {
    "type": "yarn",
    "factor": 4.0,  # 扩展倍数
    "original_max_position_embeddings": 8192,
    "beta_fast": 32,
    "beta_slow": 1,
}

# PI 配置
config_pi = AutoConfig.from_pretrained("meta-llama/Llama-3-8B-Instruct")
config_pi.rope_scaling = {
    "type": "linear",  # Position Interpolation
    "factor": 4.0,
}

# NTK-aware 配置 (部分框架支持)
config_ntk = AutoConfig.from_pretrained("meta-llama/Llama-3-8B-Instruct")
config_ntk.rope_scaling = {
    "type": "ntk",
    "factor": 4.0,
}

# KV Cache 显存计算公式
def kv_cache_memory(num_layers, num_kv_heads, d_k, seq_len, batch_size, dtype_bytes=2):
    """
    计算 KV Cache 的显存占用 (FP16/BF16)
    """
    # 每层: K cache + V cache
    per_layer = 2 * batch_size * num_kv_heads * seq_len * d_k * dtype_bytes
    total = num_layers * per_layer
    return total / (1024**3)  # GB

# Llama-3-8B 的 KV Cache 对比
# num_layers=32, num_kv_heads=8, d_k=128
for seq_len in [8192, 32768, 131072, 1048576]:
    mem = kv_cache_memory(32, 8, 128, seq_len, 1)
    print(f"Seq Len {seq_len:>7d}: KV Cache = {mem:.2f} GB")

# 输出:
# Seq Len    8192: KV Cache = 0.50 GB
# Seq Len   32768: KV Cache = 2.00 GB
# Seq Len  131072: KV Cache = 8.00 GB
# Seq Len 1048576: KV Cache = 64.00 GB  ← 需要量化或分页!