大模型训练与 AI Infra

今天复习模型训练这一块,主要是针对预训练还有后训练使用的方法框架,AI Infra相关就偏向于了解了

训练速度

从训练速度开始复习起,数据量越来越多,训练的速度就成了问题,硬件方面是算力资源,网络连接,存储系统三者结合(了解)。

软件方面

并行计算

  • 张量并行,Megatron-LM 将 Transformer 权重矩阵按行列拆分,不同 GPU 计算矩阵乘法子集,通过通信整合结果,解决单 GPU 容纳难题。我使用的训练框架是ms-swift,这个框架一样可以做到并行计算
  • 数据并行, 将数据分拆到不同 GPU,同步梯度更新参数,实现简单但受限于单 GPU 内存。DeepSpeed ZeRO 优化梯度同步,扩展其适用范围。ms-swift也有
  • 流水线并行:按层拆分模型到不同 GPU,采用微批次(micro-batch)使不同批次并行计算。Microsoft 交错式调度可在内存不变下提升 10% 吞吐量。就是把模型拆开,分别放在不同gpu上逐个运行(举个例子,类似于工程,gpu0只负责第一步,gpu1只负责第二步,gpu2只负责第三部,假设一个数据经历3s才能前向完成,流水线并行1s出一个数据,如果一个gpu负责三步,3s出一个数据因而可以加速)

内存优化

  • 参数分区:ZeRO 三阶段分别分区优化器状态、梯度、参数,ZeRO-3 使单 GPU 支持模型规模比传统数据并行提升 8 倍(Rajbhandari et al., 2020; 2023)。(现算现用)

  • 内存卸载:ZeRO-Infinity 将不常用模型状态卸载到 CPU/NVMe,单节点可训练超 1 万亿参数模型;2025 年新增动态内存卸载技术,通过实时计算任务优先级调整卸载策略,内存利用率提升 90%,避免“过度卸载”导致的通信延迟。

  • 计算融合:合并 LayerNorm、dropout 等操作成单算子,Megatron-LM 实现 20% 内存减少与 15% 速度提升。

  • 混合精度训练:NVIDIA AMP 自动选择 16/8 位精度,内存减半且速度提升 2 倍;FP8 量化结合精度补偿算法(如 AdaRound),模型精度损失<0.5% ,进一步降低内存需求;INT4 量化技术,使模型大小减少 90%,推理速度提升 4 倍,且支持量化感知训练(QAT),压缩后模型性能保留率达 98%。

  • 稀疏内存访问:MoE 模型通过动态路由算法,激活参数占比仅 7%,推理延迟降至 0.8ms/token(华为盘古实测);内存安全增强方面,同态加密(HEAAN 库)应用于模型 Checkpoint 存储,实现加密状态下的参数更新与梯度计算,敏感任务数据泄露风险降为 0。

ms-swift都有,内部实现还是使用的deepspeed,张量并行也是使用的Megatron-LM

模型本身优化

  • Transformer 优化:从原始结构到 GPT 因果语言模型、T5 encoder-decoder,注意力机制、层归一化不断优化。FlashAttention 通过内存访问重构,速度提升 2-4 倍(Dao et al., 2022);Flash Attention-4 进一步优化 IO 效率,速度再提升 30%。(flash attention 实现的是将Qkv分块进行计算,局部实现softmax(与全局效果一致),为什么:计算硬件有两种SRAM计算速度快但是内存小,因此分块后放在这个上面进行计算,能达到加速的效果)

  • 混合专家(MoE)模型:Switch Transformer 等通过 “专家集合 + 动态路由”,保持参数规模同时降低计算量(Fedus et al., 2022);2023 年 Switch Transformers 2.0 支持 100 万亿参数模型训练,通过动态专家选择算法,计算成本仅为密集模型的 1/5(Fedus et al., 2023)。AI Infra 需优化专家负载均衡与动态路由,如采用“预测性路由”减少专家切换延迟,实现硬件 - 软件协同。

  • 模型压缩技术:INT8 量化减 75% 模型大小,结构化剪枝减 50% 参数;2025 年INT4 量化技术商用化,模型大小减少 90%,且通过量化感知训练(QAT)使性能保留率达 98%;AI Infra 支持量化算子库与混合精度训练框架,压缩模型训练效率提升 2 倍。

预训练与后训练

预训练

预训练的算法,对于GPT模型来说就是预测下一个token,对于bert模型来说有两种,一个是给予上下文预测mask位置的token,另一个是预测两句话是否属于上下连贯的一句,在bert的变种中,robert,macbert,bart,都删除了第二个任务,还有优化就是把mask位置改成近义词,扩大了数据集。

Bert及变种
模型 原第二任务 是否删除 替代任务 Mask 位置改动 数据集扩大
RoBERTa NSP 删除 无(仅用 MLM) 改为动态 mask(每轮重新采样) 中文:160 GiB → 全词 n-gram 掩码
MacBERT NSP 删除 SOP(句子顺序预测) 近义词替 [MASK];
80 % 同义词 + 10 % 随机 + 10 % 原词
中文:继续放大,SOP 负例翻倍
BART NSP 删除 Text Infilling(跨段空白填充) Span-mask(连续片段→[MASK]) 英文:160 GB → CC-News + OpenWebText

后训练

  • 参数高效微调:
    主要是指PEFT的几种高效微调的方式,

  • 指令微调sft:
    对于预训练模型来说,sft的作用是让模型获得对话的能力,准确来说是获得指令遵循的能力,因为在预训练中,模型的训练数据大部分是文章,与对话有差距,因此通过构造对话数据进行微调让模型进行对话的能力;对于chat模型来说,chat模型一般经历了上述过程,因此对他进行sft是为了让模型在垂直领域更好的回答。

  • RLHF :
    人类偏好对齐,一般使用强化学习的方式,强化学习的底层原理是马尔科夫决策过程,马尔科夫决策过程由⟨S,A,P,r,γ⟩⟨S,A,P,r,γ⟩S,A,P,r,γ组成。
    SSS是状态的集合;
    A\mathcal{A}A是动作的集合;
    γ\gammaγ是折扣因子;
    r(s,a)r ( s , a )r(s,a)是奖励函数,此时奖励可以同时取决于状态s ss和动作a aa,在奖励函数只取决于状态s ss时,则退化为r ( s ) r(s)r(s)
    P(s′∣s,a)P ( s ′ ∣ s , a )P(ss,a)是状态转移函数,表示在状态sss执行动作aaa之后到达状态s′s^{\prime}s 的概率。
    以试错的方式学习到更好的策略。简单而言就是模型根据策略做出动作,奖励给出打分,模型根据打分调整策略,在NLP中,通过之前的token预测下一个token,这个过程和模型根据策略做出动作的过程非常相似。
    下面主要复习NLP领域中的RL是怎么做的包括DPO,PPO,GRPO的实现

NLP中的RL

模型根据prompt生成一句话,奖励模型根据模型的输出进行打分。在其他领域中的RL,模型的决策是一个多步的过程,要给每一步打分(其实就是状态奖励和动作奖励,理解为一个为长远考虑,另一个为当前考虑),一般是给每一步都乘折扣因子,最后综合的结果就是奖励,在NLP中,是只针对最后的结果进行打分。
如果仅仅使用奖励模型进行训练,很容易训飞,输出不连贯很奇怪的句子,因此一般加入kl散度作为正则化项,这里kl散度衡量的是sft模型与RL模型之间的差异,限制模型不能偏移sft模型太远。

奖励模型的训练

奖励模型的训练是通过对模型生成的多个答案进行排序来进行训练的,为什么用排序代替打分呢,因为使用模型直接打分对于每个答案来说会有较大差异,使用排序来代替能很好的减少这种差异。
当然最后给模型的奖励仍然得是数字,那他是怎么做的呢就要使用以下的损失函数
在这里插入图片描述

这个损失函数实现了模型的打分要符合排序,对不同答案打分的差值使用sigmod取概率(归一化),取负对数作为损失(概率越接近1loss越低,概率越接近0loss接近无穷),负对数图像如下,前面的1/(k_2)是指取出正负样本
负对数图像

RLHF的过程

RLHF有三个数据集:

  • 偏好数据集:训练奖励模型
  • 提示数据集:训练llm
  • 评估数据集:格式同提示数据集
PPO

PPO涉及四个模型

  • 策略模型(Policy Model,或Actor Model),生成模型回复。它就是RLHF训练希望产出的模型。在PPO训练中更新梯度。

  • 奖励模型(Reward Model),输出奖励分数来评估回复质量的好坏。在PPO训练中不更新梯度。

  • 评论模型(Critic Model,或Value Network),来预测回复的好坏,可以在训练过程中实时调整模型,选择对未来累积收益最大的行为。在PPO训练中更新梯度。

  • 参考模型(Reference Model)提供了一个 SFT 模型的备份,帮助模型不会出现过于极端的变化。在PPO训练中不更新梯度。
    PPO(Proximal Policy Optimization)是一种流行的强化学习算法,它旨在解决传统策略梯度方法中的一些问题,如策略更新时的不稳定性和高方差。PPO通过限制策略更新的幅度来实现更稳定和高效的学习。

PPO的目标函数(也称为目标策略)是:

LCLIP(θ)=E^t[min⁡(rt(θ)A^t,clip(rt(θ),1−ϵ,1+ϵ)A^t)]L^{CLIP}(\theta) = \hat{\mathbb{E}}_{t}\left[ \min \left( r_t(\theta) \hat{A}_t, \text{clip}(r_t(\theta), 1 - \epsilon, 1 + \epsilon) \hat{A}_t \right) \right]LCLIP(θ)=E^t[min(rt(θ)A^t,clip(rt(θ),1ϵ,1+ϵ)A^t)]

其中:

  • θ\thetaθ 是策略网络的参数。

  • E^t\hat{\mathbb{E}}_{t}E^t表示对时间步 ( t ) 的期望值,通常是对一批样本的平均。

  • $r_t(\theta) $是策略比率,定义为:

    rt(θ)=πθ(at∣st)πθold(at∣st)r_t(\theta) = \frac{\pi_\theta(a_t | s_t)}{\pi_{\theta_{\text{old}}}(a_t | s_t)}rt(θ)=πθold(atst)πθ(atst)

    其中 πθ\pi_\thetaπθ 是新策略,πθold\pi_{\theta_{\text{old}}}πθold是旧策略,ata_tat是在状态sts_tst 下采取的动作。

  • ( A^t\hat{A}_tA^t ) 是优势函数的估计值,通常使用广义优势估计(Generalized Advantage Estimation, GAE)来计算。

  • clip(rt(θ),1−ϵ,1+ϵ)\text{clip}(r_t(\theta), 1 - \epsilon, 1 + \epsilon)clip(rt(θ),1ϵ,1+ϵ) 是一个裁剪函数,它将 rt(θ)r_t(\theta)rt(θ) 的值限制在1−ϵ1 - \epsilon1ϵ1+ϵ1 + \epsilon1+ϵ 之间,其中 ϵ\epsilonϵ是一个超参数,用于控制策略更新的幅度。
    以上的损失函数也体现了PPO两个有助于模型训练稳定性和效率的两个重要机制:clip机制和重要性采样,重要性采样就是策略比例。

ms-swift-PPO

在ms-swift中PPO的实现流程是这样的,先采样 → 算奖励/GAE → 用 PPO-clip 更新策略 + 用 clipped-MSE 更新价值函数 → 循环直到 episode 足够。
算奖励:将生成的证据给奖励模型获得奖励,计算KL散度,给每个token位置-kl散度作为惩罚,只给最后一个token位置加reward,GAE反向计算每个位置的回报。最终的结果就是,每个token位置除最后token位置之外,都是-kl+回报,最后一个位置是-kl+奖励。

logr        = ref_logprobs - logprobs                # 重要性采样比的对数
kl          = -logr                                   # k1 估计器
non_score_reward = -args.kl_coef * kl               # 负号→惩罚
rewards     = non_score_reward.clone()              # 先填上每 token 的 KL 惩罚
actual_end  = 句子真实长度(不含 pad)
rewards[ [batch_idx, actual_end] ] += scores        # 只有最后一个有效位置加“大奖”
lastgaelam = 0
advantages_reversed = []
gen_length = responses.shape[1]          # 最大生成长度
for t in reversed(range(gen_length)):    # ← 从后往前扫
    nextvalues = values[:, t+1] if t < gen_length-1 else 0.0
    delta = rewards[:, t] + args.gamma * nextvalues - values[:, t]
    lastgaelam = delta + args.gamma * args.lam * lastgaelam
    advantages_reversed.append(lastgaelam)
advantages = torch.stack(advantages_reversed[::-1], 1)  # 再翻回来
returns = advantages + values # (价值网络的目标)
DPO

DPO是通过直接学习偏好数据的差异来更新模型参数,将奖励隐含在loss中
LDPO(θ)=−logσ(β[log(πθ(yw∣x)/πref(yw∣x))−log(πθ(yl∣x)/πref(yl∣x))])L_DPO(θ) = -log σ( β [log(πθ(y_w|x)/πref(y_w|x)) - log(πθ(y_l|x)/πref(y_l|x))] )LDPO(θ)=logσ(β[log(πθ(ywx)/πref(ywx))log(πθ(ylx)/πref(ylx))])

DPO变种 LD-DPO
ms-swift-LD_DPO

ms-swift使用的LD-DPO,混杂了sft的损失,用于稳定训练

GRPO
ms-swift GRPO

评估

基础能力评估:Perplexity 测预测能力,BLEU/ROUGE 测生成一致性;AI Infra 通过并行计算,将千亿参数模型 Perplexity 计算从数小时缩至分钟级;2025 年实时评估工具上线,支持训练过程中每 1000 步自动计算 Perplexity,评估延迟<30 秒。

综合能力评估:MMLU(57 科目)、GPQA(专业知识)、SWE-bench(代码)等基准需大量计算;2025 年评估基准新增多模态能力维度(如 M3LUE),涵盖图文音视频理解与生成;AI Infra 通过批量处理与结果缓存,降低重复评估成本,相同模型多次评估时间缩短 80%。

安全与对齐评估:HELM Safety、AIR-Bench 测安全性、偏见与事实性;2025 年动态安全测试技术,通过对抗样本生成自动检测模型漏洞,安全风险识别率达 95%;AI Infra 通过专用环境与隐私计算(如差分隐私+同态加密),保障敏感内容评估安全,数据泄露风险降为 0。

评估自动化:将评估整合进开发流水线,自动生成报告;斯坦福报告显示,自动化评估使迭代周期缩短 40%;2025 年评估结果可视化平台,支持多维指标(性能、效率、安全性)实时展示与对比,决策效率提升 50%。

参考:
https://infrasys-ai.github.io/aiinfra-docs/00Summary/04TrainingStack.html
https://blog.csdn.net/hffwj/article/details/121423957?ops_request_misc=elastic_search_misc&request_id=a4ebfc23b93b4bdbe2fcf0b34beff8fe&biz_id=0&utm_medium=distribute.pc_search_result.none-task-blog-2allsobaiduend~default-2-121423957-null-null.142v102pc_search_result_base3&utm_term=NVMe%20SSD&spm=1018.2226.3001.4187

Logo

中国智能体开发者社区,聚焦智能体与大模型开发,提供前沿资讯、实用工具链、开源项目及行业案例。通过技术沙龙、开发者大赛等活动,促进经验交流与协作,助力开发者快速构建创新智能应用。

更多推荐