讲义原文一、课前概览上一讲内容为概述与分词Tokenization本讲核心是讲解训练模型所需的所有基础组件将从张量tensors自底向上讲解至模型、优化器再到训练循环全程重点关注效率资源使用主要核算两类资源内存GB和计算量FLOPs。启发性问题粗略估算先通过两个实际问题建立对资源消耗的直观认知问题1使用1024张H100显卡训练700亿参数的模型处理15万亿个token需要多久total_flops6*70e9*15e12# 总计算量asserth100_flop_per_sec1979e12/2# H100每秒计算量非稀疏模式mfu0.5# 模型FLOPs利用率flops_per_dayh100_flop_per_sec*mfu*1024*60*60*24# 单日总计算量daystotal_flops/flops_per_day# 计算所需天数问题2使用8张H100显卡基于AdamW优化器朴素实现能训练的最大模型有多大h100_bytes80e9# 单张H100显存字节# 每个参数占用字节参数(4) 梯度(4) 优化器状态(44)bytes_per_parameter44(44)num_parameters(h100_bytes*8)/bytes_per_parameter# 最大参数数量注意1默认对参数和梯度使用float32精度也可对参数/梯度用bf1622字节同时保留float32参数副本4字节虽不省内存但运算更快。注意2未核算激活值的内存占用取决于批次大小和序列长度。说明本讲不会讲解Transformer模型已有优质资料Assignment 1 handout、Mathematical description、Illustrated Transformer、Illustrated GPT-2而是用更简单的模型讲解核心知识点操作机制仅使用PyTorch简单直接核心思维资源核算务必掌握核心直觉整体框架不涉及大模型细节二、内存核算2.1 tensor基础tensor是存储所有数据的基本单元参数、梯度、优化器状态、数据、激活值均以tensor存储PyTorch tensor文档。创建张量的多种方式# 直接创建xtorch.tensor([[1.,2,3],[4,5,6]])# 全0张量xtorch.zeros(4,8)# 全1张量xtorch.ones(4,8)# 标准正态分布张量xtorch.randn(4,8)# 分配内存但不初始化值后续自定义赋值xtorch.empty(4,8)nn.init.trunc_normal_(x,mean0,std1,a-2,b2)# 截断正态分布初始化2.2 张量内存几乎所有数据参数、梯度、激活值、优化器状态都以浮点数存储不同浮点类型的内存占用和特性差异显著1float32单精度是PyTorch默认类型占4字节/元素维基百科。科学计算中是基准深度学习中可灵活选用更低精度。内存计算示例xtorch.zeros(4,8)assertx.dtypetorch.float32# 验证默认类型assertx.numel()4*8# 元素总数assertx.element_size()4# 单元素字节数assertget_memory_usage(x)4*8*4# 总内存128字节# GPT-3前馈层单个矩阵的内存占用assertget_memory_usage(torch.empty(12288*4,12288))2304*1024*1024# 2.3 GB2float16半精度占2字节/元素内存减半维基百科。动态范围差小数值易下溢可能导致训练不稳定xtorch.zeros(4,8,dtypetorch.float16)assertx.element_size()2xtorch.tensor([1e-8],dtypetorch.float16)assertx0# 下溢为03bfloat16脑浮点Google Brain 2018年推出占2字节/元素动态范围与float32一致维基百科。分辨率降低但对深度学习影响小无下溢问题xtorch.tensor([1e-8],dtypetorch.bfloat16)assertx!0# 无下溢4fp82022年标准化专为机器学习设计H100支持E4M3范围[-448, 448]和E5M2[-57344, 57344]两种变体参考。训练启示float32训练稳定但内存占用大fp8/float16/bfloat16训练有不稳定风险需用混合精度训练解决。2.3 张量与GPU默认张量存储在CPU内存需移至GPU以利用并行计算能力# 验证CPU默认存储xtorch.zeros(32,32)assertx.devicetorch.device(cpu)# 检查GPU可用性iftorch.cuda.is_available():num_gpustorch.cuda.device_count()foriinrange(num_gpus):propertiestorch.cuda.get_device_properties(i)# 查看GPU属性memory_allocatedtorch.cuda.memory_allocated()# 已分配显存# 移至GPU0号设备yx.to(cuda:0)asserty.devicetorch.device(cuda,0)# 直接在GPU创建张量ztorch.zeros(32,32,devicecuda:0)# 计算显存占用new_memory_allocatedtorch.cuda.memory_allocated()memory_usednew_memory_allocated-memory_allocatedassertmemory_used2*(32*32*4)# 2个32x32 float32张量1024*4*28192字节三、计算量核算3.1 张量操作多数张量由其他张量运算生成不同操作的内存和计算开销不同1张量存储原理PyTorch张量是指向内存的指针附带元数据描述元素访问方式xtorch.tensor([[0.,1,2,3],[4,5,6,7],[8,9,10,11],[12,13,14,15],])# 维度0行的步长跳过4个元素assertx.stride(0)4# 维度1列的步长跳过1个元素assertx.stride(1)1# 计算元素索引行r1列c2r,c1,2indexr*x.stride(0)c*x.stride(1)assertindex6# 对应元素62张量切片视图多数切片操作返回张量的视图无拷贝修改原张量会同步修改视图xtorch.tensor([[1.,2,3],[4,5,6]])# 取第0行视图yx[0]asserttorch.equal(y,torch.tensor([1.,2,3]))assertsame_storage(x,y)# 共享内存# 取第1列视图yx[:,1]asserttorch.equal(y,torch.tensor([2,5]))assertsame_storage(x,y)# 维度重塑视图yx.view(3,2)asserttorch.equal(y,torch.tensor([[1,2],[3,4],[5,6]]))assertsame_storage(x,y)# 转置视图yx.transpose(1,0)asserttorch.equal(y,torch.tensor([[1,4],[2,5],[3,6]]))assertsame_storage(x,y)# 修改原张量视图同步变化x[0][0]100asserty[0][0]100# 非连续视图无法直接重塑需先转为连续张量会拷贝xtorch.tensor([[1.,2,3],[4,5,6]])yx.transpose(1,0)assertnoty.is_contiguous()try:y.view(2,3)assertFalseexceptRuntimeErrorase:assertview size is not compatible with input tensors size and strideinstr(e)# 转为连续张量后重塑内存拷贝yx.transpose(1,0).contiguous().view(2,3)assertnotsame_storage(x,y)关键视图操作无开销拷贝操作占用额外内存和计算资源。3按元素操作对每个元素执行操作返回同形状新张量xtorch.tensor([1,4,9])# 幂/平方根/倒数平方根asserttorch.equal(x.pow(2),torch.tensor([1,16,81]))asserttorch.equal(x.sqrt(),torch.tensor([1,2,3]))asserttorch.equal(x.rsqrt(),torch.tensor([1,1/2,1/3]))# 算术运算asserttorch.equal(xx,torch.tensor([2,8,18]))asserttorch.equal(x*2,torch.tensor([2,8,18]))asserttorch.equal(x/0.5,torch.tensor([2,8,18]))# 上三角矩阵用于因果注意力掩码xtorch.ones(3,3).triu()asserttorch.equal(x,torch.tensor([[1,1,1],[0,1,1],[0,0,1]],))4矩阵乘法深度学习核心# 基础矩阵乘法xtorch.ones(16,32)wtorch.ones(32,2)yx wasserty.size()torch.Size([16,2])# 批量序列维度的矩阵乘法xtorch.ones(4,8,16,32)# batch, sequence, ..., hiddenwtorch.ones(32,2)yx wasserty.size()torch.Size([4,8,16,2])# 前2维迭代与w相乘3.2 Einops命名维度的张量操作Einops基于爱因斯坦求和符号设计通过命名维度简化张量操作解决传统PyTorch维度易混淆问题。1传统方式的痛点xtorch.ones(2,2,3)# batch, sequence, hiddenytorch.ones(2,2,3)# batch, sequence, hiddenzx y.transpose(-2,-1)# 维度索引(-2/-1)易混淆2Jaxtyping维度标注仅文档无强制校验x:Float[torch.Tensor,batch seq heads hidden]torch.ones(2,2,1,3)3Einops Einsum通用矩阵乘法x:Float[torch.Tensor,batch seq1 hidden]torch.ones(2,3,4)y:Float[torch.Tensor,batch seq2 hidden]torch.ones(2,3,4)# 传统方式zx y.transpose(-2,-1)# Einops方式维度显式命名zeinsum(x,y,batch seq1 hidden, batch seq2 hidden - batch seq1 seq2)# 省略通用维度...表示任意数量的前置维度zeinsum(x,y,... seq1 hidden, ... seq2 hidden - ... seq1 seq2)4Einops Reduce张量归约x:Float[torch.Tensor,batch seq hidden]torch.ones(2,3,4)# 传统方式最后一维求均值yx.mean(dim-1)# Einops方式归约hidden维度求和yreduce(x,... hidden - ...,sum)5Einops Rearrange维度重组x:Float[torch.Tensor,batch seq total_hidden]torch.ones(2,3,8)# total_hidden heads * hidden1w:Float[torch.Tensor,hidden1 hidden2]torch.ones(4,4)# 拆分total_hidden为heads和hidden1xrearrange(x,... (heads hidden1) - ... heads hidden1,heads2)# 矩阵乘法xeinsum(x,w,... hidden1, hidden1 hidden2 - ... hidden2)# 合并heads和hidden2xrearrange(x,... heads hidden2 - ... (heads hidden2))3.3 张量操作的FLOPs1FLOPs定义FLOPs浮点运算次数计算量指标FLOP/s每秒浮点运算次数硬件速度指标也写作FLOPS。2直观认知GPT-32020训练需3.14e23 FLOPsGPT-42023训练推测需2e25 FLOPsA100峰值性能312 teraFLOP/sH100峰值性能1979 teraFLOP/s稀疏模式非稀疏模式减半。3线性模型的FLOPs计算# 根据GPU可用性设置参数iftorch.cuda.is_available():B16384# 样本数D32768# 特征维度K8192# 输出维度else:B1024D256K64deviceget_device()xtorch.ones(B,D,devicedevice)wtorch.randn(D,K,devicedevice)yx w# 每个(i,j,k)三元组1次乘法 1次加法 → 总FLOPs 2*B*D*Kactual_num_flops2*B*D*K4其他操作的FLOPs按元素操作m×n矩阵O(mn) FLOPs矩阵加法m×nmn FLOPs核心结论大矩阵场景下矩阵乘法是深度学习中计算量最大的操作。5FLOPs与实际耗时# 计时矩阵乘法actual_timetime_matmul(x,w)# 实际FLOP/sactual_flop_per_secactual_num_flops/actual_time# 硬件标称FLOP/s依赖设备和数据类型promised_flop_per_secget_promised_flop_per_sec(device,x.dtype)# 模型FLOPs利用率MFU实际FLOP/s / 标称FLOP/smfuactual_flop_per_sec/promised_flop_per_secMFU ≥ 0.5即为优秀矩阵乘法占比越高MFU越高。四、模型、训练循环与最佳实践注原代码中该部分仅定义函数未实现核心前置知识已覆盖模型基于nn.Module定义管理参数训练循环包含数据加载、优化器如AdamW、梯度更新、断点续训、混合精度训练等关键注意点随机种子固定保证可复现、激活值内存核算、MFU优化等。核心总结张量是深度学习的基础需熟练掌握其创建、存储、操作和设备迁移内存核算需关注浮点类型float32/bf16/fp8的取舍计算量核算核心是矩阵乘法的FLOPsEinops简化张量维度操作提升代码可读性MFU是衡量硬件利用率的核心指标优化MFU是提升训练效率的关键。