第18章：AI前沿技术趋势

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple, Optional
import numpy as np

class Expert(nn.Module):
    """单个专家网络 - 标准的FFN"""
    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [batch_size, seq_len, d_model]
        return self.w2(self.dropout(F.gelu(self.w1(x))))

class TopKGate(nn.Module):
    """Top-K门控网络"""
    def __init__(
        self,
        d_model: int,
        num_experts: int,
        top_k: int = 2,
        capacity_factor: float = 1.25,
        noisy_gate: bool = True
    ):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.capacity_factor = capacity_factor
        self.noisy_gate = noisy_gate

        # 门控权重
        self.w_gate = nn.Linear(d_model, num_experts, bias=False)

        # 噪声参数（用于训练时的探索）
        if noisy_gate:
            self.w_noise = nn.Linear(d_model, num_experts, bias=False)

    def forward(
        self,
        x: torch.Tensor,
        train: bool = True
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Args:
            x: [batch_size, seq_len, d_model]

        Returns:
            gates: [batch_size, seq_len, top_k] - 门控权重
            indices: [batch_size, seq_len, top_k] - 选中的专家索引
            load: [num_experts] - 每个专家的负载
        """
        # 计算门控logits
        logits = self.w_gate(x)  # [batch_size, seq_len, num_experts]

        # 添加噪声（训练时）
        if self.noisy_gate and train:
            noise_logits = self.w_noise(x)
            noise = torch.randn_like(logits) * F.softplus(noise_logits)
            logits = logits + noise

        # Top-K选择
        gates = F.softmax(logits, dim=-1)
        top_k_gates, top_k_indices = torch.topk(gates, self.top_k, dim=-1)

        # 重新归一化
        top_k_gates = top_k_gates / (top_k_gates.sum(dim=-1, keepdim=True) + 1e-8)

        # 计算负载（用于负载均衡）
        load = self._calculate_load(top_k_indices)

        return top_k_gates, top_k_indices, load

    def _calculate_load(self, indices: torch.Tensor) -> torch.Tensor:
        """计算每个专家的负载"""
        batch_size, seq_len, top_k = indices.shape

        # 统计每个专家被选中的次数
        load = torch.zeros(
            self.num_experts,
            device=indices.device,
            dtype=torch.float32
        )

        indices_flat = indices.view(-1)
        load.scatter_add_(
            0,
            indices_flat,
            torch.ones_like(indices_flat, dtype=torch.float32)
        )

        # 归一化
        load = load / (batch_size * seq_len * top_k)
        return load

class MoELayer(nn.Module):
    """完整的MoE层"""
    def __init__(
        self,
        d_model: int,
        d_ff: int,
        num_experts: int = 8,
        top_k: int = 2,
        dropout: float = 0.1,
        capacity_factor: float = 1.25,
        balance_loss_weight: float = 0.01
    ):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.balance_loss_weight = balance_loss_weight

        # 创建专家网络
        self.experts = nn.ModuleList([
            Expert(d_model, d_ff, dropout)
            for _ in range(num_experts)
        ])

        # 门控网络
        self.gate = TopKGate(d_model, num_experts, top_k, capacity_factor)

    def forward(
        self,
        x: torch.Tensor,
        train: bool = True
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            x: [batch_size, seq_len, d_model]

        Returns:
            output: [batch_size, seq_len, d_model]
            aux_loss: 负载均衡损失
        """
        batch_size, seq_len, d_model = x.shape

        # 门控
        gates, indices, load = self.gate(x, train)
        # gates: [batch_size, seq_len, top_k]
        # indices: [batch_size, seq_len, top_k]

        # 初始化输出
        output = torch.zeros_like(x)

        # 为每个专家处理分配的tokens
        for i in range(self.num_experts):
            # 找到分配给当前专家的所有位置
            expert_mask = (indices == i).any(dim=-1)  # [batch_size, seq_len]

            if expert_mask.any():
                # 提取分配给该专家的tokens
                expert_input = x[expert_mask]  # [num_tokens, d_model]

                # 专家处理
                expert_output = self.experts[i](expert_input)

                # 获取对应的门控权重
                expert_gates = gates[expert_mask]  # [num_tokens, top_k]
                expert_indices = indices[expert_mask]  # [num_tokens, top_k]

                # 找到当前专家在top_k中的位置
                expert_positions = (expert_indices == i).nonzero(as_tuple=True)[1]
                expert_weights = expert_gates[
                    torch.arange(expert_gates.size(0)),
                    expert_positions
                ].unsqueeze(-1)

                # 加权累加到输出
                output[expert_mask] += expert_weights * expert_output

        # 计算负载均衡损失
        aux_loss = self._calculate_balance_loss(load)

        return output, aux_loss

    def _calculate_balance_loss(self, load: torch.Tensor) -> torch.Tensor:
        """
        计算负载均衡损失
        目标：让每个专家的负载接近 1/num_experts
        """
        # 目标负载
        target_load = 1.0 / self.num_experts

        # 使用CV (Coefficient of Variation) 作为均衡指标
        balance_loss = torch.sum(load * torch.log(load * self.num_experts + 1e-8))

        return self.balance_loss_weight * balance_loss

class MoETransformerBlock(nn.Module):
    """包含MoE的Transformer块"""
    def __init__(
        self,
        d_model: int = 768,
        num_heads: int = 12,
        d_ff: int = 3072,
        num_experts: int = 8,
        top_k: int = 2,
        dropout: float = 0.1
    ):
        super().__init__()

        # 自注意力
        self.attention = nn.MultiheadAttention(
            d_model,
            num_heads,
            dropout=dropout,
            batch_first=True
        )
        self.norm1 = nn.LayerNorm(d_model)

        # MoE层
        self.moe = MoELayer(
            d_model=d_model,
            d_ff=d_ff,
            num_experts=num_experts,
            top_k=top_k,
            dropout=dropout
        )
        self.norm2 = nn.LayerNorm(d_model)

    def forward(
        self,
        x: torch.Tensor,
        mask: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # 自注意力 + 残差
        attn_out, _ = self.attention(x, x, x, attn_mask=mask)
        x = self.norm1(x + attn_out)

        # MoE + 残差
        moe_out, aux_loss = self.moe(x)
        x = self.norm2(x + moe_out)

        return x, aux_loss

# 使用示例
def example_moe_usage():
    """MoE使用示例"""
    batch_size = 4
    seq_len = 128
    d_model = 768

    # 创建MoE Transformer块
    moe_block = MoETransformerBlock(
        d_model=d_model,
        num_heads=12,
        d_ff=3072,
        num_experts=8,
        top_k=2
    )

    # 随机输入
    x = torch.randn(batch_size, seq_len, d_model)

    # 前向传播
    output, aux_loss = moe_block(x)

    print(f"输入形状: {x.shape}")
    print(f"输出形状: {output.shape}")
    print(f"负载均衡损失: {aux_loss.item():.6f}")

    # 训练循环示例
    optimizer = torch.optim.Adam(moe_block.parameters(), lr=1e-4)

    for step in range(10):
        optimizer.zero_grad()

        # 前向传播
        output, aux_loss = moe_block(x)

        # 主任务损失（这里用随机目标演示）
        target = torch.randn_like(output)
        main_loss = F.mse_loss(output, target)

        # 总损失 = 主任务损失 + 负载均衡损失
        total_loss = main_loss + aux_loss

        # 反向传播
        total_loss.backward()
        optimizer.step()

        if step % 5 == 0:
            print(f"Step {step}: Main Loss={main_loss.item():.4f}, "
                  f"Aux Loss={aux_loss.item():.6f}")

if __name__ == "__main__":
    example_moe_usage()

import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModelForCausalLM
from typing import Optional, List

class MixtralConfig:
    """Mixtral模型配置"""
    def __init__(
        self,
        vocab_size: int = 32000,
        hidden_size: int = 4096,
        intermediate_size: int = 14336,
        num_hidden_layers: int = 32,
        num_attention_heads: int = 32,
        num_key_value_heads: int = 8,  # GQA
        num_experts: int = 8,
        num_experts_per_token: int = 2,
        max_position_embeddings: int = 32768,
        rope_theta: float = 10000.0,
        rms_norm_eps: float = 1e-5,
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.num_experts = num_experts
        self.num_experts_per_token = num_experts_per_token
        self.max_position_embeddings = max_position_embeddings
        self.rope_theta = rope_theta
        self.rms_norm_eps = rms_norm_eps

class MixtralSparseMoeBlock(nn.Module):
    """Mixtral的MoE实现"""
    def __init__(self, config: MixtralConfig):
        super().__init__()
        self.hidden_dim = config.hidden_size
        self.ffn_dim = config.intermediate_size
        self.num_experts = config.num_experts
        self.top_k = config.num_experts_per_token

        # 门控网络
        self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False)

        # 专家网络（每个专家是一个完整的FFN）
        self.experts = nn.ModuleList([
            MixtralExpert(self.hidden_dim, self.ffn_dim)
            for _ in range(self.num_experts)
        ])

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, hidden_dim = hidden_states.shape

        # 重塑为2D以便处理
        hidden_states_flat = hidden_states.view(-1, hidden_dim)

        # 计算路由权重
        router_logits = self.gate(hidden_states_flat)
        routing_weights = torch.softmax(router_logits, dim=-1)

        # 选择top-k专家
        routing_weights, selected_experts = torch.topk(
            routing_weights, self.top_k, dim=-1
        )

        # 重新归一化
        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)

        # 初始化输出
        final_hidden_states = torch.zeros_like(hidden_states_flat)

        # 为每个专家处理tokens
        expert_mask = torch.nn.functional.one_hot(
            selected_experts,
            num_classes=self.num_experts
        ).permute(2, 1, 0)

        for expert_idx in range(self.num_experts):
            expert = self.experts[expert_idx]
            idx, top_x = torch.where(expert_mask[expert_idx])

            if top_x.shape[0] == 0:
                continue

            # 处理分配给该专家的tokens
            current_state = hidden_states_flat[None, top_x].reshape(-1, hidden_dim)
            current_hidden_states = expert(current_state)

            # 应用路由权重
            current_hidden_states *= routing_weights[top_x, idx, None]

            # 累加到输出
            final_hidden_states.index_add_(
                0, top_x, current_hidden_states.to(hidden_states.dtype)
            )

        # 恢复原始形状
        final_hidden_states = final_hidden_states.reshape(
            batch_size, seq_len, hidden_dim
        )

        return final_hidden_states

class MixtralExpert(nn.Module):
    """Mixtral专家网络"""
    def __init__(self, hidden_dim: int, ffn_dim: int):
        super().__init__()
        self.w1 = nn.Linear(hidden_dim, ffn_dim, bias=False)
        self.w2 = nn.Linear(ffn_dim, hidden_dim, bias=False)
        self.w3 = nn.Linear(hidden_dim, ffn_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # SwiGLU activation
        return self.w2(nn.functional.silu(self.w1(x)) * self.w3(x))

# 使用Hugging Face的Mixtral模型
def use_mixtral_model():
    """使用预训练的Mixtral模型"""
    model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"

    # 加载tokenizer和模型
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # 使用8bit量化加载（节省内存）
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        device_map="auto",
        load_in_8bit=True,
        torch_dtype=torch.float16
    )

    # 推理示例
    prompt = "Explain the concept of Mixture of Experts in AI:"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    outputs = model.generate(
        **inputs,
        max_new_tokens=200,
        temperature=0.7,
        top_p=0.9,
        do_sample=True
    )

    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(generated_text)

# 标准自注意力的复杂度
import torch
import torch.nn as nn
import math

def standard_attention(Q, K, V, mask=None):
    """
    标准缩放点积注意力

    复杂度：
    - 时间: O(n²d)，其中n是序列长度，d是特征维度
    - 空间: O(n²)，需要存储注意力矩阵
    """
    d_k = Q.size(-1)

    # QK^T: [batch, heads, n, n]
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)

    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)

    # Softmax: [batch, heads, n, n]
    attn_weights = torch.softmax(scores, dim=-1)

    # 输出: [batch, heads, n, d]
    output = torch.matmul(attn_weights, V)

    return output, attn_weights

# 复杂度分析
def analyze_attention_complexity():
    """分析不同序列长度的内存和计算需求"""
    import numpy as np

    batch_size = 1
    num_heads = 32
    d_head = 128

    seq_lengths = [512, 2048, 8192, 32768, 131072]

    print("序列长度 | 注意力矩阵大小 | 内存占用(GB) | FLOPs(B)")
    print("-" * 65)

    for seq_len in seq_lengths:
        # 注意力矩阵大小: [batch, heads, n, n]
        attn_matrix_size = batch_size * num_heads * seq_len * seq_len

        # 内存占用（float16）
        memory_gb = attn_matrix_size * 2 / (1024**3)

        # FLOPs: QK^T + Softmax + AttnV
        flops = (
            2 * seq_len * seq_len * d_head +  # QK^T
            4 * seq_len * seq_len +            # Softmax
            2 * seq_len * seq_len * d_head     # AttnV
        ) * num_heads
        flops_b = flops / 1e9

        print(f"{seq_len:8d} | {attn_matrix_size:14,d} | "
              f"{memory_gb:11.2f} | {flops_b:10.2f}")

if __name__ == "__main__":
    analyze_attention_complexity()

import torch
import torch.nn as nn
import math
from typing import Tuple

class RotaryPositionEmbedding(nn.Module):
    """
    旋转位置编码（RoPE）
    论文: RoFormer: Enhanced Transformer with Rotary Position Embedding
    """
    def __init__(
        self,
        dim: int,
        max_seq_len: int = 2048,
        base: float = 10000.0
    ):
        super().__init__()
        self.dim = dim
        self.max_seq_len = max_seq_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,
        seq_len: int,
        device: torch.device
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        生成cos和sin缓存

        Returns:
            cos_cached: [seq_len, dim]
            sin_cached: [seq_len, dim]
        """
        # 位置索引
        t = torch.arange(seq_len, device=device).type_as(self.inv_freq)

        # 计算频率 * 位置
        freqs = torch.einsum('i,j->ij', t, self.inv_freq)

        # 拼接以匹配完整维度
        emb = torch.cat((freqs, freqs), dim=-1)

        return emb.cos(), emb.sin()

def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """将特征维度分为两半并旋转"""
    x1, x2 = x[..., :x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    应用旋转位置编码

    Args:
        q, k: [batch, heads, seq_len, dim]
        cos, sin: [seq_len, dim]
    """
    # 扩展维度以匹配q, k
    cos = cos[None, None, :, :]  # [1, 1, seq_len, dim]
    sin = sin[None, None, :, :]

    # 应用旋转
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)

    return q_embed, k_embed

class RoPEScaledAttention(nn.Module):
    """带有RoPE的注意力层，支持长度外推"""
    def __init__(
        self,
        d_model: int,
        num_heads: int,
        max_seq_len: int = 2048,
        rope_base: float = 10000.0,
        rope_scaling: Optional[dict] = None
    ):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_head = d_model // num_heads

        # QKV投影
        self.q_proj = nn.Linear(d_model, d_model, bias=False)
        self.k_proj = nn.Linear(d_model, d_model, bias=False)
        self.v_proj = nn.Linear(d_model, d_model, bias=False)
        self.o_proj = nn.Linear(d_model, d_model, bias=False)

        # RoPE
        self.rope = RotaryPositionEmbedding(
            self.d_head,
            max_seq_len,
            rope_base
        )

        # RoPE缩放（用于长度外推）
        self.rope_scaling = rope_scaling
        if rope_scaling:
            self.rope_scale_factor = rope_scaling.get('factor', 1.0)
            self.rope_scale_type = rope_scaling.get('type', 'linear')
        else:
            self.rope_scale_factor = 1.0
            self.rope_scale_type = None

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        batch_size, seq_len, _ = hidden_states.shape

        # QKV投影
        Q = self.q_proj(hidden_states)
        K = self.k_proj(hidden_states)
        V = self.v_proj(hidden_states)

        # 重塑为多头
        Q = Q.view(batch_size, seq_len, self.num_heads, self.d_head)
        K = K.view(batch_size, seq_len, self.num_heads, self.d_head)
        V = V.view(batch_size, seq_len, self.num_heads, self.d_head)

        # 转置: [batch, heads, seq_len, d_head]
        Q = Q.transpose(1, 2)
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)

        # 应用RoPE
        cos, sin = self.rope(seq_len, hidden_states.device)

        # RoPE缩放（用于外推）
        if self.rope_scaling:
            cos, sin = self._apply_rope_scaling(cos, sin, seq_len)

        Q, K = apply_rotary_pos_emb(Q, K, cos, sin)

        # 计算注意力
        attn_output = self._compute_attention(Q, K, V, attention_mask)

        # 输出投影
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, self.d_model)
        output = self.o_proj(attn_output)

        return output

    def _apply_rope_scaling(
        self,
        cos: torch.Tensor,
        sin: torch.Tensor,
        seq_len: int
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """应用RoPE缩放以支持长度外推"""
        if self.rope_scale_type == 'linear':
            # 线性插值
            max_seq_len = self.rope.max_seq_len
            if seq_len > max_seq_len:
                # 插值到更长的序列
                scale = seq_len / max_seq_len
                positions = torch.arange(seq_len, device=cos.device) / scale

                # 重新计算频率
                freqs = torch.einsum(
                    'i,j->ij',
                    positions,
                    self.rope.inv_freq
                )
                emb = torch.cat((freqs, freqs), dim=-1)
                cos = emb.cos()
                sin = emb.sin()

        elif self.rope_scale_type == 'dynamic':
            # 动态NTK插值
            scale = self.rope_scale_factor
            max_seq_len = self.rope.max_seq_len * scale

            # 调整base以保持高频信息
            new_base = self.rope.base * (
                (scale * seq_len / max_seq_len) - (scale - 1)
            ) ** (self.d_head / (self.d_head - 2))

            inv_freq = 1.0 / (
                new_base ** (
                    torch.arange(0, self.d_head, 2, device=cos.device).float()
                    / self.d_head
                )
            )

            t = torch.arange(seq_len, device=cos.device).type_as(inv_freq)
            freqs = torch.einsum('i,j->ij', t, inv_freq)
            emb = torch.cat((freqs, freqs), dim=-1)
            cos = emb.cos()
            sin = emb.sin()

        return cos, sin

    def _compute_attention(
        self,
        Q: torch.Tensor,
        K: torch.Tensor,
        V: torch.Tensor,
        mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """计算缩放点积注意力"""
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_head)

        if mask is not None:
            scores = scores + mask

        attn_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, V)

        return output

# 测试RoPE外推
def test_rope_extrapolation():
    """测试RoPE的长度外推能力"""
    d_model = 768
    num_heads = 12
    train_seq_len = 2048
    test_seq_len = 8192  # 4倍于训练长度

    # 创建模型（带RoPE缩放）
    rope_scaling = {
        'type': 'dynamic',
        'factor': 4.0
    }

    model = RoPEScaledAttention(
        d_model=d_model,
        num_heads=num_heads,
        max_seq_len=train_seq_len,
        rope_scaling=rope_scaling
    )

    # 测试短序列（训练长度）
    short_input = torch.randn(1, train_seq_len, d_model)
    short_output = model(short_input)
    print(f"短序列 ({train_seq_len}): 输出形状 {short_output.shape}")

    # 测试长序列（外推）
    long_input = torch.randn(1, test_seq_len, d_model)
    long_output = model(long_input)
    print(f"长序列 ({test_seq_len}): 输出形状 {long_output.shape}")

    print("\nRoPE外推成功！")

if __name__ == "__main__":
    test_rope_extrapolation()

import torch
import torch.nn as nn
from typing import Optional

def sliding_window_attention(
    Q: torch.Tensor,
    K: torch.Tensor,
    V: torch.Tensor,
    window_size: int = 512
) -> torch.Tensor:
    """
    滑动窗口注意力

    每个token只关注窗口内的其他tokens
    复杂度: O(n * w * d)，其中w是窗口大小

    Args:
        Q, K, V: [batch, heads, seq_len, d_head]
        window_size: 窗口大小
    """
    batch_size, num_heads, seq_len, d_head = Q.shape

    # 创建滑动窗口mask
    mask = torch.ones(seq_len, seq_len, device=Q.device)
    for i in range(seq_len):
        start = max(0, i - window_size // 2)
        end = min(seq_len, i + window_size // 2 + 1)
        mask[i, start:end] = 0

    # 注意力计算
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_head ** 0.5)
    scores = scores.masked_fill(mask.bool(), float('-inf'))
    attn_weights = torch.softmax(scores, dim=-1)
    output = torch.matmul(attn_weights, V)

    return output

class FlashAttention(nn.Module):
    """
    FlashAttention-2 简化实现

    核心思想：
    1. 分块计算，减少HBM访问
    2. 在线softmax，避免存储完整注意力矩阵
    3. 重计算而非存储中间结果
    """
    def __init__(self, dropout: float = 0.0):
        super().__init__()
        self.dropout = dropout

    def forward(
        self,
        Q: torch.Tensor,
        K: torch.Tensor,
        V: torch.Tensor,
        causal: bool = False,
        block_size: int = 256
    ) -> torch.Tensor:
        """
        FlashAttention前向传播

        Args:
            Q, K, V: [batch, heads, seq_len, d_head]
            causal: 是否使用因果mask
            block_size: 分块大小
        """
        batch_size, num_heads, seq_len, d_head = Q.shape

        # 分块数量
        num_blocks = (seq_len + block_size - 1) // block_size

        # 输出缓冲区
        O = torch.zeros_like(Q)
        l = torch.zeros(batch_size, num_heads, seq_len, 1, device=Q.device)
        m = torch.full(
            (batch_size, num_heads, seq_len, 1),
            float('-inf'),
            device=Q.device
        )

        # 分块处理
        for i in range(num_blocks):
            # Q块
            q_start = i * block_size
            q_end = min((i + 1) * block_size, seq_len)
            Q_block = Q[:, :, q_start:q_end, :]

            # 初始化块级统计
            O_block = torch.zeros_like(Q_block)
            l_block = torch.zeros(
                batch_size, num_heads, q_end - q_start, 1,
                device=Q.device
            )
            m_block = torch.full(
                (batch_size, num_heads, q_end - q_start, 1),
                float('-inf'),
                device=Q.device
            )

            for j in range(num_blocks):
                # 因果mask检查
                if causal and j > i:
                    continue

                # K, V块
                k_start = j * block_size
                k_end = min((j + 1) * block_size, seq_len)
                K_block = K[:, :, k_start:k_end, :]
                V_block = V[:, :, k_start:k_end, :]

                # 计算注意力分数
                S_block = torch.matmul(Q_block, K_block.transpose(-2, -1))
                S_block = S_block / (d_head ** 0.5)

                # 因果mask
                if causal:
                    block_mask = torch.triu(
                        torch.ones(
                            q_end - q_start, k_end - k_start,
                            device=Q.device
                        ),
                        diagonal=k_start - q_start + 1
                    )
                    S_block = S_block.masked_fill(
                        block_mask.bool(),
                        float('-inf')
                    )

                # 在线softmax更新
                m_block_new = torch.maximum(
                    m_block,
                    S_block.max(dim=-1, keepdim=True)[0]
                )

                # 更新指数和
                alpha = torch.exp(m_block - m_block_new)
                beta = torch.exp(S_block - m_block_new)

                l_block_new = alpha * l_block + beta.sum(dim=-1, keepdim=True)

                # 更新输出
                O_block = (
                    alpha * O_block +
                    torch.matmul(beta, V_block)
                )

                m_block = m_block_new
                l_block = l_block_new

            # 归一化
            O_block = O_block / l_block

            # 写回全局输出
            O[:, :, q_start:q_end, :] = O_block

        return O

# 性能对比
def compare_attention_methods():
    """对比不同注意力方法的性能"""
    import time

    batch_size = 2
    num_heads = 12
    seq_len = 4096
    d_head = 64

    Q = torch.randn(batch_size, num_heads, seq_len, d_head, device='cuda')
    K = torch.randn(batch_size, num_heads, seq_len, d_head, device='cuda')
    V = torch.randn(batch_size, num_heads, seq_len, d_head, device='cuda')

    # 标准注意力
    torch.cuda.synchronize()
    start = time.time()
    standard_out, _ = standard_attention(Q, K, V)
    torch.cuda.synchronize()
    standard_time = time.time() - start

    # 滑动窗口注意力
    torch.cuda.synchronize()
    start = time.time()
    window_out = sliding_window_attention(Q, K, V, window_size=512)
    torch.cuda.synchronize()
    window_time = time.time() - start

    # FlashAttention
    flash_attn = FlashAttention().cuda()
    torch.cuda.synchronize()
    start = time.time()
    flash_out = flash_attn(Q, K, V)
    torch.cuda.synchronize()
    flash_time = time.time() - start

    print(f"序列长度: {seq_len}")
    print(f"标准注意力: {standard_time:.4f}s")
    print(f"窗口注意力: {window_time:.4f}s ({standard_time/window_time:.2f}x)")
    print(f"FlashAttention: {flash_time:.4f}s ({standard_time/flash_time:.2f}x)")

if __name__ == "__main__":
    if torch.cuda.is_available():
        compare_attention_methods()

import torch
import torch.nn as nn
import numpy as np
from typing import Tuple

def quantize_tensor(tensor: torch.Tensor, num_bits: int = 8) -> Tuple[torch.Tensor, float, float]:
    """
    对称量化

    量化公式:
    q = round(x / scale)
    x_reconstructed = q * scale

    Args:
        tensor: 原始浮点张量
        num_bits: 量化位数

    Returns:
        quantized: 量化后的整数张量
        scale: 量化比例
        zero_point: 零点
    """
    # 计算范围
    qmin = -(2 ** (num_bits - 1))
    qmax = 2 ** (num_bits - 1) - 1

    # 计算scale
    min_val = tensor.min().item()
    max_val = tensor.max().item()

    scale = (max_val - min_val) / (qmax - qmin)
    zero_point = qmin - min_val / scale

    # 量化
    quantized = torch.clamp(
        torch.round(tensor / scale + zero_point),
        qmin, qmax
    ).to(torch.int8)

    return quantized, scale, zero_point

def dequantize_tensor(
    quantized: torch.Tensor,
    scale: float,
    zero_point: float
) -> torch.Tensor:
    """反量化"""
    return (quantized.float() - zero_point) * scale

# INT8量化示例
def int8_quantization_example():
    """INT8量化示例"""
    # 原始权重
    weight = torch.randn(512, 512) * 0.1

    # 量化
    weight_q, scale, zero_point = quantize_tensor(weight, num_bits=8)

    # 反量化
    weight_dq = dequantize_tensor(weight_q, scale, zero_point)

    # 计算误差
    error = (weight - weight_dq).abs().mean()

    # 内存节省
    original_size = weight.element_size() * weight.nelement() / 1024  # KB
    quantized_size = weight_q.element_size() * weight_q.nelement() / 1024

    print(f"原始权重: {original_size:.2f} KB (FP32)")
    print(f"量化权重: {quantized_size:.2f} KB (INT8)")
    print(f"压缩比: {original_size / quantized_size:.2f}x")
    print(f"平均误差: {error:.6f}")

if __name__ == "__main__":
    int8_quantization_example()

import torch
import torch.nn as nn
from typing import List, Tuple
import math

class GPTQ:
    """
    GPTQ量化器

    论文: GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers

    核心思想:
    1. 逐层量化，最小化重构误差
    2. 使用Hessian信息指导量化
    3. 贪心逐列优化
    """
    def __init__(
        self,
        layer: nn.Linear,
        bits: int = 4,
        group_size: int = 128,
        actorder: bool = True
    ):
        self.layer = layer
        self.bits = bits
        self.group_size = group_size
        self.actorder = actorder

        self.dev = layer.weight.device
        self.rows = layer.weight.shape[0]
        self.columns = layer.weight.shape[1]

        # Hessian矩阵（近似）
        self.H = torch.zeros(
            (self.columns, self.columns),
            device=self.dev
        )
        self.nsamples = 0

    def add_batch(self, inp: torch.Tensor):
        """
        添加一批样本来估计Hessian矩阵

        Args:
            inp: [batch_size, seq_len, in_features]
        """
        if len(inp.shape) == 3:
            inp = inp.reshape(-1, inp.shape[-1])

        tmp = inp.shape[0]

        # 计算H = 2 * X^T X
        inp = math.sqrt(2 / self.nsamples) * inp.float()
        self.H += inp.t().matmul(inp)
        self.nsamples += tmp

    def quantize(self) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        执行GPTQ量化

        Returns:
            quantized_weight: 量化后的权重
            scales: 量化scale
        """
        W = self.layer.weight.data.clone()

        # 确保Hessian已经计算
        if self.nsamples == 0:
            raise ValueError("请先调用add_batch添加样本")

        # 计算量化参数
        qmin = -(2 ** (self.bits - 1))
        qmax = 2 ** (self.bits - 1) - 1

        # 初始化
        H = self.H
        dead = torch.diag(H) == 0
        H[dead, dead] = 1
        W[:, dead] = 0

        # Cholesky分解
        try:
            damp = 0.01 * torch.mean(torch.diag(H))
            diag = torch.arange(self.columns, device=self.dev)
            H[diag, diag] += damp
            H = torch.linalg.cholesky(H)
            H = torch.cholesky_inverse(H)
            H = torch.linalg.cholesky(H, upper=True)
            Hinv = H
        except:
            print("Cholesky分解失败，使用对角近似")
            Hinv = torch.diag(1.0 / torch.diag(H))

        # 量化顺序
        if self.actorder:
            perm = torch.argsort(torch.diag(H), descending=True)
            W = W[:, perm]
            Hinv = Hinv[perm][:, perm]

        # 分组量化
        Q = torch.zeros_like(W)
        scales = torch.zeros(
            (self.rows, (self.columns + self.group_size - 1) // self.group_size),
            device=self.dev
        )

        Err = torch.zeros_like(W)

        for i1 in range(0, self.columns, self.group_size):
            i2 = min(i1 + self.group_size, self.columns)
            count = i2 - i1

            W1 = W[:, i1:i2].clone()
            Q1 = torch.zeros_like(W1)
            Err1 = torch.zeros_like(W1)
            H1 = Hinv[i1:i2, i1:i2]

            for i in range(count):
                w = W1[:, i]
                d = H1[i, i]

                # 计算scale
                scale = w.abs().max() / qmax
                scales[:, i1 // self.group_size] = scale

                # 量化
                q = torch.clamp(
                    torch.round(w / scale),
                    qmin, qmax
                )
                Q1[:, i] = q * scale

                # 更新误差
                err = (w - Q1[:, i]) / d
                Err1[:, i] = err

                # 传播误差到后续列
                W1[:, i:] -= err.unsqueeze(1).matmul(H1[i, i:].unsqueeze(0))

            Q[:, i1:i2] = Q1
            Err[:, i1:i2] = Err1

        # 恢复原始顺序
        if self.actorder:
            invperm = torch.argsort(perm)
            Q = Q[:, invperm]

        # 转换为整数
        Q_int = torch.zeros_like(Q, dtype=torch.int8)
        for i in range(scales.shape[1]):
            i1 = i * self.group_size
            i2 = min(i1 + self.group_size, self.columns)
            scale = scales[:, i:i+1]
            Q_int[:, i1:i2] = torch.clamp(
                torch.round(Q[:, i1:i2] / scale),
                qmin, qmax
            ).to(torch.int8)

        return Q_int, scales

# 使用GPTQ量化模型
def quantize_model_with_gptq(
    model: nn.Module,
    calibration_data: torch.Tensor,
    bits: int = 4
) -> nn.Module:
    """
    使用GPTQ量化整个模型

    Args:
        model: 要量化的模型
        calibration_data: 校准数据
        bits: 量化位数
    """
    # 收集所有Linear层
    linear_layers = []
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            linear_layers.append((name, module))

    print(f"找到 {len(linear_layers)} 个线性层")

    # 逐层量化
    for name, layer in linear_layers:
        print(f"量化层: {name}")

        # 创建GPTQ量化器
        gptq = GPTQ(layer, bits=bits)

        # 收集激活值（这里简化，实际需要hook）
        # 实际应用中需要在前向传播时收集
        with torch.no_grad():
            fake_input = torch.randn(
                calibration_data.shape[0],
                layer.in_features,
                device=layer.weight.device
            )
            gptq.add_batch(fake_input)

        # 量化
        q_weight, scales = gptq.quantize()

        # 替换权重（这里简化，实际需要自定义量化层）
        print(f"  原始: {layer.weight.shape}, 量化: {q_weight.shape}")

    return model

# 示例
def gptq_example():
    """GPTQ量化示例"""
    # 创建简单模型
    model = nn.Sequential(
        nn.Linear(512, 1024),
        nn.ReLU(),
        nn.Linear(1024, 512)
    )

    # 校准数据
    calib_data = torch.randn(100, 512)

    # 量化
    quantized_model = quantize_model_with_gptq(model, calib_data, bits=4)

    print("\nGPTQ量化完成!")

if __name__ == "__main__":
    gptq_example()

import torch
import torch.nn as nn
from typing import Dict, List

class AWQ:
    """
    AWQ量化器

    论文: AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration

    核心思想:
    1. 识别对激活值影响大的权重通道
    2. 对这些通道应用较小的量化力度
    3. 使用逐通道scale
    """
    def __init__(
        self,
        layer: nn.Linear,
        bits: int = 4,
        group_size: int = 128,
        alpha: float = 0.5
    ):
        self.layer = layer
        self.bits = bits
        self.group_size = group_size
        self.alpha = alpha

        self.dev = layer.weight.device

        # 激活值统计
        self.act_scales = None
        self.nsamples = 0

    def collect_activations(self, inp: torch.Tensor):
        """
        收集激活值统计

        Args:
            inp: [batch_size, seq_len, in_features]
        """
        if len(inp.shape) == 3:
            inp = inp.reshape(-1, inp.shape[-1])

        # 计算每个通道的激活幅度
        scales = inp.abs().mean(dim=0)

        if self.act_scales is None:
            self.act_scales = scales
        else:
            # 指数移动平均
            self.act_scales = 0.9 * self.act_scales + 0.1 * scales

        self.nsamples += inp.shape[0]

    def compute_scaling_factors(self) -> torch.Tensor:
        """
        计算每个通道的缩放因子

        Returns:
            s: [in_features] 缩放因子
        """
        if self.act_scales is None:
            raise ValueError("请先收集激活值")

        # 基于激活值重要性的缩放
        s = self.act_scales.pow(self.alpha)

        # 归一化（保持整体scale）
        s = s / s.mean()

        return s

    def quantize(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        执行AWQ量化

        Returns:
            quantized_weight: 量化后的权重
            scales: 量化scale
            scaling_factors: 通道缩放因子
        """
        W = self.layer.weight.data.clone()

        # 计算缩放因子
        s = self.compute_scaling_factors()

        # 应用缩放
        W_scaled = W * s.view(1, -1)

        # 分组量化
        qmin = -(2 ** (self.bits - 1))
        qmax = 2 ** (self.bits - 1) - 1

        out_features, in_features = W_scaled.shape
        num_groups = (in_features + self.group_size - 1) // self.group_size

        Q = torch.zeros_like(W_scaled, dtype=torch.int8)
        scales = torch.zeros(out_features, num_groups, device=self.dev)

        for g in range(num_groups):
            start = g * self.group_size
            end = min(start + self.group_size, in_features)

            W_group = W_scaled[:, start:end]

            # 计算scale（逐通道）
            scale = W_group.abs().max(dim=1)[0] / qmax
            scales[:, g] = scale

            # 量化
            Q[:, start:end] = torch.clamp(
                torch.round(W_group / scale.view(-1, 1)),
                qmin, qmax
            ).to(torch.int8)

        return Q, scales, s

class AWQLinear(nn.Module):
    """AWQ量化的线性层"""
    def __init__(
        self,
        in_features: int,
        out_features: int,
        bits: int = 4,
        group_size: int = 128
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.bits = bits
        self.group_size = group_size

        # 量化权重
        self.register_buffer(
            'qweight',
            torch.zeros(
                (out_features, in_features),
                dtype=torch.int8
            )
        )

        # Scales
        num_groups = (in_features + group_size - 1) // group_size
        self.register_buffer(
            'scales',
            torch.zeros((out_features, num_groups))
        )

        # 通道缩放因子
        self.register_buffer(
            'scaling_factors',
            torch.ones(in_features)
        )

        if hasattr(nn.Linear(1, 1), 'bias'):
            self.register_buffer('bias', None)

    def pack_weights(
        self,
        qweight: torch.Tensor,
        scales: torch.Tensor,
        scaling_factors: torch.Tensor
    ):
        """打包量化权重"""
        self.qweight = qweight
        self.scales = scales
        self.scaling_factors = scaling_factors

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        前向传播

        Args:
            x: [batch_size, seq_len, in_features]
        """
        # 应用输入缩放
        x_scaled = x / self.scaling_factors.view(1, 1, -1)

        # 反量化权重并计算
        out_features, in_features = self.qweight.shape
        num_groups = self.scales.shape[1]

        output = torch.zeros(
            x.shape[0], x.shape[1], out_features,
            device=x.device, dtype=x.dtype
        )

        for g in range(num_groups):
            start = g * self.group_size
            end = min(start + self.group_size, in_features)

            # 反量化
            W_dequant = (
                self.qweight[:, start:end].float() *
                self.scales[:, g:g+1]
            )

            # 计算
            output += torch.matmul(
                x_scaled[:, :, start:end],
                W_dequant.t()
            )

        if self.bias is not None:
            output += self.bias

        return output

# 使用示例
def awq_example():
    """AWQ量化示例"""
    # 原始层
    layer = nn.Linear(512, 1024).cuda()

    # 创建AWQ量化器
    awq = AWQ(layer, bits=4, group_size=128)

    # 收集激活值
    for _ in range(10):
        fake_input = torch.randn(32, 128, 512).cuda()
        awq.collect_activations(fake_input)

    # 量化
    qweight, scales, scaling_factors = awq.quantize()

    # 创建量化层
    q_layer = AWQLinear(512, 1024, bits=4, group_size=128).cuda()
    q_layer.pack_weights(qweight, scales, scaling_factors)

    # 测试
    test_input = torch.randn(1, 10, 512).cuda()

    with torch.no_grad():
        original_output = layer(test_input)
        quantized_output = q_layer(test_input)

    error = (original_output - quantized_output).abs().mean()
    print(f"量化误差: {error:.6f}")

    # 内存节省
    original_size = layer.weight.nelement() * layer.weight.element_size() / 1024 / 1024
    quantized_size = (
        qweight.nelement() * qweight.element_size() +
        scales.nelement() * scales.element_size()
    ) / 1024 / 1024

    print(f"原始大小: {original_size:.2f} MB")
    print(f"量化大小: {quantized_size:.2f} MB")
    print(f"压缩比: {original_size / quantized_size:.2f}x")

if __name__ == "__main__":
    if torch.cuda.is_available():
        awq_example()