主页

Megatron Float16Module 的参数传法

Megatron 的 Float16Module 只会把位置参数转换成半精度(bf16/fp16)。一旦只用关键字参数,输入会保持原 dtype,导致前向混用 fp32/bf16。 1 相关代码 代码位置:megatron/core/transformer/module.py#L404 https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/transformer/module.py#L404 def forward(self, *inputs, fp32_output=True, **kwargs): if parallel_state.is_pipeline_first_stage(i...

阅读更多

ApexRMSNorm vs TENorm:bf16 大 weight 精度坑

bf16 下使用 ApexRMSNorm,一旦 weight/gamma/scale 放大,输出标准差看似稳定,但均值会随机机漂移;换成 Transformer Engine 的 TENorm 就恢复稳定。 1 RMSNorm RMSNorm 只做均方根归一化,不减均值、也没有 beta 偏置: \[y = \frac{x}{\sqrt{\frac{1}{d} \sum_{i=1}^{d} x_i^2 + \epsilon}} \odot \gamma\] d 为 hidden size,γ 逐通道缩放,通常初始化为 1。 由于不减均值,输出是否“居零”高度依赖舍入误差;γ 越大,误差越容易被放大。 2 现象 在 bf16 下把 γ 放大(例如 >3.5,或...

阅读更多

TP 下的激活显存公式

大模型训练时,激活显存往往比参数显存更先成为瓶颈。为了在同样的硬件预算下堆更长的序列或更大的 batch,我们会把 Transformer 的矩阵乘拆到多张 GPU 上,也就是常说的 Tensor Parallel (TP)。这篇笔记把原始草稿整理成一篇小博文,记录 TP 如何影响每层激活的内存占用,以及从公式里可以读出的工程直觉。 1 每层激活显存的估算式 \[\text{Activations memory per layer} = s b h \left( 10 + \frac{24}{t} + 5 \frac{a s}{h t} \right)\] 这个式子描述了 单层 Transformer 激活显存 与多种变量之间的关系:序列长度 $s$、micro-batch 大小...

阅读更多

ZeRO 究竟算数据并行还是模型并行?

Zero(ZeRO, Zero Redundancy Optimizer)总被问:它究竟是数据并行(DP)的技巧,还是模型并行(MP)的变体? 本文把 ZeRO 放回数据并行的语境,解释它如何削减显存冗余、为何仍维持 DP 语义,以及和真正模型并行的根本区别。 1 数据并行的痛点:冗余状态 传统 DP 在每张卡上复制 完整模型: 参数、梯度、优化器状态都各存一份; Backward 结束后 All-Reduce 梯度,保持副本一致。 随着模型变大,显存中堆着 N 份(N=GPU 数)相同状态,浪费巨大。ZeRO 的目标就是保持 DP 训练语义不变,却让这些状态不再 N 份复制。 2 ZeRO 的核心:分区状态而非拆计算 ZeRO 把模型状态切成片,平均分到不同 ...

阅读更多

MoE 实践梳理:并行策略、随机性与 DeepSeek

MoE 三个问题:如何并行、为何随机、怎样稳定 1. MoE 的模型并行 vs 数据并行 MoE 训练通常把“专家”视为模型并行单元,再与数据并行组合。下图直观展示了权重如何在 data parallel、model parallel 与 expert parallel 之间拆分: Model Parallel:同一批 token 被切到不同设备,各自负责部分 FFN/attention 计算;适合巨型专家参数。 Data Parallel:复制一份完整模型,分批处理不同样本;用来堆 throughput。 Expert Parallel:把专家实例平均分发到多卡,token 依据路由结果被 permute 到对应设备。 实战中常出现“三明治”拓扑:DP 包...

阅读更多

注意力头设计

在不牺牲模型质量的前提下,如何从注意力头这一层入手,压缩推理显存、带宽与算力成本 1 绝大多数LLM为何仍守着标准多头 Transformer 诞生至今,最常见的还是经典多头注意力(MHA): \(\text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V\) 每个 head 拥有独立的 Q、K、V,与其他 head 并行计算。GPT-3、LLaMA、Qwen 这些主流 LLM 之所以鲜少改动,原因主要有三: 训练可复现:生态里所有框架、算子库都为 MHA 做了极致优化,工程风险最小。 表达力充足:每个 head 面向不同子空间,模型容量不受限制。 推理成本仍可接受...

阅读更多

Softmax 数值稳定性的危机

大模型训练里,最常被忽略的炸弹往往埋在最熟悉的算子下。Softmax 同时包含指数与除法,一旦输入 logits 偏离正常尺度,就会把梯度链条整段炸成 NaN。本文把课堂随笔整理成一篇可查的博客,集中梳理 softmax 在注意力与输出层中的不稳定来源,以及工业界常用的工程缓解手段 1 Softmax 为什么会失稳 Softmax 把向量 $\mathbf{z}$ 转成概率分布: \(p_i = \frac{e^{z_i}}{\sum_j e^{z_j}}\) 在自注意力中它直接作用在 $\frac{QK^\top}{\sqrt{d_k}}$ 上。问题在于,$QK^\top$ 的取值范围会随 hidden size、head dim、训练阶段而大幅波动,导致 softmax 出现...

阅读更多

Transformer 超参数的取舍

从 small model 到万亿 token 级的大模型,表现和效率往往被少量关键超参数左右:hidden size、前馈层扩张倍数、注意力头数与 head dim、词表大小、正则化策略以及深宽比例(aspect ratio)。随着参数预算与训练数据呈指数级增长,工程团队反而更依赖这些“旋钮”来稳定收益——一次错误的配置可能让数百万美元的训练算力付诸流水。本文把我在课程里的零散记录整理成一篇可查的博客,并加入实践中反复被问到的背景解释,方便之后设计或复现 Transformer。 阅读顺序上,可以先浏览第 1 节的速览表迅速建立数量级感知,再根据实际需求跳到单独的小节。如果你正打算把某篇论文的配置迁移到自己的项目里,建议直接对照第 7 节的顺序进行 sanity check。 1...

阅读更多