ttt_cross_layer.py

import torch
import torch.nn as nn
import torch.nn.functional as F


def scan(f, init, xs, out, checkpoint_group=0):
    """
    模拟JAX中的lax.scan函数，用于序列化处理数据。
    
    参数:
        f: 处理函数，接收(carry, x)作为输入，返回(new_carry, y)
        init: 初始状态值
        xs: 输入序列，可以是字典或列表
        out: 输出结果的存储张量
        checkpoint_group: 梯度检查点分组数量，用于节省内存
    
    返回:
        carry: 最终的状态值
        out: 填充好的输出张量
    """
    # 初始化状态值
    carry = init
    
    # 确定输入序列的长度
    if isinstance(xs, dict):
        # 如果输入是字典，取第一个键对应值的长度
        num_items = len(next(iter(xs.values())))
    else:
        # 如果输入是列表，取第一个元素的长度
        num_items = len(xs[0])

    def scan_fn(carry, i_start, i_end):
        """内部扫描函数，处理从i_start到i_end的元素"""
        for i in range(i_start, i_end):
            # 提取当前位置的输入
            if isinstance(xs, dict):
                # 字典情况：创建包含每个键在位置i处值的新字典
                x = {key: tensor[i] for key, tensor in xs.items()}
            else:
                # 列表情况：创建包含每个列表在位置i处值的新列表
                x = [x[i] for x in xs]
            
            # 调用处理函数f，获取新的状态和输出
            carry, y = f(carry, x)
            
            # 将输出存储到结果张量中
            out[i] = y
        
        # 返回最终状态
        return carry

    # 根据checkpoint_group决定是否使用梯度检查点
    if checkpoint_group > 0:
        # 计算每个检查点组包含的元素数量
        ckpt_every_n = num_items // checkpoint_group
        
        # 按组处理数据
        for k in range(0, num_items, ckpt_every_n):
            # 使用torch.utils.checkpoint节省内存
            carry = torch.utils.checkpoint.checkpoint(
                scan_fn, carry, k, min(k + ckpt_every_n, num_items), use_reentrant=False
            )
    else:
        # 不使用检查点，直接处理所有数据
        carry = scan_fn(carry, 0, num_items)

    # 返回最终状态和填充好的输出张量
    return carry, out

def ln_fwd(x, gamma, beta, eps=1e-6):
    "Batch forward for LayerNorm."

    # Mean and variance computation
    mu = x.mean(dim=-1, keepdim=True)
    var = x.var(dim=-1, keepdim=True, unbiased=False)

    # Normalization
    std = torch.sqrt(var + eps)
    x_hat = (x - mu) / std

    # Scale and shift
    y = gamma * x_hat + beta

    return y

def ln_fused_l2_bwd(x, l2_target, gamma, beta, eps=1e-6):
    """
    层归一化(LayerNorm)与L2损失融合的反向传播函数。
    
    这个函数执行两个操作：
    1. 前向传播：对输入x进行层归一化，得到输出y
    2. 反向传播：计算L2损失(y - l2_target)对输入x的梯度
    
    参数:
        x: 输入张量
        l2_target: L2损失的目标值
        gamma: 层归一化的缩放参数
        beta: 层归一化的偏移参数
        eps: 数值稳定性的小常数
        
    返回:
        z: 损失对输入x的梯度
    """
    D = x.shape[-1]  # 获取特征维度

    # 计算均值和方差
    mu = x.mean(dim=-1, keepdim=True)  # 沿特征维度计算均值
    var = x.var(dim=-1, keepdim=True, unbiased=False)  # 计算方差

    # 归一化处理
    std = torch.sqrt(var + eps)  # 计算标准差
    x_hat = (x - mu) / std  # 标准化输入

    # 缩放和偏移
    y = gamma * x_hat + beta  # 层归一化的输出

    # 计算梯度
    grad_output = y - l2_target  # L2损失的梯度
    grad_x_hat = grad_output * gamma  # 对标准化输入的梯度
    
    # 完整的反向传播公式，考虑了归一化操作的链式法则
    z = (
        (1.0 / D)
        * (
            D * grad_x_hat
            - grad_x_hat.sum(dim=-1, keepdim=True)  # 均值项的梯度贡献
            - x_hat * (grad_x_hat * x_hat).sum(dim=-1, keepdim=True)  # 方差项的梯度贡献
        )
        / std  # 除以标准差完成梯度计算
    )

    return z

from torch.autograd import Function
class MyLinearFunction(Function):
    @staticmethod
    def forward(ctx, input, weight, bias):
        """
        正向计算： y = x * W^T + b
        参数：
            ctx    ：上下文对象，用于保存反向传播时需要的信息。
            input  ：输入 tensor, 尺寸为 (N, in_features)
            weight ：权重 tensor, 尺寸为 (out_features, in_features)
            bias   ：偏置 tensor, 尺寸为 (out_features)
        返回：
            输出 tensor, 尺寸为 (N, out_features)
        """
        # 保存必要的中间变量，供 backward 时使用
        ctx.save_for_backward(input, weight, bias)
        
        # 计算输出
        output = input.matmul(weight.t()) + bias
        return output

    @staticmethod
    def backward(ctx, grad_output):
        """
        反向传播：计算正向计算中各个输入的梯度。
        参数：
            grad_output：从上层传回来的梯度，形状与 forward 的输出相同 (N, out_features)
        返回：
            grad_input  ：关于 input 的梯度，形状 (N, in_features)
            grad_weight ：关于 weight 的梯度，形状 (out_features, in_features)
            grad_bias   ：关于 bias 的梯度，形状 (out_features)
        """
        # 从上下文中取出保存的变量
        input, weight, bias = ctx.saved_tensors
        
        # 链式法则：已知 output = input.matmul(weight.t()) + bias
        # 关于 input 的梯度：
        # ∂L/∂input = ∂L/∂output * ∂output/∂input = grad_output.matmul(weight)
        grad_input = grad_output.matmul(weight)
        
        # 关于 weight 的梯度：
        # ∂L/∂weight = ∂L/∂output^T * ∂output/∂weight
        # 注意到 output 对 weight 的导数为 input 的转置，此处：
        # grad_weight 的计算通常为：grad_output^T.matmul(input)
        grad_weight = grad_output.t().matmul(input)
        
        # 关于 bias 的梯度：
        # 因为 output = ... + bias，因此每个 bias 项对应所有样本的梯度和
        grad_bias = grad_output.sum(dim=0)
        
        # 注意：返回的梯度顺序必须与 forward 中参数的顺序一致
        return grad_input, grad_weight, grad_bias

class TTT_Cross_Layer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.input_size = config.concept_dim   # 128
        self.concept_dim = config.concept_dim  # 128
        # self.linear = nn.Linear(self.input_size, self.hidden_size)
        # self.ln = nn.LayerNorm(self.hidden_size)

        # self.logit_dim = 32
        self.logit_dim = config.logit_dim

        self.weight_linear = nn.Parameter(torch.empty(self.concept_dim, self.input_size, self.logit_dim))
        self.weight_ln = nn.Parameter(torch.empty(self.concept_dim, self.logit_dim))
        self.bias_linear = nn.Parameter(torch.empty(self.concept_dim, self.logit_dim))
        self.bias_ln = nn.Parameter(torch.empty(self.concept_dim, self.logit_dim))

        # self.weight_linear_tmp = torch.empty_like(self.weight_linear)
        # self.weight_ln_tmp = torch.empty_like(self.weight_ln)
        # self.bias_linear_tmp = torch.empty_like(self.bias_linear)
        # self.bias_ln_tmp = torch.empty_like(self.bias_ln)
        
        self.config = config
        self.init_weights()
    # def init_tmp_weights(self):
    #     weight_linear_tmp = self.weight_linear.clone().to(self.weight_linear.device).to(self.weight_linear.dtype)
    #     weight_ln_tmp = self.weight_ln.clone().to(self.weight_ln.device).to(self.weight_ln.dtype)
    #     bias_linear_tmp = self.bias_linear.clone().to(self.bias_linear.device).to(self.bias_linear.dtype)
    #     bias_ln_tmp = self.bias_ln.clone().to(self.bias_ln.device).to(self.bias_ln.dtype)
    #     params = {
    #         'weight_linear_tmp': weight_linear_tmp,
    #         'weight_ln_tmp': weight_ln_tmp,
    #         'bias_linear_tmp': bias_linear_tmp,
    #         'bias_ln_tmp': bias_ln_tmp
    #     }
    #     return params
    
    def init_params_as_logits(self, batch_size, sequence_length):
        weight_linear_tmp = torch.ones(batch_size, sequence_length, self.logit_dim).to(self.weight_linear.device).to(self.weight_linear.dtype)
        weight_ln_tmp = torch.ones(batch_size, sequence_length, self.logit_dim).to(self.weight_linear.device).to(self.weight_linear.dtype)
        bias_linear_tmp = torch.ones(batch_size, sequence_length, self.logit_dim).to(self.weight_linear.device).to(self.weight_linear.dtype)
        bias_ln_tmp = torch.ones(batch_size, sequence_length, self.logit_dim).to(self.weight_linear.device).to(self.weight_linear.dtype)
        
        params = {
            'weight_linear_tmp': weight_linear_tmp,
            'weight_ln_tmp': weight_ln_tmp,
            'bias_linear_tmp': bias_linear_tmp,
            'bias_ln_tmp': bias_ln_tmp
        }
        return params

    def init_weights(self):
        # torch.manual_seed(42)  # 固定随机种子可能导致可预测性
        nn.init.normal_(self.weight_linear, mean=0.0, std=self.config.initializer_range)
        nn.init._no_grad_fill_(self.weight_ln, 1.0 / self.logit_dim)
        # nn.init.zeros_(self.bias_linear)
        # nn.init.zeros_(self.bias_ln)
        nn.init.normal_(self.bias_linear, mean=0.0, std=self.config.initializer_range / self.logit_dim)
        nn.init.normal_(self.bias_linear, mean=0.0, std=self.config.initializer_range / self.logit_dim)

    def get_weight_per_token(self, params):
        
        weight_linear_tmp = torch.einsum('iol,bsl->bsio', self.weight_linear, params['weight_linear_tmp'])
        weight_ln_tmp = torch.einsum('ol,bsl->bso', self.weight_ln, params['weight_ln_tmp'])
        bias_linear_tmp = torch.einsum('ol,bsl->bso', self.bias_linear, params['bias_linear_tmp'])
        bias_ln_tmp = torch.einsum('ol,bsl->bso', self.bias_ln, params['bias_ln_tmp'])

        return weight_linear_tmp, weight_ln_tmp, bias_linear_tmp, bias_ln_tmp

    def learn(self, k, v, params, lr_linear=1, lr_ln=1):
        # k和v形状: [batch_size, length, channel_dim]
        # batch_size, seq_length, channel_dim = k.shape
        # weight_linear_tmp = params['weight_linear_tmp']
        # weight_ln_tmp = params['weight_ln_tmp']
        # bias_linear_tmp = params['bias_linear_tmp']
        # bias_ln_tmp = params['bias_ln_tmp']
        weight_linear_tmp, weight_ln_tmp, bias_linear_tmp, bias_ln_tmp = self.get_weight_per_token(params)
        # 1. 将输入重塑为二维以进行预测
        # k_reshaped = k.reshape(-1, channel_dim)  # [batch_size*length, channel_dim]
        
        # output_reshaped = self.predict(k_reshaped, params)  # [batch_size*length, channel_dim]
        # z = F.linear(k_reshaped, params['weight_linear_tmp'], params['bias_linear_tmp'])
        # mu = z.mean(dim=-1, keepdim=True)
        # var = z.var(dim=-1, keepdim=True, unbiased=False)

        z = torch.einsum('bsi,bsio->bso', k, weight_linear_tmp) + bias_linear_tmp
        mu = z.mean(dim=-1, keepdim=True)
        var = z.var(dim=-1, keepdim=True, unbiased=False)

        # Normalization
        eps = 1e-6
        std = torch.sqrt(var + eps)
        z_hat = (z - mu) / std     
        # output_reshaped = params['weight_ln_tmp'] * z_hat + params['bias_ln_tmp'] + k
        output_reshaped = weight_ln_tmp * z_hat + bias_ln_tmp + k

        # # 计算误差
        # v_reshaped = v.reshape(-1, channel_dim)
        # error_reshaped = output_reshaped - v_reshaped  # [batch_size*length, channel_dim]
        error_reshaped = output_reshaped - v
        # 计算层归一化梯度
        # 层归一化参数更新
        # ln_rate = learning_rate * 0.1  # 降低LN学习率        
        grad_weight_ln_temp = error_reshaped * z_hat
        # grad_weight_ln = grad_weight_ln_temp.mean(dim=0) # 
        # weight_ln_tmp = weight_ln_tmp - ln_rate * grad_weight_ln # sequence length, channel_dim
        grad_weight_ln = grad_weight_ln_temp
        # batch_size, sequence length, logit_dim
        params0 = params['weight_ln_tmp'] - lr_ln * torch.einsum('ol,bso->bsl', self.weight_ln, grad_weight_ln)
        
        # bias_update = ln_rate * error_reshaped # .mean(dim=0)
        # bias_ln_tmp = bias_ln_tmp - bias_update # batch_size, sequence length, concept_dim
        grad_bias_ln = error_reshaped
        params1 = params['bias_ln_tmp'] - lr_ln * torch.einsum('ol,bso->bsl', self.bias_ln, grad_bias_ln)

        # 线性层权重梯度: [out_dim, in_dim]
        # grad_linear_temp = error_reshaped - error_reshaped.mean(dim=-1, keepdim=True) - z_hat * grad_weight_ln_temp.mean(dim=-1, keepdim=True)
        grad_linear = weight_ln_tmp * error_reshaped / std # batch_size, sequence length, concept_dim
        # grad_weight_linear = grad_linear.t() @ k  # [channel_dim, channel_dim]
        grad_weight_linear = torch.einsum('bsi,bso->bsio', k, grad_linear)
        # 应用梯度 (避免使用原地操作 -=)
        # weight_linear_tmp = weight_linear_tmp - learning_rate * grad_weight_linear.mean(dim=0)
        params2 = params['weight_linear_tmp'] - lr_linear * torch.einsum('iol,bsio->bsl', self.weight_linear, grad_weight_linear)
        # 更新偏置(如果存在) (避免使用原地操作 -=)
        grad_b = grad_linear #.mean(dim=0)  # [channel_dim]
        # bias_linear_tmp = bias_linear_tmp - learning_rate * grad_b
        params3 = params['bias_linear_tmp'] - lr_linear * torch.einsum('ol,bso->bsl', self.bias_linear, grad_b)
        
        params_new = {
            'weight_linear_tmp': params2,
            'weight_ln_tmp': params0,
            'bias_linear_tmp': params3,
            'bias_ln_tmp': params1
        }

        return params_new

    def predict(self, q, params):
        weight_linear_tmp, weight_ln_tmp, bias_linear_tmp, bias_ln_tmp = self.get_weight_per_token(params)
        z = torch.einsum('bsi,bsio->bso', q, weight_linear_tmp) + bias_linear_tmp
        mu = z.mean(dim=-1, keepdim=True)
        var = z.var(dim=-1, keepdim=True, unbiased=False)

        # Normalization
        eps = 1e-6
        std = torch.sqrt(var + eps)
        z_hat = (z - mu) / std     
        # output_reshaped = params['weight_ln_tmp'] * z_hat + params['bias_ln_tmp'] + k
        output = weight_ln_tmp * z_hat + bias_ln_tmp + q

        return output