AI 计算模式（下）：轻量化网络与大模型分布式并行的计算模式分析

文档教程人工智能【免费下载链接】AISystemAISystem 主要是指AI系统包括AI芯片、AI编译器、AI推理和训练框架等AI全栈底层技术项目地址https://gitcode.com/GitHub_Trending/ai/AISystem点击查看免费下载AI 计算模式决定了 AI 芯片的计算单元设计、存储层次与互联带宽需求。本文基于 AISystem 项目《AI 计算体系》课程中的 AI 计算模式下 一节聚焦轻量化网络模型与大模型分布式并行两大主题从 Params/FLOPs 量化指标出发逐层拆解深度可分离卷积、特征复用等结构设计对芯片算力特征的影响并结合 数据并行、FSDP/ZeRO、张量并行 等仓库文档梳理集合通信原语与 DP/DDP/FSDP/TP/PP 全谱系并行策略最终落到芯片架构应当支持的互联、异构资源管理与端到端优化等计算模式要求上。轻量化网络模型在精度与算力之间做减法随着神经网络应用从云端下沉到移动端与嵌入式设备硬件平台普遍面临内存资源少、处理器性能不高、功耗受限的约束。轻量化网络模型的设计方向由此而来在保持模型精度的基础上进一步减少模型参数量与计算量让模型能够以更小的存储占用、更低的访存开销部署到受限设备上。轻量级的两把标尺Params 与 FLOPs衡量网络模型轻量级程度的两项核心指标是网络参数量Params与浮点运算数FLOPs。对于卷积神经网络中的卷积层二者的定义如下。对于输入为 $w \times h \times Ci$ 的图像、卷积核大小为 $k \times k$、输出特征图尺寸为 $W \times H \times Co$ 的卷积操作参数量$Params (k \times k \times Ci 1) \times Co$浮点运算数$FLOPs W \times H \times (k \times k \times Ci 1) \times Co$其中 $Ci$ 为输入通道数$Co$ 为输出通道数即卷积核个数$1$ 对应偏置项。直观来看参数量和计算量越小模型速度越快。但模型快慢不仅取决于参数量与 FLOPs还与内存访问次数、网络结构本身强相关——这正是从 AI 计算模式角度审视轻量化设计的出发点芯片不仅要算得快还要访存省、结构灵活。轻量化三大设计流派轻量化网络从三个维度改造卷积结构减少内存空间卷积核分解、减少通道数深度可分离卷积、减少卷积核个数特征复用。减少内存空间小卷积核分解大卷积核在 VGG 与 InceptionNet 系列中经典做法是用两个 $3 \times 3$ 卷积核替代一个 $5 \times 5$ 卷积核或用一个 $5 \times 1$ 卷积核加一个 $1 \times 5$ 卷积核替代一个 $5 \times 5$ 卷积核。以 $3 \times 3$ 替代 $5 \times 5$ 为例其核心收益有二感知野不变、网络更深两个串联的 $3 \times 3$ 卷积叠加后的感知野等价于一个 $5 \times 5$ 卷积但网络深度增加非线性表达能力更强一定程度上提升网络效果参数显著下降模型参数由 $5 \times 5 \times Ci \times Co$ 变为 $3 \times 3 \times Ci \times Co 3 \times 3 \times Ci \times Co$。假设 $Ci Co$该层参数可减小为原来的 $18/25$。这种大卷积核分解为小卷积核的模式要求芯片能够高效执行大量小尺寸卷积核的乘加运算并配合残差Skip Connection等连接结构提升深层网络的可训练性。减少通道数深度可分离卷积MobileNet 系列提出了深度可分离卷积Depthwise Separable Convolution通过 Depthwise 逐通道卷积加 $1 \times 1$ Pointwise 卷积来替代标准卷积。其中 Depthwise 卷积在每个输入通道上独立执行 $k \times k$ 卷积负责空间特征提取$1 \times 1$ 的 Pointwise 卷积负责通道维度的特征融合与通道数缩减从而大幅降低参数量。以 $3 \times 3$ 卷积核、输入通道 16、输出通道 32 的卷积层为例标准卷积参数量$3 \times 3 \times 16 \times 32 4608$深度可分离卷积参数量$3 \times 3 \times 16 1 \times 1 \times 16 \times 32 144 512 656$参数量从 4608 降至 656减少约 85%。从计算模式看深度可分离卷积将稠密的空间-通道联合卷积拆解为稀疏的空间卷积 稠密的 $1 \times 1$ 通道卷积对芯片意味着需要高效支持 $1 \times 1$ 卷积本质是大规模矩阵乘与低计算强度的 Depthwise 卷积并存二者对计算阵列的利用率与访存模式要求截然不同。减少卷积核个数特征图复用DenseNet 与 GhostNet 采用Reuse Feature Map的设计思路来减少参数与运算量。以 DenseNetV1 为例第 $n$ 层的输入不仅包含上一层输出还复用之前所有层的 Feature Map。由于复用了前 $n-1$ 层的 Feature Map第 $n$ 层的参数量由 $k \times k \times C1 \times (C1C2)$ 变为 $k \times k \times C1 \times C2$即原来的 $C2/(C1C2)$。其中 $k$ 为卷积核尺寸$C1$ 为前 $n-1$ 层 Feature Map 的总个数$C2$ 为第 $n$ 层输出 Feature Map 个数且 $C2 \ll C1$因此参数压缩非常可观。特征复用模式对芯片访存系统提出了要求同一份 Feature Map 需要被多个后续层反复读取这要求芯片具备足够大的片上缓存与高带宽的片上互联NoC以支撑数据的多次复用降低对外部存储的访问压力。AI 计算模式思考轻量化结构对芯片的启示从轻量化网络的设计分类可以看到AI 网络对卷积层存在多种不同的设计方法而这些方法正是芯片设计时需要重点支持的AI 计算模式特性可以归纳为三个层面卷积核尺寸层面小卷积核替代用多个小卷积核代替单个大卷积核降低计算成本多尺寸卷积核采用不同尺寸的卷积核捕捉多尺度特征可变形卷积核从固定形状转向可变形卷积核适应不同输入特征1×1 卷积核使用 1×1 卷积核构建 bottleneck 结构有效减少参数和计算量。卷积层运算层面Depthwise 卷积用 Depthwise 卷积代替标准卷积减少参数、保持特征表达Group 卷积应用分组卷积提高计算效率、降低模型复杂度Channel Shuffle通过通道混洗增强特征融合提升模型性能通道加权实施通道加权计算动态调整通道贡献优化特征表示。卷积层连接层面Skip Connection采用跳跃连接使网络更深同时避免梯度消失问题Dense Connection利用密集连接整合不同层特征增强特征融合与信息流。对 AI 芯片设计而言这些模式意味着计算单元需要支持异构的卷积形态标准卷积/深度可分离/分组/1×1、灵活的通道组织shuffle、加权以及大量重复的 Feature Map 访存。这与仓库中 AI 计算模式上 对经典网络结构全连接、卷积、循环、注意力与量化剪枝的分析一脉相承共同构成芯片指令集与数据通路设计的输入。大模型分布式并行超越单芯片的算力边界大模型具有超高的模型参数量与计算量。在单芯片或单加速卡无法提供所需算力与内存的前提下分布式并行成为必由之路。分布式并行总体分为数据并行与模型并行两大类模型并行又可细分为张量并行与流水线并行。在展开之前先理解并行计算中无处不在的基础设施——集合通信原语。集合通信原语分布式并行的消息机制通信原语是不同计算节点或设备之间进行数据传输与同步的基本操作是实现复杂并行算法与应用的基础。常见原语包括All-reduce将所有节点上的数据收集起来执行某种操作通常是求和或求平均再将结果广播回每个节点。常用于全局梯度更新。All-gather每个节点上的数据被广播到其他所有节点最终每个节点都收到来自所有节点的数据集合用于收集局部数据进行全局聚合或分析。Broadcast将一台节点上的数据广播到其他所有节点通常用于分发模型参数或全局数据。Reduce将所有节点数据执行某种操作求和、求平均、取最大值等后将结果发送回指定节点用于局部聚合。Scatter将一个节点的数据集合切分并分发到其他节点通常用于将大集合拆分后并行处理。Gather将各节点数据收集到一个节点上用于局部数据的汇总与分析。这些原语是 DP/DDP/FSDP/TP/PP 各类并行策略的通信基石。例如 DDP 依赖All-reduceRing-AllReduce做梯度同步FSDP/ZeRO 依赖All-gather / Reduce-scatter做参数与梯度分片交换张量并行依赖All-gather / All-reduce做部分结果拼接。芯片与网络必须高效实现这些原语分布式并行才有性能可言。数据并行DP → DDP → FSDP 的演进路径根据模型在设备之间的通信程度数据并行技术可分为DP、DDP、FSDP三种形态详见仓库 数据并行 与 FSDP/ZeRO 全分片数据并行。DP 数据并行最简单的起点DP 是最简单的分布式并行将大规模数据集分割为多个小批量每个批量被发送到不同计算设备如 NPU上并行处理。每个设备持有完整模型副本独立计算梯度再通过All-reduce机制同步更新模型参数以保持一致性。DP 存在明显局限单机单进程多线程受 Python GIL 限制梯度汇总集中在单一节点易造成负载不均扩展性与性能都受限。DDP 分布式数据并行多进程 Ring-AllReduceDDP 允许模型在多个计算节点上并行训练每个节点有本地模型副本与本地数据。训练流程为每个节点的模型副本执行前向/反向传播并计算梯度 → 梯度在节点间通信并取平均默认使用Ring-AllReduce算法→ 所有节点用全局梯度更新本地参数。DDP 相比 DP 的关键改进在于多进程实现避开 Python GIL 限制可扩展到多机多卡Ring-AllReduce 与计算通信重叠反向传播过程中梯度就绪即触发集合通信通过 autograd hook 参数分桶 reducer 实现将通信延迟隐藏在计算中大幅提高集群利用率负载均衡不再依赖单一节点聚合。DDP 通常与 AI 框架深度集成例如 PyTorch 的torch.nn.parallel.DistributedDataParallel模块可自动处理模型与梯度同步及分布式通信。仓库 数据并行 给出了完整的 2 卡 DDP 训练示例先通过dist.init_process_group(nccl, rankrank, world_sizeworld_size)初始化进程组用DDP(model, device_ids[rank])封装模型配合torch.multiprocessing.spawn多进程启动真实场景还需DistributedSampler保证各设备拿到互不重叠的数据分片。FSDP 全分片数据并行DP 与 DDP 的融合升级FSDPFully Sharded Data Parallelism是 DP/DDP 的结合升级版实现更高效的训练与更好的横向扩展。其核心思想是将神经网络的权重、梯度、优化器状态全部进行分片shard分配到不同设备上并行处理需要时再通过集合通信获取。由于分享全部模型参数、梯度与优化状态FSDP 在计算阶段需要频繁执行参数、梯度和优化状态的同步通信操作——如图所示不同计算节点之间存在大量虚线连接的通信操作。FSDP 是零冗余优化器 ZeRO 在 AI 框架中的落地实现ZeRO-DP 通过优化器状态分区ZeRO-1内存降至约 1/4、梯度分区ZeRO-2内存降至约 1/8、参数分区ZeRO-3内存与并行度 $N_d$ 成反比三阶段将每卡内存占用从 $4\Psi K\Psi$ 逐步降至 $(4\Psi K\Psi)/N_d$ZeRO-R 则针对激活检查点分区、恒定大小缓冲区与内存碎片整理做剩余状态优化ZeRO-Infinity 进一步将参数在 NPU/CPU/NVMe 间卸载迁移支撑万亿参数级模型训练。通信代价上ZeRO-3 相比标准 DP 的通信量约为 1.5 倍——以通信换显存是 FSDP 的根本权衡。仓库 FSDP/ZeRO 全分片数据并行 详细给出了torch.distributed.fsdp.FullyShardedDataParallel的封装用法以及通过ShardingStrategy.SHARD_GRAD_OP对应 ZeRO-2与FULL_SHARD对应 ZeRO-3选择分片策略的代码示例。模型并行张量并行与流水线并行模型并行将模型本身拆分到多设备分为张量并行TP与流水线并行PP两种详见仓库 张量并行。张量并行层内切分张量并行将模型的张量操作分解为多个子张量操作在不同设备上并行执行。其收益是将大模型计算负载分布到多设备提高计算效率与训练速度代价是必须设计合适的通信机制来交换数据、同步参数通常使用All-reduce / All-gather等原语完成梯度聚合与参数同步。以矩阵乘算子的张量并行为例设 $X$ 为激活输入、$A$ 为算子权重将 $A$按列切分每个计算节点保留完整的 $X$ 与部分的 $A$各自计算部分乘积最后通过All-gather通信将两节点数据同步拼接为完整的 $Y$ 输出供下一层使用。仓库 张量并行 进一步介绍了 Megatron-LM 提出的细粒度切分方案MLP 层对权重 $A$ 列切分、$B$ 行切分自注意力层对 Q/K/V 按列切分配合列切分ColwiseParallel、行切分RowwiseParallel与序列并行SequenceParallel等 DeviceMesh/DTensor 原语以及 Embedding 的表切分/列切分与 Cross Entropy Loss 并行完整支撑大模型各层级的并行化。流水线并行层间切分 空泡优化流水线并行将模型的不同层划分为多个阶段每个阶段部署在不同设备上执行设备各自计算模型的一部分把结果传递给下一设备形成计算流水线。设计要点包括合适的数据流与通信机制缓冲区与流水线控制器确保计算正确性与一致性。以一个具有 Forward/Backward 两个阶段、0-3 共 4 层网络、分布在 4 个计算设备的流水线为例右图展示了时间维度下不同层、不同阶段的执行顺序。每个设备在切换阶段时会存在等待空闲区域称为Bubble 空泡降低设备利用率。一个简单而有效的优化是增加数据并行度让每层数据切分为若干个 batch 送入流水线使各设备在等待时处理其他 batch从而提高流水线设备利用率。这正是 Megatron-LM 等大模型训练框架的实践方向——以更多的通信量换取更低的空泡率。AI 计算模式思考并行策略对芯片架构的要求不同并行策略深刻揭示了 AI 计算模式在硬件设计上的体现芯片架构设计可从以下四个方面考量模型并行与数据并行支持AI 芯片需同时支持模型并行与数据并行。对模型并行芯片需具备灵活的计算资源分配与通信机制支撑模型不同部分在多个设备上计算对数据并行芯片需提供高带宽、低延迟的通信与同步机制支撑设备间的数据交换与同步。异构计算资源管理AI 芯片通常包含 CPU、GPU、TPU 等多种计算资源。分布式并行计算要求芯片提供统一的异构资源管理机制实现不同计算资源的协同工作与资源调度。高效的通信与同步机制分布式并行伴随大量数据交换与同步操作。芯片需提供高效的通信与同步机制降低通信延迟、提高通信带宽从而支撑高效分布式计算。端到端的优化AI 芯片需支持从模型设计、算法优化到系统设计的端到端优化。例如 Transformer 是众多大模型结构的基础组件芯片可提供专用高速 Transformer 引擎如对矩阵乘、注意力计算的硬件加速从计算模式源头提升大模型训练与推理效率。小结与思考从轻量化网络到大模型分布式并行AI 计算模式对芯片设计提出了系统性的要求AI 芯片需具备对多样神经网络模型架构的灵活支持与高效执行特有计算逻辑的能力以适应不同应用需求AI 芯片应支持模型压缩算法如量化、剪枝提高模型终端部署时的推理性能实现软件算法与硬件执行的高效协同AI 芯片设计应考虑轻量化网络结构支持更复杂的卷积运算与数据逻辑以适应算力和带宽受限的场景AI 芯片需支持大模型的分布式并行策略包括高效的片上网络接口与总线设计以及大内存容量和高速互联带宽以应对多芯片堆叠提高性能的需求。这些思考与仓库中 AI 计算模式上经典网络结构与量化剪枝、关键设计指标、核心计算矩阵乘、数据单位比特位 等课程内容前后呼应共同勾勒出算法演进 → 计算模式 → 芯片设计的完整链路。读者可结合 02Hardware 硬件基础目录 与 05Framework 框架并行目录 继续深入对应主题。赞分享文档教程人工智能【免费下载链接】AISystemAISystem 主要是指AI系统包括AI芯片、AI编译器、AI推理和训练框架等AI全栈底层技术项目地址https://gitcode.com/GitHub_Trending/ai/AISystem点击查看免费下载相关推荐ETL容器性能优化缓存友好的内存布局设计ETL容器性能优化缓存友好的内存布局设计 Embedded Template LibraryETL作为轻量级嵌入式开发工具包其容器性能优化的核心在于 缓嵌入式开发工具ORB_SLAM2并行BA优化分布式计算加速大规模优化ORB_SLAM2并行BA优化分布式计算加速大规模优化 SLAMSimultaneous Localization and Mapping同步定位与地图构计算机视觉机器人科研ollama模型并行计算多GPU分布式训练方案ollama模型并行计算多GPU分布式训练方案 引言多GPU并行的必要性与挑战 随着大语言模型LLM参数量呈指数级增长从Llama 2的70B到GPT人工智能大模型模型推理服务本地部署创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考