张量并行Megatron-LM

张量并行详解:Megatron-LM

为什么需要张量并行? 当模型参数量超过单GPU显存容量时,模型并行成为必需。张量并行(Tensor Parallelism, TP)是最常用的模型并行方式——将单个层的权重矩阵切分到多个GPU上,每个GPU计算一部分结果,然后通过通信合并。 NVIDIA的Megatron-LM提出了张量并行的标准方案,成为2026年大模型训练的基础设施。 张量并行的核心原理 线性层的切分 一个标准线性层 Y = XW,其中 W ∈ ℝ^{d×h}。将W切分为两个子矩阵有两种方式: 列切分(Column Parallel) 将W按列切分:W = [W₁, W₂],每个GPU持有 W_i ∈ ℝ^{d×h/N} GPU 1: Y₁ = XW₁ ∈ ℝ^{n×h/N} GPU 2: Y₂ = XW₂ ∈ ℝ^{n×h/N} Y = [Y₁, Y₂] ∈ ℝ^{n×h} 列切分后,每个GPU独立计算Y的一部分列,最终通过All-Gather合并。但更聪明的方式是——不需要显式合并,因为后续操作可以利用分块结构。 行切分(Row Parallel) 将W按行切分:W = [W₁; W₂],每个GPU持有 W_i ∈ ℝ^{d/N×h} GPU 1: Y₁ = X₁W₁ (X₁是X的前d/N列) GPU 2: Y₂ = X₂W₂ (X₂是X的后d/N列) Y = Y₁ + Y₂ (All-Reduce) 行切分后需要All-Reduce求和。 ...

2026-07-02 · 3 min · 461 words · 硅基 AGI 探索者
梯度检查点

梯度检查点原理与实现

训练大模型的显存困境 训练一个70B参数的模型,仅模型参数(FP16)就需要140GB显存。加上梯度(140GB)、优化器状态(Adam需要280GB)和激活值,总显存需求轻松超过600GB——即使8卡A100 80GB也无法满足。 在这些显存消耗中,激活值是最容易被忽视的部分。在标准反向传播中,前向传播的所有中间激活值都需要保存下来供反向传播计算梯度使用。对于70B模型、batch_size=4、seq_len=4096,激活值占用可达100GB以上。 梯度检查点(Gradient Checkpointing)通过"以时间换空间"的策略,大幅减少激活值的显存占用。 标准反向传播的显存分析 在标准反向传播中,每个操作的前向输出都需要被保存: 前向: x → Layer1 → a1 → Layer2 → a2 → Layer3 → a3 → Loss 保存: a1, a2, a3 (以及各层的中间状态) 反向: Loss ← ∂L/∂a3 ← ∂L/∂a2 ← ∂L/∂a1 ← ∂L/∂x 使用: 重新使用a3计算a2的梯度,重新使用a2计算a1的梯度,... 对于N层Transformer,每层的激活值大小为 O(batch × seq_len × hidden_dim)。所有层的激活值总和为 O(N × batch × seq_len × hidden_dim),与层数线性增长。 梯度检查点的核心思想 选择性重计算 梯度检查点的核心洞察:不需要保存所有层的激活值,只需保存部分"检查点"层的输入,其余层的激活值在反向传播时重新计算。 检查点设置(每2层一个检查点): 保存: x, a2, a4, a6, ... 前向: x → L1 → a1 → L2 → a2* → L3 → a3 → L4 → a4* → ... (* = 保存的检查点) 反向: 从a6*出发,重新前向计算a5, a4→计算梯度 从a4*出发,重新前向计算a3, a2→计算梯度 ... 显存-计算权衡 假设N层Transformer,每隔k层设置一个检查点: ...

2026-07-02 · 3 min · 514 words · 硅基 AGI 探索者
激活函数综述

激活函数综述2026

激活函数的角色 激活函数是神经网络中引入非线性的关键组件。没有激活函数,多层线性变换的堆叠等价于单层线性变换——网络的表达能力将被严重限制。激活函数的选择直接影响模型的训练动态、收敛速度和最终性能。 在LLM时代,激活函数的演进从ReLU到GELU再到SwiGLU,每一步都带来了可测量的性能提升。 ReLU时代 ReLU的革命性 ReLU(Rectified Linear Unit)的定义极其简单: ReLU(x) = max(0, x) 在ReLU之前,sigmoid和tanh是主流选择,但它们存在梯度消失问题——深层网络中梯度指数衰减。ReLU的梯度在正区间恒为1,有效解决了梯度消失问题,使得深层网络的训练成为可能。 ReLU的缺陷 Dead ReLU问题:当输入持续为负时,ReLU的梯度为零,神经元将永久"死亡"无法恢复。这在学习率设置不当时尤为严重。 非零中心化:ReLU的输出始终非负,导致后续层的输入分布偏向正方向,影响梯度下降效率。 Leaky ReLU与变体 为解决Dead ReLU问题,多种变体被提出: def leaky_relu(x, negative_slope=0.01): """Leaky ReLU: 负区间保留小梯度""" return torch.where(x > 0, x, negative_slope * x) def prelu(x, alpha): """Parametric ReLU: 负区间斜率可学习""" return torch.where(x > 0, x, alpha * x) def elu(x, alpha=1.0): """ELU: 负区间平滑过渡到指数""" return torch.where(x > 0, x, alpha * (torch.exp(x) - 1)) 这些变体在CV领域有一定应用,但在LLM中几乎未被采用——LLM的激活函数走上了另一条路。 GELU:Transformer的原始选择 定义 GELU(Gaussian Error Linear Unit)将输入的高斯分布概率与输入本身相乘: GELU(x) = x · Φ(x) 其中 Φ(x) 是标准正态分布的累积分布函数。实践中常使用近似: def gelu(x): """精确GELU""" return 0.5 * x * (1 + torch.erf(x / math.sqrt(2))) def gelu_tanh_approx(x): """tanh近似(更快)""" return 0.5 * x * (1 + torch.tanh(math.sqrt(2/math.pi) * (x + 0.044715 * x**3))) GELU vs ReLU GELU相比ReLU有两个关键优势: 平滑过渡:在零点附近,GELU是平滑的而非硬截断。这使得梯度更连续,训练更稳定 随机正则化:GELU隐含了一种随机dropout机制——输入越大,被保留的概率越高。这在一定程度上起到了自正则化的作用 原始Transformer(Attention is All You Need)选择了GELU,此后BERT、GPT系列也沿用至今。 ...

2026-07-02 · 2 min · 369 words · 硅基 AGI 探索者
神经网络归一化

神经网络归一化:LN vs BN vs RMSNorm

归一化:深度学习的"稳定器" 深度神经网络的训练面临一个根本性挑战:随着层数加深,激活值的分布会发生剧烈变化(Internal Covariate Shift)。这种分布偏移导致上层需要不断适应下层的输出分布变化,学习变得困难。归一化层通过将激活值"拉回"到稳定分布,有效缓解了这一问题。 在Transformer时代,Layer Normalization(LN)已成为标配,而RMSNorm等轻量化变体正在成为新的主流选择。 Batch Normalization的局限性 BN的工作原理 Batch Normalization对batch维度进行归一化: def batch_norm(x, gamma, beta, eps=1e-5, momentum=0.1, running_mean=None, running_var=None): """ x: [batch_size, channels, height, width] """ if running_mean is not None: # 推理模式:使用预计算的统计量 mean = running_mean var = running_var else: # 训练模式:计算当前batch的统计量 mean = x.mean(dim=(0, 2, 3), keepdim=True) var = x.var(dim=(0, 2, 3), keepdim=True) # 更新running统计量 running_mean = momentum * mean + (1 - momentum) * running_mean running_var = momentum * var + (1 - momentum) * running_var # 归一化 x_norm = (x - mean) / torch.sqrt(var + eps) # 仿射变换 return gamma.view(1, -1, 1, 1) * x_norm + beta.view(1, -1, 1, 1) BN在LLM中的问题 BN在LLM训练中有几个致命缺陷: Batch Size依赖:LLM通常用很小的batch size(甚至batch_size=1),统计量不稳定 RNN不兼容:RNN的时间步之间需要sequential处理,无法跨batch统计 分布式训练复杂:不同GPU上的batch统计量不一致,需要同步 序列长度变化:NLP任务中序列长度经常变化,padding影响统计 Layer Normalization LN的工作原理 Layer Normalization对单个样本的所有特征进行归一化,独立于batch维度: def layer_norm(x, gamma, beta, eps=1e-5): """ x: [batch_size, seq_len, hidden_size] """ # 对最后一个维度计算均值和方差 mean = x.mean(dim=-1, keepdim=True) var = x.var(dim=-1, keepdim=True) # 归一化 x_norm = (x - mean) / torch.sqrt(var + eps) # 仿射变换 return gamma * x_norm + beta LN的优势 Batch无关:每个样本独立归一化,不依赖batch size 序列处理友好:NLP和RNN任务天然适用 实现简单:无需维护running统计量 Transformer标配:Pre-LN Transformer几乎成为标准 Pre-LN vs Post-LN 原始Transformer使用Post-LN(归一化在残差连接之后),但训练不稳定。Pre-LN将归一化移到残差分支内部: ...

2026-07-02 · 2 min · 412 words · 硅基 AGI 探索者
LLM剪枝技术

LLM剪枝2026:结构化vs非结构化

剪枝:去除冗余参数 神经网络的过参数化是公认的事实——大量参数对输出的贡献微乎其微。剪枝(Pruning)通过移除这些"冗余"参数来减少模型大小和计算量。 在LLM时代,剪枝面临新的挑战:模型规模巨大使得重新训练成本高昂,且LLM的参数分布与传统CNN有显著不同。2026年的剪枝研究集中在如何在不重训或少量微调的前提下实现高稀疏率。 非结构化剪枝 基本原理 非结构化剪枝将单个权重置零,不改变模型结构。最经典的方法是幅度剪枝(Magnitude Pruning)——移除绝对值最小的权重: def magnitude_prune(weight, sparsity=0.5): """幅度剪枝:将最小幅度的权重置零""" # 计算阈值 threshold = torch.quantile(weight.abs().flatten(), sparsity) # 创建掩码 mask = (weight.abs() > threshold).float() # 应用掩码 return weight * mask, mask 稀疏模式 非结构化剪枝产生的稀疏模式是"随机"的——零权重分布在矩阵的任意位置。这种模式虽然理论上能减少参数量,但在标准GPU上无法获得实际加速——因为稀疏矩阵的乘法效率远低于密集矩阵。 2:4稀疏 NVIDIA Ampere及以后架构支持2:4结构化稀疏——每4个连续元素中恰好2个为零。这种模式有硬件原生支持,可以获得约2倍加速: def apply_2_4_sparsity(weight): """应用2:4稀疏模式""" output = weight.clone() # 将权重重排为 [N, 4] 的块 blocks = output.view(-1, 4) for i in range(blocks.shape[0]): block = blocks[i].abs() # 保留最大的2个,其余置零 top2_indices = block.topk(2).indices mask = torch.zeros(4, device=weight.device) mask[top2_indices] = 1 blocks[i] *= mask return output.view_as(weight) 2:4稀疏是非结构化和结构化之间的折中——有一定的结构约束(硬件加速),但保持了一定的灵活性。 SparseGPT SparseGPT是2023年提出的LLM后训练剪枝方法,能在不需要重训的情况下将模型稀疏化到50%: def sparsegpt_prune(layer, calibration_data, sparsity=0.5): """SparseGPT单层剪枝""" W = layer.weight.data # [out, in] H = compute_hessian(layer, calibration_data) # [in, in] H += torch.eye(H.shape[0]) * 0.01 # 正则化 # 逐列处理 for i in range(W.shape[1]): # 计算每个权重的"重要性"分数 importance = W[:, i].abs() ** 2 / H[i, i] # 选择保留的权重 n_keep = int(W.shape[0] * (1 - sparsity)) keep_indices = importance.topk(n_keep).indices # 对非保留权重进行补偿 mask = torch.zeros(W.shape[0], device=W.device) mask[keep_indices] = 1 # 重建误差补偿 err = (W[:, i] * (1 - mask)) / H[i, i] W[:, i+1:] -= err.unsqueeze(1) @ H[i, i+1:].unsqueeze(0) # 应用掩码 W[:, i] *= mask layer.weight.data = W SparseGPT的核心创新是在剪枝的同时通过重建补偿来减少精度损失——类似于GPTQ的思想,但用于剪枝而非量化。 ...

2026-07-02 · 3 min · 437 words · 硅基 AGI 探索者
LLM蒸馏技术

LLM蒸馏技术2026实践

为什么LLM需要蒸馏? 训练一个超大模型(如700B)然后部署它,成本极其高昂。知识蒸馏(Knowledge Distillation)提供了一条务实的路径:先用大模型(Teacher)的输出作为信号训练小模型(Student),让小模型在更小参数量下接近大模型的性能。 2026年,蒸馏已经成为大模型工程化的标准环节。DeepSeek-V3、Qwen-3等模型都大量使用了蒸馏技术,将超大模型的能力迁移到可部署的尺寸。 蒸馏的理论基础 软标签的信息优势 硬标签(one-hot)只包含"正确答案"的信息,而软标签(softmax概率分布)还包含"错误答案之间的关系"。例如,在分类"猫"时,软标签可能同时给出"狗"的概率0.1——这告诉Student模型"猫和狗在某种特征上是相似的"。 这种"暗知识"(Dark Knowledge)是蒸馏有效性的核心。Teacher模型的输出分布包含了其学到的类别间关系,这些信息在硬标签中完全丢失。 温度参数 温度T控制软标签的"软度": soft_label = softmax(logits / T) 高温使分布更平滑(暴露更多暗知识),低温使分布更尖锐(接近one-hot)。实践中T通常设置为2-10。 LLM蒸馏的主要方法 1. Logit级蒸馏 最经典的蒸馏方式——Student直接学习Teacher的输出概率分布: def logit_distillation_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.7): """ student_logits, teacher_logits: [batch, seq_len, vocab_size] labels: [batch, seq_len] T: 温度参数 alpha: 蒸馏损失权重 """ # 蒸馏损失:KL散度 student_log_probs = F.log_softmax(student_logits / T, dim=-1) teacher_probs = F.softmax(teacher_logits / T, dim=-1) distill_loss = F.kl_div( student_log_probs.reshape(-1, student_logits.size(-1)), teacher_probs.reshape(-1, teacher_logits.size(-1)), reduction='batchmean' ) * (T ** 2) # 梯度缩放补偿 # 任务损失:交叉熵 task_loss = F.cross_entropy( student_logits.reshape(-1, student_logits.size(-1)), labels.reshape(-1) ) return alpha * distill_loss + (1 - alpha) * task_loss 关键点: KL散度损失需要乘以 T² 来补偿温度对梯度的影响 alpha控制蒸馏与任务学习的平衡 需要Teacher和Student的词表对齐 2. 序列级蒸馏(Sequence-Level KD) 不让Student逐token模仿Teacher,而是让Student学习Teacher生成的完整序列。具体做法是先用Teacher生成大量数据,然后用这些数据训练Student: ...

2026-07-02 · 3 min · 438 words · 硅基 AGI 探索者
量化技术对比

量化技术对比:INT4/INT8/FP8

量化:用更少的比特做更多的事 LLM的参数量从7B到671B不断增长,但GPU显存的增长速度远跟不上。量化是缓解这一矛盾最直接的技术——用更少的比特表示模型参数和激活值,在不修改模型架构的前提下减少显存占用和加速推理。 2026年,量化技术已经形成了从INT4到FP8的完整谱系,每种精度都有其适用场景和工程权衡。 量化的基本原理 对称量化 将浮点数映射到整数的固定范围。以INT8为例: x_int8 = round(x_fp / scale) scale = max(|x_fp|) / 127 反量化:x_fp ≈ x_int8 × scale 非对称量化 引入零点偏移,处理分布不对称的情况: x_int8 = round((x_fp - zero_point) / scale) scale = (max(x) - min(x)) / 255 zero_point = min(x) 量化粒度 逐张量(per-tensor):整个张量共用一个scale,最简单但精度损失大 逐通道(per-channel):每个输出通道一个scale,精度好但额外参数多 逐组(per-group):将通道分组,组内共用scale,是精度和效率的折中 def per_group_quantize(x, group_size=128): """分组量化""" original_shape = x.shape # 重排为 [out_features // group_size, group_size, in_features] x_reshaped = x.view(-1, group_size, original_shape[-1]) # 每组计算独立的scale scales = x_reshaped.abs().max(dim=1, keepdim=True).values / 127 # 量化 x_int8 = (x_reshaped / scales).round().clamp(-128, 127).to(torch.int8) return x_int8, scales def per_group_dequantize(x_int8, scales, original_shape): """分组反量化""" x_fp = x_int8.float() * scales return x_reshaped.view(original_shape) INT8量化 适用场景 INT8是最成熟的量化方案,几乎所有推理框架都原生支持。适用于: ...

2026-07-02 · 3 min · 477 words · 硅基 AGI 探索者
线性注意力研究

线性注意力研究2026

O(n²)到O(n)的吸引力 标准注意力的O(n²)复杂度是长序列处理的核心障碍。线性注意力旨在将复杂度降至O(n),使百万级token序列成为可能。但这个目标在保持注意力核心功能的同时实现极其困难——注意力之所以是O(n²),正是因为它需要建模所有Query-Key对之间的关系。 线性注意力的核函数框架 基本公式 标准注意力:Attention(Q,K,V) = softmax(QK^T)V 如果我们将softmax替换为一个核函数 φ(·),使得: Attention(Q,K,V) = φ(Q)φ(K)^T V = φ(Q)(φ(K)^T V) 关键变换在于结合律:先计算 φ(K)^T V(与Q无关),得到一个 d×d 的矩阵,然后与 φ(Q) 相乘。这使得复杂度从 O(n²d) 降为 O(nd²)。当 n » d 时,这是显著的速度提升。 Performer的随机特征 Performer使用随机特征方法来近似softmax核: def performer_attention(Q, K, V, n_features=256): """Performer的随机特征近似""" d = Q.shape[-1] # 生成随机投影矩阵 omega = torch.randn(d, n_features) / math.sqrt(d) # 非负映射函数 def phi(x): return torch.exp(x @ omega - 0.5 * (x ** 2).sum(-1, keepdim=True)) / math.sqrt(n_features) # 线性注意力 phi_Q = phi(Q) # [batch, n, m] phi_K = phi(K) # [batch, m, m] phi_V = V # [batch, m, d] # 先算 K^T V,再算 Q (K^T V) KV = torch.einsum('bmd,bme->bde', phi_K, phi_V) # [batch, d, d] output = torch.einsum('bnd,bde->bne', phi_Q, KV) # [batch, n, d] # 归一化 normalizer = torch.einsum('bnd,bmd->bn', phi_Q, phi_K).unsqueeze(-1) output = output / (normalizer + 1e-6) return output Performer的优势是无偏近似(随着特征数增大趋于精确),但需要较大的特征维度才能接近softmax注意力的性能。 ...

2026-07-02 · 2 min · 412 words · 硅基 AGI 探索者
Ring Attention分布式注意力

Ring Attention:分布式长上下文

超长上下文的分布式挑战 随着LLM上下文长度从32K扩展到1M甚至更长,单GPU的显存和计算能力已经远远不够。一个1M token的注意力矩阵在FP16下需要2TB显存——即使分布在8个GPU上,每张卡也需要256GB。 更关键的是,标准注意力计算需要在所有Query-Key对之间计算注意力分数,这在分布式环境下意味着大量的跨GPU通信。Ring Attention通过优雅的环形通信模式解决了这个问题。 核心思想:通信-计算重叠的环形流 Ring Attention的基本思想是:将长序列均匀切分到多个GPU上,每个GPU持有一段序列的Q、K、V。然后通过环形通信传递K/V块,同时与本地Q计算部分注意力。 环形拓扑 假设有N个GPU,编号为0到N-1。GPU i 持有序列的第 i 段的Q_i、K_i、V_i。计算过程如下: 第0步: GPU_i 用本地 Q_i 和 K_i 计算注意力,结果存为 partial_O_i 同时,GPU_i 将 K_i, V_i 发送给 GPU_{(i+1) % N} 第1步: GPU_i 接收 K_{(i-1) % N}, V_{(i-1) % N} 用 Q_i 和接收的 K 计算注意力,累加到 partial_O_i 同时发送 K_{(i-1) % N}, V_{(i-1) % N} 给 GPU_{(i+1) % N} ... (重复N-1步) 第N-1步: GPU_i 已与所有K/V块计算过注意力 partial_O_i 即为最终结果 关键在于:每一步中,通信和计算是并行进行的——GPU在接收下一块K/V的同时,用当前K/V进行注意力计算。 ...

2026-07-02 · 3 min · 454 words · 硅基 AGI 探索者
混合深度MoD

混合深度(MoD):动态计算深度

计算量的另一维度的稀疏化 MoE在"宽度"维度上实现了稀疏激活——不同token使用不同的专家。而Mixture-of-Depth(MoD)在"深度"维度上实现稀疏化——不同token经过不同数量的Transformer层。 这个想法直指一个核心观察:不是所有token都需要同等深度的计算。简单token(如停用词、常见短语)可能在几层之后就已经获得了充分的表示,而复杂token(如多义词、需要推理的位置)则需要更深的处理。 MoD的工作原理 基本架构 在MoD架构中,每一层都有一个路由器,决定哪些token需要继续经过当前层的计算,哪些可以"跳过"当前层直接传递到下一层: class MoDLayer(nn.Module): def __init__(self, d_model, n_heads, d_ff, capacity_ratio=0.5): super().__init__() self.router = nn.Linear(d_model, 1) # 输出一个标量路由分数 self.self_attn = MultiHeadAttention(d_model, n_heads) self.ffn = FFN(d_model, d_ff) self.norm1 = nn.RMSNorm(d_model) self.norm2 = nn.RMSNorm(d_model) self.capacity_ratio = capacity_ratio # 参与计算的token比例 def forward(self, x): batch_size, seq_len, d_model = x.shape # 路由决策 router_scores = self.router(x).squeeze(-1) # [batch, seq_len] n_process = int(seq_len * self.capacity_ratio) # 选择top-k个token参与计算 _, process_indices = torch.topk(router_scores, n_process, dim=-1) process_mask = torch.zeros_like(router_scores, dtype=torch.bool) process_mask.scatter_(1, process_indices, True) # 只有被选中的token经过注意力计算 processed = x.clone() if process_mask.any(): selected = x[process_mask].unsqueeze(0) attn_out = self.self_attn(self.norm1(selected)) ffn_out = self.ffn(self.norm2(attn_out)) selected = selected + attn_out + ffn_out processed[process_mask] = selected.squeeze(0) # 未选中的token直接传递(残差连接的等价效果) return processed 路由器的设计 路由器极其简单——只是一个单层线性映射加sigmoid,输出一个标量分数。这个分数表示"这个token需要当前层处理的程度"。 选择策略有两种: Top-k选择:固定比例的token参与计算,保证计算量可预测 阈值选择:分数超过阈值的token参与计算,计算量动态变化 实践中,Top-k选择更受青睐,因为它保证了计算量的可预测性,便于批处理和资源调度。 与MoE的关系与区别 MoD和MoE是正交的两种稀疏化策略: 维度 MoE MoD 稀疏方向 宽度(选择不同专家) 深度(选择是否计算) 路由目标 哪个专家 是否参与 计算节省 减少FFN计算量 减少整层计算量 可组合 — 可以与MoE叠加 两者可以组合使用:在MoE层的FFN部分使用专家路由,同时在层级别使用MoD路由。这种"双重稀疏"可以进一步降低计算量。 ...

2026-07-02 · 2 min · 254 words · 硅基 AGI 探索者
鲁ICP备2026018361号