拓冰网站定制 
目录第三章PyTorch 与资源核算3.1 为什么需要资源核算计算量Compute显存Memory3.2 Tensor深度学习的基本数据结构3.3 MemoryTensor 占多少内存为什么 LLM 常使用 BF16Tensor 的底层存储3.4 Compute这些 Tensor 要计算多少FLOPs 与实际 GPU 性能MFUModel FLOPs Utilization自动求导与反向传播3.5 Model Training把所有组件组装起来数据加载Mixed Precision3.6 Arithmetic Intensity 与 RooflineArithmetic IntensityMemory-BoundCompute-BoundRoofline Model为什么训练和推理表现不同3.13.6 的整体逻辑第三章总示意图本章核心总结第三章PyTorch 与资源核算本章主线为什么要算资源 → 数据在 PyTorch 中是什么 → 数据占多少显存 → 计算需要多少算力 → 如何组成完整训练系统 → 为什么 GPU 实际性能达不到理论峰值。3.1 为什么需要资源核算训练 LLM 前一个重要的工程问题是这个模型训不训得动需要多少显存要训练多久资源核算主要关注两个维度1. 计算量Compute计算量通常使用FLOPsFloating Point Operations衡量即训练过程中需要执行多少次浮点运算。对于 LLM可以根据模型参数量训练 Token 数量GPU 理论算力实际 GPU 利用率MFU快速估算整个训练过程所需的计算量和时间。2. 显存Memory训练模型占用的显存并不只有模型参数还包括Model ParametersGradientsOptimizer StatesActivations例如使用FP32 AdamW训练时每个参数粗略需要Parameter 4 Bytes Gradient 4 Bytes Adam First Moment 4 Bytes Adam Second Moment 4 Bytes -------------------------- Total ≈ 16 Bytes / Parameter而这甚至还没有包含 Activation。因此模型“放得下”不代表模型“训得动”。核心理解训练前先做 Napkin Math算力决定训练多久显存决定能不能训练。3.2 Tensor深度学习的基本数据结构有了资源核算意识下一个问题是GPU 到底在存储和计算什么答案是Tensor张量。PyTorch 中几乎所有核心数据最终都是 TensorData Parameters Gradients Optimizer States Activations ↓ Tensor理解 Tensor 时需要重点关注shaperankdtypedevicestride例如 Transformer 中经常出现[Batch, Sequence, Heads, Hidden]Tensor 之间会进行viewreshapetransposeelement-wise operationMatrix Multiplication其中最核心的计算是Matrix MultiplicationMatMul此外jaxtyping和einops可以帮助我们更加明确地表达 Tensor 的维度变化例如batch × sequence × heads × hidden核心理解Tensor 是深度学习世界的数据载体MatMul 是最核心的计算。3.3 MemoryTensor 占多少内存既然模型中的参数、梯度和 Activation 都是 Tensor那么自然会产生下一个问题一个 Tensor 到底占多少显存基本公式Tensor Memory ≈ Number of Elements × Bytes per Element因此dtype会直接影响显存占用。Data Type每元素大小特点FP324 Bytes精度高、稳定但显存开销大FP162 Bytes快、省显存但动态范围较小BF162 Bytes动态范围接近 FP32LLM 训练常用FP81 Byte更省显存但数值稳定性要求更高为什么 LLM 常使用 BF16BF16 与 FP16 都只需要 2 Bytes但 BF16 保留了更大的 exponent 范围。因此FP16 精度较高 但数值范围较小 BF16 精度略低 但数值范围更大对于容易出现极大或极小数值的深度神经网络BF16 通常具有更好的训练稳定性。Tensor 的底层存储一个 Tensor 并不只是数据本身还可以理解为Tensor ├── Data Pointer ├── Shape ├── Stride ├── Dtype └── Device很多 Tensor 操作例如view()transpose()slice()可能只是改变“如何解释底层数据”而不会真正复制数据。但contiguous()可能产生新的内存复制。此外CPU 与 GPU 之间的数据传输本身也存在成本CPU RAM │ │ PCIe ↓ GPU HBM因此优化 LLM 不仅需要减少计算也需要尽量减少不必要的数据搬运。核心理解显存优化的本质是管理 Tensor 的数量、dtype、布局以及数据移动。3.4 Compute这些 Tensor 要计算多少解决 Memory 问题之后下一个问题是这些 Tensor 到底需要多少计算这里最核心的指标就是FLOPs。对于矩阵乘法[M, K] × [K, N]计算量大约为2 × M × K × N FLOPs因此模型规模、Batch Size、Sequence Length 等都会直接影响训练计算量。FLOPs 与实际 GPU 性能理论 FLOPs 并不代表 GPU 实际能够达到的性能。因此需要引入MFUModel FLOPs UtilizationMFU 实际 FLOP/s ──────────── GPU 理论 FLOP/sMFU 衡量的是GPU 的理论计算能力到底真正利用了多少。例如 GPU 理论峰值为1000 TFLOPS而模型实际只跑到500 TFLOPS那么MFU 50%自动求导与反向传播PyTorch 会自动构建计算图Input ↓ Forward ↓ Loss ↓ Backward ↓ Gradient通过loss.backward()PyTorch 根据链式法则计算每个参数的梯度。训练过程因此包含Forward Compute Backward Compute通常反向传播所需计算量还会高于单纯的 Forward。核心理解FLOPs 告诉我们理论工作量MFU 告诉我们 GPU 实际完成了多少工作。3.5 Model Training把所有组件组装起来前面学习的是训练系统的各个“零件”Tensor Memory Compute Gradient这一节开始把它们组合成完整的训练流程。典型的 PyTorch Training LoopInitialize Parameters ↓ Load Data ↓ Forward ↓ Loss ↓ Backward ↓ Gradient ↓ Optimizer ↓ Update Parameters ↓ Next Batch模型中的可学习参数通过nn.Parameter进行管理。参数初始化也非常重要例如 Xavier Initialization其目标之一是避免深层网络中出现Activation Explosion or Activation Vanishing数据加载当训练数据大于 CPU 内存时可以使用memmap避免一次性将整个数据集加载进 RAM。GPU 训练时还希望形成CPU 准备 Batch N1 │ │ asynchronous transfer ↓ GPU 计算 Batch N通过Pinned MemoryAsynchronous Transfer让 CPU 数据准备与 GPU 计算重叠从而减少 GPU 等待时间。Mixed Precision训练时还可以混合使用FP32 BF16 FP16在尽量保持数值稳定性的同时降低显存占用提高 Tensor Core 利用率提升训练速度核心理解3.5 是整个章节的组装环节Tensor Memory Compute Data Gradient Optimizer → Training。3.6 Arithmetic Intensity 与 Roofline最后需要回答一个非常重要的问题为什么 GPU 理论上有 1000 TFLOPS我的程序却远远跑不到 1000 TFLOPS原因是 GPU 不仅需要计算还需要不断搬运数据HBM ↓ Load Data ↓ Compute Core ↓ Compute ↓ Store Data ↓ HBM所以程序性能同时受到两个因素影响Compute Capability Memory BandwidthArithmetic Intensity定义Arithmetic Intensity FLOPs ────── Bytes它表示每从 Memory 搬运 1 Byte 数据可以完成多少次浮点计算。Memory-Bound如果搬很多数据 ↓ 只做很少计算Arithmetic Intensity 就很低。此时性能主要受Memory Bandwidth限制。例如ReLU Element-wise Operations Vector Operations通常更容易 Memory-Bound。Compute-Bound如果搬一次数据 ↓ 进行大量计算Arithmetic Intensity 较高。例如Large Matrix Multiplication这时性能主要受 GPU 的计算能力限制Compute Capability即 Compute-Bound。Roofline ModelRoofline Model 将两种瓶颈统一起来转折点可以理解为Arithmetic Intensity ≈ Peak FLOPs ────────────── Memory Bandwidth低于这个点Memory-Bound高于这个点Compute-Bound为什么训练和推理表现不同这个模型还能解释 LLM 中一个非常重要的现象。Training训练中大量执行Matrix × MatrixArithmetic Intensity 较高因此更容易Compute-BoundGPU 的 Tensor Core 可以得到较充分利用。Autoregressive Inference逐 Token 推理时经常更接近Matrix × Vector每生成一个 Token 都需要读取大量模型参数但计算量相对有限。因此更容易Memory-Bound这也是为什么LLM 推理性能不仅取决于 GPU FLOPs也高度依赖显存带宽。3.13.6 的整体逻辑-问题形式整个第三章可以看成六个连续的问题3.1 Resource Accounting 为什么要做资源核算 │ ↓ LLM 很大、训练很贵 │ ├───────────────┐ ↓ ↓ Memory Compute 能不能放下 要计算多久 │ │ ↓ │ 3.2 Tensor │ 训练中的基本对象 │ │ │ ↓ │ Parameters / Gradients │ Activations / Data │ │ │ ↓ ↓ 3.3 Memory 3.4 Compute Tensor占多少显存 需要多少FLOPs │ │ └───────┬───────┘ ↓ 3.5 Training │ Data → Model → Loss ↓ Backward ↓ Optimizer ↓ Update Weight │ ↓ 为什么GPU还是不够快 │ ↓ 3.6 Roofline │ Arithmetic Intensity FLOPs / Bytes │ ┌───────┴────────┐ ↓ ↓ Memory-Bound Compute-Bound │ │ └───────┬────────┘ ↓ GPU Performance第三章总示意图LLM TRAINING │ ▼ ┌─────────────────────────┐ │ 3.1 Resource Accounting │ │ 训练需要多少资源 │ └────────────┬────────────┘ │ ┌───────────┴───────────┐ │ │ ▼ ▼ MEMORY COMPUTE 能不能训 要训多久 │ │ ▼ │ ┌─────────────┐ │ │ 3.2 Tensor │ │ └──────┬──────┘ │ │ │ ┌──────────┼──────────┐ │ │ │ │ │ ▼ ▼ ▼ │ Parameters Gradients Activations │ │ │ │ │ └──────────┼──────────┘ │ ▼ ▼ ┌────────────────┐ ┌────────────────┐ │ 3.3 Memory │ │ 3.4 Compute │ │ │ │ │ │ Shape × Dtype │ │ MatMul → FLOPs │ │ FP32/BF16/FP8 │ │ MFU │ └────────┬───────┘ └────────┬───────┘ │ │ └───────────┬───────────┘ ▼ ┌──────────────────┐ │ 3.5 Training │ │ │ │ Data → Forward │ │ ↓ │ │ Loss │ │ ↓ │ │ Backward │ │ ↓ │ │ Optimizer │ │ ↓ │ │ Update Weight │ └────────┬─────────┘ │ ▼ 实际 GPU 为什么不够快 │ ▼ ┌──────────────────┐ │ 3.6 Roofline │ │ │ │ FLOPs / Bytes │ │ ↓ │ │ Arithmetic │ │ Intensity │ └────────┬─────────┘ │ ┌───────────┴───────────┐ │ │ ▼ ▼ MEMORY-BOUND COMPUTE-BOUND 显存带宽限制 GPU算力限制 │ │ └───────────┬───────────┘ ▼ GPU UTILIZATION本章核心总结第三章真正想建立的不是单纯的PyTorch API 使用能力而是一套LLM Systems Thinking模型训练的本质是 Tensor 在 Memory 中存储和移动同时 GPU 对 Tensor 执行大量计算。因此整章可以压缩成一条逻辑链Resource ↓ Tensor ↓ Memory Compute ↓ Training ↓ Performance ↓ Roofline进一步来看后续很多 LLM 系统优化技术其实都可以放进这个框架Mixed Precision ↓ 减少 Memory 提升 Compute Gradient Checkpointing ↓ 用 Compute 换 Memory FlashAttention ↓ 减少 HBM 数据搬运 KV Cache ↓ 用 Memory 换重复 Compute Distributed Training ↓ 多 GPU 分摊 Memory Compute所以理解这一章之后再学习Transformer、FlashAttention、KV Cache、Mixed Precision、Gradient Checkpointing 和 Distributed Training可以始终追问两个核心问题1. 这个技术改变了多少 Memory2. 这个技术改变了多少 Compute / Data Movement这两个问题基本贯穿整个 LLM Systems 优化体系。 diy-llm 第3章 PyTorch 与资源核算
目录第三章PyTorch 与资源核算3.1 为什么需要资源核算计算量Compute显存Memory3.2 Tensor深度学习的基本数据结构3.3 MemoryTensor 占多少内存为什么 LLM 常使用 BF16Tensor 的底层存储3.4 Compute这些 Tensor 要计算多少FLOPs 与实际 GPU 性能MFUModel FLOPs Utilization自动求导与反向传播3.5 Model Training把所有组件组装起来数据加载Mixed Precision3.6 Arithmetic Intensity 与 RooflineArithmetic IntensityMemory-BoundCompute-BoundRoofline Model为什么训练和推理表现不同3.13.6 的整体逻辑第三章总示意图本章核心总结第三章PyTorch 与资源核算本章主线为什么要算资源 → 数据在 PyTorch 中是什么 → 数据占多少显存 → 计算需要多少算力 → 如何组成完整训练系统 → 为什么 GPU 实际性能达不到理论峰值。3.1 为什么需要资源核算训练 LLM 前一个重要的工程问题是这个模型训不训得动需要多少显存要训练多久资源核算主要关注两个维度1. 计算量Compute计算量通常使用FLOPsFloating Point Operations衡量即训练过程中需要执行多少次浮点运算。对于 LLM可以根据模型参数量训练 Token 数量GPU 理论算力实际 GPU 利用率MFU快速估算整个训练过程所需的计算量和时间。2. 显存Memory训练模型占用的显存并不只有模型参数还包括Model ParametersGradientsOptimizer StatesActivations例如使用FP32 AdamW训练时每个参数粗略需要Parameter 4 Bytes Gradient 4 Bytes Adam First Moment 4 Bytes Adam Second Moment 4 Bytes -------------------------- Total ≈ 16 Bytes / Parameter而这甚至还没有包含 Activation。因此模型“放得下”不代表模型“训得动”。核心理解训练前先做 Napkin Math算力决定训练多久显存决定能不能训练。3.2 Tensor深度学习的基本数据结构有了资源核算意识下一个问题是GPU 到底在存储和计算什么答案是Tensor张量。PyTorch 中几乎所有核心数据最终都是 TensorData Parameters Gradients Optimizer States Activations ↓ Tensor理解 Tensor 时需要重点关注shaperankdtypedevicestride例如 Transformer 中经常出现[Batch, Sequence, Heads, Hidden]Tensor 之间会进行viewreshapetransposeelement-wise operationMatrix Multiplication其中最核心的计算是Matrix MultiplicationMatMul此外jaxtyping和einops可以帮助我们更加明确地表达 Tensor 的维度变化例如batch × sequence × heads × hidden核心理解Tensor 是深度学习世界的数据载体MatMul 是最核心的计算。3.3 MemoryTensor 占多少内存既然模型中的参数、梯度和 Activation 都是 Tensor那么自然会产生下一个问题一个 Tensor 到底占多少显存基本公式Tensor Memory ≈ Number of Elements × Bytes per Element因此dtype会直接影响显存占用。Data Type每元素大小特点FP324 Bytes精度高、稳定但显存开销大FP162 Bytes快、省显存但动态范围较小BF162 Bytes动态范围接近 FP32LLM 训练常用FP81 Byte更省显存但数值稳定性要求更高为什么 LLM 常使用 BF16BF16 与 FP16 都只需要 2 Bytes但 BF16 保留了更大的 exponent 范围。因此FP16 精度较高 但数值范围较小 BF16 精度略低 但数值范围更大对于容易出现极大或极小数值的深度神经网络BF16 通常具有更好的训练稳定性。Tensor 的底层存储一个 Tensor 并不只是数据本身还可以理解为Tensor ├── Data Pointer ├── Shape ├── Stride ├── Dtype └── Device很多 Tensor 操作例如view()transpose()slice()可能只是改变“如何解释底层数据”而不会真正复制数据。但contiguous()可能产生新的内存复制。此外CPU 与 GPU 之间的数据传输本身也存在成本CPU RAM │ │ PCIe ↓ GPU HBM因此优化 LLM 不仅需要减少计算也需要尽量减少不必要的数据搬运。核心理解显存优化的本质是管理 Tensor 的数量、dtype、布局以及数据移动。3.4 Compute这些 Tensor 要计算多少解决 Memory 问题之后下一个问题是这些 Tensor 到底需要多少计算这里最核心的指标就是FLOPs。对于矩阵乘法[M, K] × [K, N]计算量大约为2 × M × K × N FLOPs因此模型规模、Batch Size、Sequence Length 等都会直接影响训练计算量。FLOPs 与实际 GPU 性能理论 FLOPs 并不代表 GPU 实际能够达到的性能。因此需要引入MFUModel FLOPs UtilizationMFU 实际 FLOP/s ──────────── GPU 理论 FLOP/sMFU 衡量的是GPU 的理论计算能力到底真正利用了多少。例如 GPU 理论峰值为1000 TFLOPS而模型实际只跑到500 TFLOPS那么MFU 50%自动求导与反向传播PyTorch 会自动构建计算图Input ↓ Forward ↓ Loss ↓ Backward ↓ Gradient通过loss.backward()PyTorch 根据链式法则计算每个参数的梯度。训练过程因此包含Forward Compute Backward Compute通常反向传播所需计算量还会高于单纯的 Forward。核心理解FLOPs 告诉我们理论工作量MFU 告诉我们 GPU 实际完成了多少工作。3.5 Model Training把所有组件组装起来前面学习的是训练系统的各个“零件”Tensor Memory Compute Gradient这一节开始把它们组合成完整的训练流程。典型的 PyTorch Training LoopInitialize Parameters ↓ Load Data ↓ Forward ↓ Loss ↓ Backward ↓ Gradient ↓ Optimizer ↓ Update Parameters ↓ Next Batch模型中的可学习参数通过nn.Parameter进行管理。参数初始化也非常重要例如 Xavier Initialization其目标之一是避免深层网络中出现Activation Explosion or Activation Vanishing数据加载当训练数据大于 CPU 内存时可以使用memmap避免一次性将整个数据集加载进 RAM。GPU 训练时还希望形成CPU 准备 Batch N1 │ │ asynchronous transfer ↓ GPU 计算 Batch N通过Pinned MemoryAsynchronous Transfer让 CPU 数据准备与 GPU 计算重叠从而减少 GPU 等待时间。Mixed Precision训练时还可以混合使用FP32 BF16 FP16在尽量保持数值稳定性的同时降低显存占用提高 Tensor Core 利用率提升训练速度核心理解3.5 是整个章节的组装环节Tensor Memory Compute Data Gradient Optimizer → Training。3.6 Arithmetic Intensity 与 Roofline最后需要回答一个非常重要的问题为什么 GPU 理论上有 1000 TFLOPS我的程序却远远跑不到 1000 TFLOPS原因是 GPU 不仅需要计算还需要不断搬运数据HBM ↓ Load Data ↓ Compute Core ↓ Compute ↓ Store Data ↓ HBM所以程序性能同时受到两个因素影响Compute Capability Memory BandwidthArithmetic Intensity定义Arithmetic Intensity FLOPs ────── Bytes它表示每从 Memory 搬运 1 Byte 数据可以完成多少次浮点计算。Memory-Bound如果搬很多数据 ↓ 只做很少计算Arithmetic Intensity 就很低。此时性能主要受Memory Bandwidth限制。例如ReLU Element-wise Operations Vector Operations通常更容易 Memory-Bound。Compute-Bound如果搬一次数据 ↓ 进行大量计算Arithmetic Intensity 较高。例如Large Matrix Multiplication这时性能主要受 GPU 的计算能力限制Compute Capability即 Compute-Bound。Roofline ModelRoofline Model 将两种瓶颈统一起来转折点可以理解为Arithmetic Intensity ≈ Peak FLOPs ────────────── Memory Bandwidth低于这个点Memory-Bound高于这个点Compute-Bound为什么训练和推理表现不同这个模型还能解释 LLM 中一个非常重要的现象。Training训练中大量执行Matrix × MatrixArithmetic Intensity 较高因此更容易Compute-BoundGPU 的 Tensor Core 可以得到较充分利用。Autoregressive Inference逐 Token 推理时经常更接近Matrix × Vector每生成一个 Token 都需要读取大量模型参数但计算量相对有限。因此更容易Memory-Bound这也是为什么LLM 推理性能不仅取决于 GPU FLOPs也高度依赖显存带宽。3.13.6 的整体逻辑-问题形式整个第三章可以看成六个连续的问题3.1 Resource Accounting 为什么要做资源核算 │ ↓ LLM 很大、训练很贵 │ ├───────────────┐ ↓ ↓ Memory Compute 能不能放下 要计算多久 │ │ ↓ │ 3.2 Tensor │ 训练中的基本对象 │ │ │ ↓ │ Parameters / Gradients │ Activations / Data │ │ │ ↓ ↓ 3.3 Memory 3.4 Compute Tensor占多少显存 需要多少FLOPs │ │ └───────┬───────┘ ↓ 3.5 Training │ Data → Model → Loss ↓ Backward ↓ Optimizer ↓ Update Weight │ ↓ 为什么GPU还是不够快 │ ↓ 3.6 Roofline │ Arithmetic Intensity FLOPs / Bytes │ ┌───────┴────────┐ ↓ ↓ Memory-Bound Compute-Bound │ │ └───────┬────────┘ ↓ GPU Performance第三章总示意图LLM TRAINING │ ▼ ┌─────────────────────────┐ │ 3.1 Resource Accounting │ │ 训练需要多少资源 │ └────────────┬────────────┘ │ ┌───────────┴───────────┐ │ │ ▼ ▼ MEMORY COMPUTE 能不能训 要训多久 │ │ ▼ │ ┌─────────────┐ │ │ 3.2 Tensor │ │ └──────┬──────┘ │ │ │ ┌──────────┼──────────┐ │ │ │ │ │ ▼ ▼ ▼ │ Parameters Gradients Activations │ │ │ │ │ └──────────┼──────────┘ │ ▼ ▼ ┌────────────────┐ ┌────────────────┐ │ 3.3 Memory │ │ 3.4 Compute │ │ │ │ │ │ Shape × Dtype │ │ MatMul → FLOPs │ │ FP32/BF16/FP8 │ │ MFU │ └────────┬───────┘ └────────┬───────┘ │ │ └───────────┬───────────┘ ▼ ┌──────────────────┐ │ 3.5 Training │ │ │ │ Data → Forward │ │ ↓ │ │ Loss │ │ ↓ │ │ Backward │ │ ↓ │ │ Optimizer │ │ ↓ │ │ Update Weight │ └────────┬─────────┘ │ ▼ 实际 GPU 为什么不够快 │ ▼ ┌──────────────────┐ │ 3.6 Roofline │ │ │ │ FLOPs / Bytes │ │ ↓ │ │ Arithmetic │ │ Intensity │ └────────┬─────────┘ │ ┌───────────┴───────────┐ │ │ ▼ ▼ MEMORY-BOUND COMPUTE-BOUND 显存带宽限制 GPU算力限制 │ │ └───────────┬───────────┘ ▼ GPU UTILIZATION本章核心总结第三章真正想建立的不是单纯的PyTorch API 使用能力而是一套LLM Systems Thinking模型训练的本质是 Tensor 在 Memory 中存储和移动同时 GPU 对 Tensor 执行大量计算。因此整章可以压缩成一条逻辑链Resource ↓ Tensor ↓ Memory Compute ↓ Training ↓ Performance ↓ Roofline进一步来看后续很多 LLM 系统优化技术其实都可以放进这个框架Mixed Precision ↓ 减少 Memory 提升 Compute Gradient Checkpointing ↓ 用 Compute 换 Memory FlashAttention ↓ 减少 HBM 数据搬运 KV Cache ↓ 用 Memory 换重复 Compute Distributed Training ↓ 多 GPU 分摊 Memory Compute所以理解这一章之后再学习Transformer、FlashAttention、KV Cache、Mixed Precision、Gradient Checkpointing 和 Distributed Training可以始终追问两个核心问题1. 这个技术改变了多少 Memory2. 这个技术改变了多少 Compute / Data Movement这两个问题基本贯穿整个 LLM Systems 优化体系。