智能计算系统复习第4-6章

发表于 分类于 本文字数： 8.8k

摘要

第四章 编程框架使用

编程框架概述

为什么需要编程框架

  • 深度学习算法具有多层结构，每层的运算由一些基本操作构成
  • 这些基本操作中存在大量共性运算及操作，如基础操作、大模型算子、分布式逻辑等。将这些共性操作封装成标准化组件，可以提高编程实现效率
  • 面向这些封装起来的操作，硬件程序员可以基于硬件特征，有针对性的进行充分优化，使其能充分发挥硬件的效率

定义

  • 将深度学习算法中的基本操作封装成一系列组件，这一系列深度学习组件，即构成一套深度学习框架
  • 编程框架能够帮助算法开发人员更简单的实现已有算法，或设计新的算法。也有助于硬件程序员更有针对性的对关键操作进行优化，使其能充分发挥硬件效率

代表性框架分类：

类型 代表性框架 特点
通用深度学习框架 PyTorch、TensorFlow 开发灵活
大模型专用框架 Megatron-LM、DeepSpeed 基于 PyTorch 二次开发，大规模训练
推理框架 vLLM 大模型推理

PyTorch概述

PyTorch起源-Torch

  • Torch是瑞士亚普研究所（IDIAP）在2002年发布的一款机器学习框架，采用LuaJIT脚本编程语言为接口，内核采用C/C++来实现
  • 核心是易于使用的神经网络和优化库，同时在实现复杂的神经网络拓扑结构方面具有最大的灵活性
  • 在GPU上具有较高的性能

PyTorch简介

  • PyTorch = Py + Torch
  • PyTorch是一个基于Torch的Python开源机器学习库，用于自然语言处理等应用程序，具有强大的GPU加速的张量计算（如Numpy）能力
  • 产生背景：编程语言提供了面向深度学习的高效编程库（NumPy、Eigen、Torch等）；Python开源生态蓬勃发展

2.3 PyTorch 2.x

  • 编译优化：torch.compile，兼顾动态图灵活度以及静态图高性能
  • 大模型生态支持：用于大模型训练的DeepSpeed、Megatron-LM，用于大模型推理的vLLM

PyTorch编程模型及基本用法

张量（tensor）

  • 张量是计算图上的数据载体，用张量统一表示所有的数据，张量在计算图的节点之间传递
  • 张量对应了神经网络中在各个节点之间传递、流动的数据
  • 张量可以看做是n维的数组，数组的维数即为张量的阶数
  • 与NumPy中的ndarray不同，PyTorch中的张量可以运行在GPU或深度学习处理器上，因此具有较高的性能

张量的阶数与数据形式：

阶数 对应数据形式
0 标量
1 向量
2 矩阵
n n维数组

张量的常用属性：

属性名 含义
dtype tensor存储的数据类型
shape tensor各阶的长度
device 存储tensor的设备对象
grad 默认为None，当调用backward()进行反向传播后为该tensor的梯度值

张量的数据类型（dtype）：

PyTorch数据类型 说明
torch.float16/torch.half 16位浮点
torch.float32/torch.float 32位浮点
torch.float64/torch.double 64位浮点
torch.float8_e4m3fn/torch.float8_e5m2 8位浮点
torch.uint8 8位无符号整型
torch.int8 8位整型
torch.int16/torch.short 16位整型
torch.int32/torch.int 32位整型
torch.int64/torch.long 64位整型
torch.bool 布尔型
torch.bfloat16 Brain Float16

智能计算涉及到的数据类型：

格式 符号位 指数位（范围） 小数位（精度） 总位数
FP64 1 11 52 64
FP32 1 8 23 32
TF32 1 8 10 19
FP16 1 5 10 16
BF16 1 8 7 16
FP8 E4M3 1 4 3 8
FP8 E5M2 1 5 2 8
FP4 1 2 1 4
  • FP8 E4M3更适合权重、激活数据，适合前向计算（DeepSeek-R1）。FP8 E5M2更适合梯度数据，适合反向计算
  • 许多框架会在一个模型中混合使用这两种格式：前向传播用E4M3，反向传播用E5M2，以平衡精度和范围需求

INT8与FP8：

  • INT8是定点数，表示范围固定[-128, 127]，精度均匀
  • FP8是浮点数，具有指数位，可表示更大范围的数值，但精度随数值变化。越靠近0，分布越稠密，越远离0，分布越稀疏
  • Transformer模型中，权重分布大部分在0附近，因此在Transformer模型中，使用FP8量化效果好于INT8

INT8和FP8不存在哪一个有绝对优势，可以根据不同的层的实际loss进行选择。相对来说，INT8量化更适合于权重，FP8-E4量化更适合于大多数层的激活值

硬件实现上，FP8和INT8的乘加（MAC）单元的面积几乎相同

微缩放（Microscaling）

由AMD、ARM、Intel、Meta、Microsoft、NVIDIA、Qualcomm联合发布的数据格式，是深度学习中一种先进的低精度数据表示方法

通过精细化调整模型参数、数据格式和计算粒度，深度优化深度学习模型训练和推理效率

将张量分割为固定大小的子块（blocks），为每个子块分配独立的缩放因子（scaling factor），并使用窄位宽格式（如FP8、INT4等）存储块内元素

与传统的全张量共享单一缩放因子相比，Microscaling能在保持模型精度的同时，大幅提升硬件效率与存储压缩比

每k个数据（P1,…,Pk）共享一个缩放因子X，这k个数据称为一个block，每个block占用（w+kd）bits

具体的MX格式：

  • 每32个数据共享相同的scale，scale数据类型为FP8（E8M0），无符号位

使用MX格式的优势：

  • 在多项任务里，MX直接替代FP32推理、用于低比特训练，能够接近/对齐FP32/BF16的精度
  • 已有多个硬件（英伟达Blackwell、寒武纪等）架构支持MX格式，在英伟达平台上，MXFP8相比BF16的矩阵核，吞吐达到近2倍

张量的形状属性（shape）

表示张量中每一个维度（axis）的尺寸

image-20260612163436164

张量的device属性

torch.device指定tensor所在的设备名、设备序号

  • 用'cuda:n'表示第n个GPU设备，用'dlp:n'表示第n个深度学习处理器

张量属性的转换：

  • tensor.to()：进行张量的数据类型或设备类型转换
  • 将GPU上的张量转换成NumPy数组，需要先将张量转换到CPU上，再转换成NumPy
  • tensor.reshape(*shape)：进行张量的形状属性转换。单个维度上的shape值可以为-1，表示该维度上的shape值需要根据其他维度的shape来推算
  • 张量从CPU转换到DLP上：a.dlp() 或 a.to(torch.device('dlp'))

张量的复制：

  • tensor.clone()：对张量进行复制，返回一个完全相同的tensor，新张量会保存在新的内存中，且仍然留在计算图中

张量的数据格式（data layout）

  • 张量数据可以有多种数据格式，代表了多维数组以何种线性存储方式在存储空间中存储
  • PyTorch、GPU中采用NCHW，TensorFlow、CPU中采用NHWC
  • N：一批次的数据个数（batch size），C：通道数（channel），H：高度（height），W：宽度（width）

一张RGB彩色图像包含3个原色（红色、绿色、蓝色）通道，对应的张量表示中，N=1，C=3

image-20260612163554420

用于大模型计算优化的张量格式NHD/HND：N为batch 或 sequence length，H为注意力头数量，D为每个注意力头的隐藏层维度。

  • NHD相邻内存属于同一个token的不同head，适合按token并行处理。

  • HND相邻内存属于同一个head的不同token，适合按head并行处理，对GPU实现更友好

张量的自动求导支持

  • PyTorch支持自动求导，用户定义好操作的前向计算和反向梯度计算规则，PyTorch能够在训练时自动调用计算图算子，完成整个网络的自动求导
  • 使用**requires_grad参数来设置张量是否需要自动求导**，默认为false
  • 对于一个计算操作来说，如果所有输入中有一个输入需要求导，则输出就需要求导；如果所有输入都不需要求导，则输出也不需要求导

操作(operation)

  • PyTorch基于张量开展各种类型的计算操作。每个操作接收若干个张量作为输入，操作完成后更新原张量或生成新张量作为输出
  • 计算操作是使用PyTorch实现模型训练和推理的基础
  • 可以采用torch.operation、tensor.operation等形式来实现张量的计算操作

操作的广播（broadcasting）机制：

  • 对于参与计算操作的多个张量，如果张量维度不匹配，可以使用PyTorch的广播机制对不匹配的张量维度进行扩展，最终将这些张量均扩展为维度相同
  • 能够进行广播机制的条件：每个张量都有至少1个维度；从张量末尾的维度开始对齐扩展，在对齐后的同一维度中，仅下列情况之一才允许进行广播操作：1）维度尺寸相同；2）维度尺寸不同但其中一个维度尺寸为1；3）其中一个张量没有该维度
  • 对于维度数量相同的张量，比较每个维度对应的维度尺寸，若维度尺寸不同但其中一个维度尺寸为1，则将其维度尺寸扩展为另一张量的维度尺寸
  • 对于维度数量不同的张量，首先从张量末尾的维度开始对齐扩展，对缺少的维度尺寸补1，再沿每个维度方向进行尺寸对比及扩展

计算图

编程框架中使用有向图来描述计算过程。有向图中包含一组节点和边

  • 计算图对应了神经网络的结构

节点和边：

  • 节点一般用来表示各类操作，包括数学运算、变量读写、数据填充等，也可以表示输入数据、模型参数、输出数据
  • 边表示”节点”之间的输入输出关系。分为两类：一类是传递具体数据的边，传递的数据即为张量（tensor）；一类是表示节点之间控制依赖关系的边，这类边不传递数据，只表示节点执行的顺序：必须前序节点计算完成，后序节点才开始计算

静态图 vs. 动态图：

  • 静态图：先定义整张图，再运行；可以对图进行全局优化，获得更快的运算速度；调试不方便
  • 动态图：即时运行，网络模型可在运行时修改；代码编写灵活，可立即获得执行结果，调试方便；优化不方便
  • TensorFlow 1.x为静态图，PyTorch为动态图
  • 静态图：构建静态图，每一次iteration中重复执行同样的图
  • 动态图：每一次iteration中构建并执行新图，在反向传播结束之后，整个计算图就在内存中被释放了

torch.compile：

  • PyTorch 2.x中的重要升级，实现了动态图语义+静态图性能的统一，一行代码即可获得显著加速

基于PyTorch的模型推理实现

image-20260612164042260

读取输入图像

构建神经网络

  • 可以自定义神经网络模型，也可以直接调用PyTorch框架中提供的模型
  • 自定义模型：PyTorch提供torch.nn、torch.nn.Module、torch.nn.functional等模块
  • 直接调用预训练模型：PyTorch提供torchvision.models，包含了用于处理不同任务的各种模型，如图像分类、语义分割、目标检测、关键点检测等

torch.nn.Module：

  • PyTorch使用模块（module）来表示神经网络
  • torch.nn.Module是用于封装PyTorch模型及组件的基类
  • 包含了__init__以及forward方法等
  • 自定义网络模型的方法：需要继承torch.nn.Module类；首先通过__init__方法初始化整个模型，定义模型结构及待学习的参数，再使用forward方法定义模型的前向计算过程；PyTorch支持模型的自动梯度计算，因此在forward()中无需定义反向计算过程

parameter与buffer：

  • 模块类（module）中包含两种不同状态的参数：
    • parameters：可学习的参数，即反向传播时可以被优化器更新的参数
    • buffers：不可学习的参数，即反向传播时不可以被优化器更新的参数
  • Parameter被保存在state_dict之中
  • Buffer分成两种：persistent和non-persistent。前者保存在state_dict中，而后者不包含在state_dict中
  • 保存模型时保存的是包含在state_dict中的参数（parameter、persistent buffer）

torch.nn中的计算功能：

计算类型 计算操作
卷积层 nn.Conv1d, nn.Conv2d, nn.Conv3d
池化层 nn.MaxPool1d, nn.MaxPool2d, nn.MaxPool3d, nn.AvgPool1d, nn.AvgPool2d, nn.AvgPool3d
非线性激活 nn.ELU, nn.ReLU, nn.SELU, nn.Sigmoid, nn.Tanh
归一化层 nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.InstanceNorm1d, nn.InstanceNorm2d, nn.InstanceNorm3d
循环神经网络层 nn.RNN, nn.LSTM, nn.GRU
线性层 nn.Linear
损失函数 nn.L1Loss, nn.MSELoss, nn.CrossEntropyLoss

直接调用预训练模型：

  • 基于某一种网络结构，首先在一个初始任务场景、初始数据集上训练好一个模型，然后再应用到目标任务上，针对目标任务的特征、数据集，对训练好的模型进行微调（fine-tune），最终满足目标任务的需求
  • 不需要从零开始训练模型，只需要调用已有的预训练模型参数，进行简单的微调，能够节省大量的计算资源和计算时间

预训练模型的应用：

  • 每种预训练模型对其输入均有各自的规格要求（尺寸、像素值等），在应用该预训练模型之前，需要根据这些要求对输入进行预处理
  • Torchvision针对每种权重对输入的要求，提供了预处理方法，可以通过使用.transforms()来实现

model.train与model.eval：

  • 在训练开始之前，代码中添加model.train()，在测试时添加model.eval()
  • 二者差别主要体现在对Batch Normalization和Dropout操作的处理上，model.eval()的作用是不启用BN和Dropout
  • Dropout：评估模式下，测试的模型应该是最终得到的模型，而这个模型应该是一个完整的模型
  • BN层主要涉及到四个需要更新的参数：均值running_mean、方差running_var、weight(gamma)、bias(beta)
  • model.train()保证训练时每经过一个mini-batch都会更新这批数据的均值和方差
  • model.eval()控制BN层中的running_mean、running_std不更新，采用训练结束后的running_mean、running_std来规范化图像

Hugging Face的transformers库：

  • Hugging Face是目前全球最活跃的AI开源社区，被称为”AI界的GitHub”
  • Transformers库是Hugging Face生态的核心Python库，专注自然语言处理类相关的处理
  • 常用API包括：AutoTokenizer、AutoModelForCausalLM、AutoModelForSequenceClassification、Trainer、model.generate()、pipeline

实例化神经网络模型

  • 采用如下步骤进行神经网络模型的实例化：
    • 完成神经网络模块（module）的定义，包括__init__()方法和forward()方法的定义
    • 对模型中的参数（weight、bias等）进行初始化
    • 实例化模型结构，将结构相关参数传递给module
    • 定义模型的输入数据
    • 将模型输入传入实例化后的模型，获取模型输出

使用torch.nn.init模块完成初始化：

  • module中的parameter和buffer参数默认为CPU上执行的32位浮点数，在定义module时可以将其设置为任意的数据类型及设备类型
  • 在module的构建函数__init__()中，可以使用torch.nn.init模块来对parameter和buffer参数进行初始化

torchmetrics库：

  • 为评估深度学习模型的功能、性能，需要计算各种指标
  • torchmetrics库提供了一组评估指标计算工具，包括准确率、精确率、召回率、F1分数、均方误差（MSE）、平均绝对误差（MAE）等

调试

  • 使用python调试工具（使用交互式调试库python的pdb，进行设置断点、单步执行、查看代码、变量、参数等调试功能）
  • 打印模型结构、参数等信息
  • 使用TensorBoard实现数据可视化

使用TensorBoard实现数据可视化：

  • TensorBoard提供了对神经网络模型结构、参数等的可视化功能，可用于查看、分析神经网络模型的结构、权重、损失值、精度等
  • PyTorch提供torch.utils.tensorboard，用于引入tensorboard工具，实现神经网络模型及张量的可视化

神经网络模型优化

  • 使用torch.nn.utils.prune对神经网络模型进行剪枝操作
  • 使用torch.ao.quantization进行神经网络模型的量化
  • 使用torch.compile进行神经网络模型的性能优化

神经网络模型剪枝：

  • 神经网络所包含的大量参数中，有一部分是冗余且对输出结果无贡献的，将这些参数裁减掉，可以提升网络的计算速度，且不影响最终的精度
  • 非结构化剪枝：按一定规则对单个参数进行裁剪
  • 结构化剪枝：按一定结构规则对一组参数进行裁剪，如裁减掉整个卷积核，或裁减掉卷积核的某一行或某一列

动态量化：

  • 神经网络模型的量化，指将模型中的权重和/或激活数据从多位宽的浮点格式，转换为低位宽的整数型格式
  • 量化过程中使用到的缩放系数与参与量化张量的数值范围有关
  • 权重参数的数值范围是固定的，其缩放系数也是固定的
  • 动态量化：每一层激活数据的数值范围随计算过程变化，需要动态确定缩放系数，根据缩放系数动态完成数据格式转换
  • PyTorch提供torch.ao.quantization.quantize_dynamic()用于模型的动态量化

性能优化工具 torch.compile：

  • 传统PyTorch的动态图模式，存在解释器开销、且算子碎片化
  • torch.compile在保留动态图的灵活性的同时，获得静态图的性能优势
  • 对计算图自动应用算子融合、内存优化、硬件适配、向量化/并行化等优化技术

大模型推理引擎：

  • HuggingFace Transformers库存在显存碎片化严重、并发能力差、计算内核利用率低等问题，无法满足生产环境高吞吐、低延迟的需求
  • 主流大模型推理引擎：vLLM（基于PagedAttention技术，高吞吐，易于使用，社区生态最活跃）、TensorRT-LLM（英伟达官网工具，深度优化GPU性能）、TGI（基于Hugging Face模型，易部署）

vLLM大模型推理引擎：

  • PagedAttention技术：借鉴操作系统的虚拟内存分页思想。将KV Cache切分为固定大小的Block（页），逻辑上连续，物理上离散存储。为每个请求按需分配块，避免预留式内存浪费
  • 高吞吐量、低延迟、高内存效率
  • PyTorch作为基础：使用PyTorch来定义和加载模型，PyTorch提供了灵活的模型定义和自动微分功能
  • 推理引擎优化：算子优化、内存管理、调度优化、量化支持等

基于PyTorch的模型训练实现

加载训练数据集

应用于大模型场景的流式数据集：

  • 大模型训练的数据集通常达到TB甚至PB级别，无法一次性加载到内存中
  • 流式数据集加载方式：不将整个数据集一次性加载到内存中，而是在训练过程中动态地、逐批次地从磁盘或网络读取数据
  • 代表性工具：Hugging Face的datasets库

模型训练

  • 使用PyTorch进行模型训练时，首先需要定义损失函数的计算方法，然后构建优化器实现对模型的梯度计算及更新
  • 反向传播过程中，可以利用内建的性能分析工具、梯度检查函数等，验证训练过程的正确性和有效性

损失函数定义：

  • 损失函数可以自己定义，也可以直接使用PyTorch提供的内建损失函数
  • 自定义损失函数可以通过定义模块实例来实现
  • 内建损失函数：nn.L1Loss()（计算平均绝对误差MAE）、nn.MSELoss()（计算均方误差MSE）、nn.CrossEntropyLoss()（计算交叉熵损失函数）、nn.NLLLoss（计算负对数似然损失）、nn.BCELoss（计算二分类交叉熵损失函数）

神经网络模型训练：

  • 损失函数定义完成后，需要使用优化器对模型进行训练，包括计算梯度，优化梯度并更新参数等步骤
  • 计算梯度时需要构建计算图并进行反向传播，根据计算图中的张量属性，控制其节点是否需要计算梯度
  • PyTorch提供torch.optim包来实现多种梯度优化算法
  • 训练流程：构建优化目标 -> 计算预测值 -> 计算损失函数 -> 对参数梯度清零 -> 计算梯度 -> 优化梯度，更新参数

构建计算图来计算梯度：

  • 计算图的主要作用在于能够在前向计算过程中保存所有中间节点的计算结果，便于反向传播时构建反向传播路径，并利用链式法则完成自动求导
  • 计算图在一次反向传播后会被立即销毁，释放存储空间，下次调用时需要再次创建
  • 只有在训练时计算图是必需的，而如果只是单纯的推理，可以选择不创建计算图，以节省存储占用和资源消耗

使用torch.no_grad禁用计算图：

  • 可以使用torch.no_grad上下文管理器，在其作用域内定义的所有计算，仍然可以前向传播得到计算输出，但不会反向传播计算梯度，也不会创建计算图

计算图的反向传播：

  • 对于requires_grad属性为True的张量，其前向计算过程中的后继张量，除了使用no_grad管理的，其它张量均默认requires_grad=True，即需要计算梯度，需要创建计算图
  • 可以使用tensor.backward()函数计算当前张量相对于计算图中所有requires_grad属性也为True张量的梯度
  • 每次调用backward()后，计算图会被释放掉

通过detach()方法修改计算图：

  • tensor.detach()：返回一个新张量，该张量从当前计算图中剥离，成为一个新的张量，新张量的requires_grad=False，即不需要计算梯度
  • 返回的张量与原张量共享相同内存，对原张量或新张量的原位修改（如尺寸、stride等的原位修改）均会报错

使用torch.optim包优化梯度：

  • 完成梯度计算后，可以使用torch.optim包来优化梯度、更新模型参数
  • torch.opim.Optimizer(params,defaults)：所有优化器的基类。其中，params为需要优化的模型参数列表，defaults为包含了如learning rate等优化选项的字典

模型的保存与恢复

使用torch.save()保存模型

  • 官方推荐：可以仅保存模型的state_dict，模型保存格式为.pt或.pth。使用model.load_state_dict()恢复模型
  • 也可以自定义保存的内容，如将模型的state_dict、优化器的state_dict、epoch、loss值等一起保存为检查点，检查点文件的常见保存格式为.tar。使用torch.load()恢复模型

大模型参数的高效保存：

  • 70B模型的state_dict约140GB（BF16）
  • 存储格式：Safetensors
  • 核心技术：分片保存、量化保存、激活检查点
  • 分片保存技术：将模型参数切分为多个连续的文件块，配合一个索引文件（model.safetensors.index.json），记录每个参数张量所在的文件分片及偏移量

激活检查点技术：

  • 在大模型训练中，正向传播过程中产生的大量激活值需要存储在内存中，以便在反向传播时计算梯度。随着模型参数量的增加（如千亿级参数），激活值的存储成为训练的主要内存瓶颈
  • 激活检查点（计算换内存）
    • 正向传播时，只保存部分层的激活值；
    • 反向传播时，如果发现需要的激活值不在显存中，则利用最近的”检查点”保存的输入，临时重新进行前向传播，计算出所需的激活值

第五章 编程框架原理

编程框架设计

  • 智能计算系统中编程框架的四大模块：
    • 必备：计算图构建模块和计算图执行模块
    • 追求更高性能：深度学习编译模块和分布式训练模块

设计原则

  • 简洁性（Simplicity）：框架提供一套抽象机制，用户仅需关心算法本身和部署策略
  • 易用性（Usability）：
    • 熟悉的开发范式：如PyTorch始于Python，忠于Python
    • 直观且用户友好的接口：如PyTorch提供了命令式的动态图编程方法
  • 高效性（Performance）：
    • 如采用静态图编程方式，可以生成完整的计算图并进行全局优化，从而尽量提高用户应用程序的运行效率
    • 支持深度学习编译技术，多层级表示优化，充分利用用户硬件的计算能力
    • 支持多机多卡条件的分布式训练，从而高效支持大规模深度学习任务

整体架构

  • 计算图构建模块：完成从输入的用户程序到编程框架内部原始计算图的转换过程，编程框架的入口模块
  • 分布式训练模块：应对更大规模的神经网络，将训练、推理任务从一台设备扩展到多台设备
  • 深度学习编译模块：对计算图分别进行图层级和算子层级的编译优化，从而提升单设备上的执行效率
  • 计算图执行模块：将优化后的计算图中的张量和操作映射到指定设备上进行具体执行，并给出编程框架的输出结果

image-20260612175041537

计算图构建

正向图与反向图构建

  • 计算图由两个基本元素构成：张量（Tensor）和张量操作（Operation）。计算图是有向图，有向边指明了张量的流动方向

正向传播

输入张量经过搭建的神经网络层层计算传递，并最终获得计算结果的过程

  • 动态图：在执行函数时，按照函数顺序逐条语句地生成节点，立即计算并返回结果；易调试但性能优化空间有限
    • 计算图在函数运行过程中逐步构建的（On-the-fly）
    • 立即（eager）模式：每次调用语句就立刻执行计算
    • PyTorch中的动态图实现：每次执行，都会重新被构建
  • 静态图：在执行计算之前构建好所有图上的节点，在图运行时才计算整个计算图并返回最终结果；不易调试但性能好
    • 整个网络的结构会在开始计算前就建立完成计算图
    • 框架执行时接收整个计算图而不是单一语句

反向传播

计算导数的方法：

自动微分

  • 手动求导：用链式法则求解出梯度公式，然后根据公式编写代码、代入数值计算得到梯度结果。缺点：对于大规模的深度学习算法非常困难
  • 数值求导：直接代入数值近似求解。优点：易操作，可对用户隐藏求解过程。缺点：计算量大，求解速度慢，可能引起舍入误差和截断误差
  • 符号求导：直接对代数表达式求解，最后才代入问题数字，出现表达式膨胀问题
  • 自动求导：用户只需描述前向计算的过程，由编程框架自动推导反向计算图，先建立表达式，再代入数值计算

自动求导法：

  • 介于数值求导和符号求导的方法：对基本算子应用符号求导法，代入数值，保留中间结果，应用于整个函数
  • 计算分两步执行：
    • 1）原始函数建立计算图，数据正向传播，计算出中间节点xi，并记录计算图中的节点依赖关系
    • 2）反向遍历计算图，计算输出对于每个节点的导数
  • 对于前向计算中一个数据连接多个输出数据的情况，自动求导中，将这些输出数据相对于该数据的导数累加

求导方式对比：

方法 对图的遍历次数 精度 备注
手动求解法 NA 高 实现复杂
数值求导法 nI+1 低 计算量大，速度慢
符号求导法 NA 高 表达式膨胀
自动求导法 nO+1 高 对输入维度较大的情况优势明显

PyTorch中的自动混合精度（AMP）：

  • AMP是一种在编程框架中混合使用高精度（FP32）和低精度（FP16、BF16）浮点数的技术，能够在保持模型精度的同时显著提升训练性能并减少内存使用
  • 动态选择数据类型：对精度不敏感的线性层、卷积层等采用低精度（FP16、BF16）；精度敏感的归约操作等采用FP32高精度

image-20260612175710353

典型使用场景：

  • 训练场景：目标是在保证模型收敛的数值稳定性的基础上，大幅提升计算速度并节省显存（允许更大Batch Size）。核心策略：遵循”低精度计算，高精度参数更新“。
  • 推理场景：目标是在保证模型推理精度的基础上，提升模型推理性能。核心策略：按需采用低精度运算

核心机制：

  • 操作分类与动态分派（torch.autocast）：PyTorch将算子预设为两类，自动判断精度选择
    • 计算密集型算子(FP16/BF16)，对精度相对不敏感：matmul、conv、linear、bmm等
    • 精度敏感型算子(保持FP32)，在FP16下极易发生数值溢出或严重损失精度：softmax、batchnorm、loss计算、log/exp、归一化层、激活函数（某些）
  • 梯度缩放(Loss Scaling)：问题来源：由于FP16指数位较短，许多微小的梯度会被截断为0。原理：缩放——在计算梯度后乘以一个巨大的缩放因子（如65536），将梯度推回FP16有效表示范围；反缩放——在利用梯度进行实际参数更新前，将梯度除以缩放因子，恢复真实量级
  • 权重副本（Master Weight Copy）：原理：在内存中同时保存FP16精度和FP32精度的副本，分别在前向和反向进行使用

计算图执行

将计算图中的张量和操作（又称算子）映射到给定设备上具体执行

  • 包括：设备管理、张量实现、算子执行（获取算子执行序列、实现算子：前端定义、后端实现、前后端绑定、查找并调用算子）

设备管理

  • 设备是编程框架中计算图执行时的硬件实体，每个设备都具体负责计算子图中的张量存放和算子运算
  • 常见设备包括通用处理器（如CPU）和领域专用处理器（如GPU和DLP等）
  • 添加对领域专用处理器的设备管理支持需要三个模块：设备操作、执行流管理、事件管理

张量实现

  • 逻辑视图：形状、布局、步长、偏移量、数据类型和设备等，是框架使用者能直接控制和表达的基本属性
  • 物理视图：设备上的物理地址空间大小、指针、数据类型等，对框架使用者不可见

一个物理视图可以对应多个逻辑视图：切片的结果不是新的物理视图，而是原本物理视图下的一个新的逻辑视图

PyTorch中的张量抽象：

PyTorch中存在与张量对应的类Tensor

  • 通过张量（Tensor）抽象类和存储（Storage）抽象类来分别表示张量数据结构中的逻辑视图和物理视图

张量内存分配：

从逻辑视图到物理视图的转换需要完成对张量的内存分配。根据设备的类型不同，张量管理的方式不同

  • 即时分配（CPU）：每当需要分配张量的内存时，就立即从系统中申请一块合适大小的内存空间。代码核心部分：malloc()和free()函数
  • 内存池分配（GPU）：预先分配一块固定大小的内存池，然后在需要时从内存池中分配内存。自我维护：内存块的拆分和合并。优点：节约设备内存使用，减少设备内存碎片化

算子执行

计算图的执行过程 = 每个算子独立执行的过程

  • 计算图 -> 执行序列（确保正确的数据流和依赖关系）
  • 针对每个算子进行算子实现：前端定义、后端实现和前后端绑定
  • 分派执行：查找适合给定输入的算子实现，并调用相应的实现来执行具体的计算任务（有点类似skills）

执行序列：

  • 分析计算图节点之间的依赖关系 -> 执行序列
  • 使用拓扑排序算法（可有多种可行的结果）

算子实现

  • 正向传播实现和反向传播实现分离
  • 用户接口（前端）和具体实现（后端）分离
  • 算子实现流程：
    • 前端定义（在编程框架中配置算子信息，包含算子的输入、输出以及相关的接口定义，最后生成前端接口如Python API）->
    • 后端实现（使用C++或其他高级的编程语言，编写算子的底层实现代码）->
    • 前后端绑定（编程框架将前端定义的算子与后端的具体实现进行绑定）

前后端绑定：同一个算子可能会有多个后端实现的代码，多种后端&多种输入，根据不同情况调用相应的后端实现。PyTorch使用分派机制来管理前后端分布式对应关系，由Dispatcher管理分派表。

分派执行：

在运行时根据输入张量的类型和设备类型查找并调用合适的算子实现方法。Dispatcher计算分派键，并由此找到对应的内核函数

  • 算子：Dispatcher的调度对象，代表了具体的计算任务
  • 分派键：根据输入张量和其他信息计算，可简单地理解为与硬件平台相关联的标识符
  • 内核函数：特定硬件平台上实现算子功能的具体代码