斯坦福大学 CS336 Lecture 02 Pytorch and Resource Accounting

1.Overview2.Resource Accounting2.1 Memory2.1.1 Tensor Concept张量是常见的数据结构的拓展例如标量零维张量向量一维张量矩阵二维张量。在机器学习中之所以使用张量是因为1.硬件加速。硬件对张量有专门的设计适配GPU进行快速计算谷歌还专门设计了TPU2.自动求导。设置了requires_grad True 时这个张量会记录自己是怎么被算出来的即记录父节点。在后续训练中只需调用.backward()梯度就能自动沿着张量传播回来3.广播机制Broadcasting。将一个标量0维张量与一个矩阵2维张量相加系统会自动把标量“扩展”成同样大小的矩阵再相加。2.1.2 Types of floatsfloat32fp32单精度一个数字32位4bytes较大自己研究了一下一个小数的存储方法有点意思。例如19303.0890是怎么存成一个float32的这下懂了为啥会出现 0.1 0.2 0.300000001这种类似的问题了float16fp16半精度动态范围较小只有5位很容易出现下溢到0或上溢到无穷大的现象bfloat16bf16fp8by Nvidia in 2022混合精度训练mix-precision-training在训练的pipeline中例如在前向传播、反向传播、优化器或是梯度累积阶段不同阶段用到的精度可能不一样。2.2 Computeimport torch memory_allocated torch.cuda.memory_allocated() z torch.zeros(32, 32, devicecuda:0) new_memory_allocated torch.cuda.memory_allocated() memory_used new_memory_allocated - memory_allocated print(memory_used)结果输出4096 32 * 32 * 4 一个float32占4bytes2.2.1 Tensor Storage一个大小为 4 x 4的矩阵在内存中是一个一维长数组。张量在 PyTorch 中是指向已分配内存的指针。张量对象本身拥有元数据指定了如何在数组中找到特定的地址。在这个二维矩阵中元数据具有两个维度stride[0] 和 stride[1]分别代表从各自的维度到下一级要跨越多少个单位即stride[0] 代表每过一行要跨越4个单位stride[0] 4; 而 stride[1] 代表每过一列要跨越1个单位stride[1] 12.2.2 Tensor Operationsa. Tensor Slicing张量的很多操作实际上并不会创建一个新的张量它们只会创建一个不同的视图而不会进行复制import torch def print_storage(x: torch.Tensor): print(x.untyped_storage().data_ptr()) x torch.tensor([[1., 2, 3], [4, 5, 6]]) y1 x[0] y2 x[:, 1] y3 x.view(3, 2) y4 x.transpose(1, 0)但 print_storage() 输出的结果一致即占用同一块内存空间。x[0, 0] 100 print(y4[0, 0]) # tensor(100.) print(y4[0, 0].item()) # 100.0所以当修改 x 时其他 Tensor 中对应元素也会修改。强制复制张量的操作1.显式调用 clone()2.当操作要求连续内存而视图不连续时调用.contiguous()某些操作例如转置会破坏原本张量的连续性contiguous。遍历具有连续性的张量意味着能够平滑地遍历存储中的一维数组反之则是跳跃访问的。对一个转置后的视图调用 .view() 改变形状会报错。.contiguous()的作用是把逻辑顺序散乱的数据重新连续排列。这个时候会发生一次复制对 x 和 y5 进行 print_storage()得到的结果不一样。b. Tensor Elementwise下列操作会得到新的张量很好理解——计算结果总是要保存的x torch.tensor([1, 4, 9]) assert torch.equal(x.pow(2), torch.tensor([1, 16, 81])) assert torch.equal(x.sqrt(), torch.tensor([1, 2, 3])) assert torch.equal(x.rsqrt(), torch.tensor([1, 1 / 2, 1 / 3])) # i - 1/sqrt(x_i) assert torch.equal(x x, torch.tensor([2, 8, 18])) assert torch.equal(x * 2, torch.tensor([2, 8, 18])) assert torch.equal(x / 0.5, torch.tensor([2, 8, 18])) y torch.ones(3, 3).triu() # triu() # 创建上三角矩阵 assert torch.equal(y, torch.tensor([ [1, 1, 1], [0, 1, 1], [0, 0, 1]], ))c. Tensor Matmul对于一个 16 x 32 和一个 32 x 2的矩阵得到的结果是一个 16 x 2 的结果。但在深度学习中常见的张量形式是一个四维矩阵w [a b cd]其中 a 为 batchb 为 sequence。c 和 d 为数据。而对 w 的乘法是将每个 batch 、每个 sequence 都进行一次矩阵的乘法即x torch.ones(4, 8, 16, 32) w torch.ones(32, 2) y x w # 只会对张量的“最后两个维度”执行矩阵乘法 # 前面的所有维度都当作“批量维度Batch”原样保留。 assert y.size() torch.Size([4, 8, 16, 2])2.2.3 Tensor Einops为了张量代码的可读性在操作张量时我们命名张量的所有维度而不像以前那样本质上只依赖于索引a.jaxtyping basics# old way : 通过注释等记录 x torch.ones(2, 2, 1, 3) # batch seq heads hidden # eniops way x: Float[torch.Tensor, batch seq heads hidden] torch.ones(2, 2, 1, 3)b. einops 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) z x y.transpose(-2, -1) # 将倒数第一维度和倒数第二维度转置 z einsum(x, y, batch seq1 hidden, batch seq2 hidden - batch seq1 seq2) # 两种方式等价 z einsum(x, y, ... seq1 hidden, ... seq2 hidden - ... seq1 seq2) # 用 ... 来表示对任意数量的维度进行广播c. einops reduce对张量的某个维度做“归约聚合”操作比如求和、求平均、取最大值等import torch from einops import reduce x torch.ones(2, 3, 4) # 对 hidden 维度求和消掉它 y reduce(x, batch seq hidden - batch seq, sum) print(y.shape) # torch.Size([2, 3]) print(y)d. einops rearrangex: Float[torch.Tensor, batch seq total_hidden] torch.ones(2, 3, 8) x rearrange(x, ... (heads hidden1) - ... heads hidden1, heads2) # 得到的结果维度为2,3,2,4 from 2,3,82.2.4 Tensor Operation Flops张量操作本质上是浮点操作floating-point operationFLOP而浮点操作是指任何涉及浮点数的运算比如加法或乘法区分 FLOPs 以及 FLOPS。前者小写 s 说明是复数形式代表浮点操作次数衡量的是完成的计算总量后者大写 S 说明是 per second代表每秒浮点运算次数衡量的是硬件的速度 。Linear Model假设一个矩阵 M 大小为BD乘以一个矩阵 D, K将 D 维向量映射到一个 K 维向量则M torch.ones(B, D, devicedevice) # 要得到的矩阵N 形状为 B x K 则线性操作矩阵T形状为 D x K T torch.randn(D, K, devicedevice) N M T要进行的FLOPs 2 * B * D * K对于 N 中每个元素都要计算 D 次乘法以及 D - 1 次加法共 2D - 1 次操作而 N 中有 B * K个元素故 FLOPs 2 * B * D * K 其他操作的浮点运算次数通常与矩阵或张量的大小呈线性关系。在深度学习中对于足够大的矩阵矩阵乘法是所有操作中最耗时的推广到 Transform——B the number of data pointsD * K : the number of parametersFLOPs for forward pass is (2 * tokens * parameters) (忽略了序列长度和其他因素)得到了浮点运算的次数计算实际运行的时间def time_matmul(a: torch.Tensor, b: torch.Tensor) - float: Return the number of seconds required to perform a b. # Wait until previous CUDA threads are done if torch.cuda.is_available(): torch.cuda.synchronize() # PyTorch 在 GPU 上运行张量运算时是异步的。 # Python 会立刻把指令 a b 发给 GPU然后不等 GPU 算完Python 就继续执行下一行代码了。 # torch.cuda.synchronize() 的作用就是强制 Python 停下来等待, 直到 GPU 上的所有线程全部执行完毕. # 保证时间计算准确 def run(): # Perform the operation a b # Wait until CUDA threads are done if torch.cuda.is_available(): torch.cuda.synchronize() # Time the operation num_trials times num_trials 5 total_time timeit.timeit(run, numbernum_trials) return total_time / num_trials # B 16384 D 32768 K 8192 actual_num_flops 2 * B * D * K actual_time time_matmul(M, T) # inspect actual_time actual_flop_per_sec actual_num_flops / actual_time print(actual_num_flops) # 8796093022208 8.80e12 print(actual_time) # 1.2285193999996409 print(actual_flop_per_sec) # 7159913813498.241 7.16e12模型浮点运算利用率 (Model Flops UtilizationMFU)即实际浮点次数除以承诺的浮点运算次数mfu actual_flop_per_sec / promised_flop_per_sec查了一下3070的promised_flop_per_sec 20.31TFLOPs但没详细写明是什么float数据格式默认利用的float类型为float32 将其换成为bfloat16得到的结果M M.to(torch.bfloat16) T T.to(torch.bfloat16) # actual_num_flops 8.80e12 不变 # actual_time 0.42 # actual_flop_per_sec 21.06 TFLOPs每秒浮点运算次数FLOP/s取决于硬件以及数据类型。5090 4090bfloat16 float322.2.5 Gradients FLOPs整个计算过程中除了计算矩阵乘法之外还有一部分用于计算梯度Example:M torch.ones(B, D, devicedevice) T1 torch.randn(D, D, devicedevice, requires_gradTrue) T2 torch.randn(D, K, devicedevice, requires_gradTrue) # 一个两层线性模型 # M --T1-- N1 --T2-- N2 - loss N1 M T1 N2 N1 T2 loss N2.pow(2).mean()前向传播计算量2 * B * D * D 2 * B * D * K计算梯度过程forward pass : 2 x ( data points) x ( parameters) FLOPSbackward pass: 4 x ( data points) x ( parameters) FLOPStotal: 6 x ( data points) x ( parameters) FLOPS3.Model3.1 Module Parametersimport torch from torch import nn input_dim 16384 output_dim 32 w nn.Parameter(torch.randn(input_dim, output_dim)) assert isinstance(w, torch.Tensor) # 检查 w 是不是 torch.Tensor 的一个实例 # Parameter 在行为上完全是一个张量可以对它做 w x、w.view()、w.transpose() 等所有张量操作 assert type(w.data) torch.Tensor # w.data 的类型是精确的 torch.Tensor而非子类不带 Parameter 的“身份标记” # 在训练过程中应通过 w.data 来修改底层数据比如做权重裁剪、初始化而不是重新赋值 w 本身。 # 否则w 会丢失 Parameter 的身份优化器将不再更新它 # nn.Parameter 相较于普通张量 torch.Tensor # 前者默认 requires_grad True # 会被自动添加到模型的 .parameters() 迭代器parameter initialization发现问题x的形状为intput_dim却能和形状为input_dimoutput_dim的w相乘得到形状为output_dim的output是PyTorch默认把一维张量当作行向量来参与矩阵乘法得到的结果较大这是因为数值增长的量级本质上与隐藏维度的平方根成正比输入数据x的标准差是 1经过一个全连接层y x w后输出y的方差会变成input_dim × 1所以y的标准差会变成sqrt(input_dim)输出结果与标准差量级相同发生数值发散的爆炸现象导致模型训练不稳定解决方法w nn.Parameter(torch.randn(input_dim, output_dim) / np.sqrt(input_dim)) # 将 w 除以输入维度的平方根 # Xavier initialization得到的结果集中在类似 Normal(0, 1) 的正态分布附近。但由于正态分布曲线永远不会碰到横轴理论上会出现及其庞大的数据特别是在大模型参数动辄几十亿数量级的情况下。故更安全的方法将尾部截断例如将范围限制在 [-3, 3]3.2 Custom Modelclass Cruncher(nn.Module): def __init__(self, dim: int, num_layers: int): super().__init__() self.layers nn.ModuleList([ Linear(dim, dim) for i in range(num_layers) # Linear(dim, dim)把输入张量的最后一个维度从 dim 映射到 dim也就是做一次线性变换 ]) # nn.ModuleList 是 PyTorch 里一个专门用来存放多个子模块如 Linear、Conv2d的容器 # 当把层放进 nn.ModuleList 时PyTorch 会自动“记住”这些层让优化器能找到它们的参数 self.final Linear(dim, 1) def forward(self, x: torch.Tensor) - torch.Tensor: # Apply linear layers B, D x.size() for layer in self.layers: x layer(x) # Apply final head x self.final(x) assert x.size() torch.Size([B, 1]) # Remove the last dimension x x.squeeze(-1) assert x.size() torch.Size([B]) return x D 64 # Dimension num_layers 2 model Cruncher(dimD, num_layersnum_layers) device get_device() model model.to(device) B 8 # Batch size x torch.randn(B, D, devicedevice) y model(x) assert y.size() torch.Size([B])当 num_layers 2 时这个模型一共有 3 层两层 Linear 和一个 final 层。前两层的参数量都为 D * D最后一层的参数量为 D * 1。3.3 Note about Randomness# Torch seed 0 torch.manual_seed(seed) # NumPy import numpy as np np.random.seed(seed) # Python import random random.seed(seed)将 Torch、NumPy 和 Python 引入随机性的地方全部“锁死”让代码每次运行时产生的随机数序列完全一样增加代码的可复现性。3.4 Data Loadingdata np.memmap(data.npy, dtypenp.int32) # data.npy这是一个已存在硬盘上的文件 # np.memmap 的返回值data 是一个内存映射对象 # 它看起来、用起来和普通 NumPy 数组几乎一模一样支持索引、切片、算术运算 # 但数据本身还留在硬盘上没有加载到内存里。 assert np.array_equal(data, orig_data)使用 np.memmap() 适合处理超大规模数据集内存无法直接 load 整个数据集时。当尝试访问数据时会按需加载文件原理 虚拟内存 缺页异常 手动管理内存3.5 Optimizarclass AdaGrad(torch.optim.Optimizer): # 父类 def __init__(self, params: Iterable[nn.Parameter], lr: float 0.01): super(AdaGrad, self).__init__(params, dict(lrlr)) # 调用 PyTorch 优化器的基类构造函数把参数和学习率注册进去这样 self.param_groups 里就有了所有参数的信息 def step(self): for group in self.param_groups: # self.param_groups 在父类定义 lr group[lr] for p in group[params]: # Optimizer state state self.state[p] # 一个字典用来记录优化器的“状态”记录每个参数的历史信息 # 在 AdaGrad 里它记录的是历史梯度平方的累加和 grad p.grad.data # 在 loss.backward() 反向传播过程中得到 # Get squared gradients g2 sum_{it} g_i^2 g2 state.get(g2, torch.zeros_like(grad)) # g2 为历史梯度平方的累加和 # Update optimizer state g2 torch.square(grad) # 把当前梯度的平方累加进去 state[g2] g2 # g2 的值只会增加不会减少 # Update parameters p.data - lr * grad / torch.sqrt(g2 1e-5) # 如果某个参数的历史梯度平方和 g2 很大说明这个参数一直在剧烈更新那么 grad / sqrt(g2) 就会变小相当于学习率被自动缩小。 optimizer AdaGrad(model.parameters(), lr0.01)x torch.randn(B, D, deviceget_device()) y torch.tensor([4., 5.], deviceget_device()) pred_y model(x) loss F.mse_loss(inputpred_y, targety) loss.backward() optimizer.step() # 在这里优化器真正发挥作用 state model.state_dict() # model.state_dict() 返回一个 Python 字典里面包含模型所有参数张量的副本 # 注意和优化器里的 state 区分 optimizer.zero_grad(set_to_noneTrue) # 释放内存优化器状态的内存需求参数数量 num_parameters D * D * num_layers D激活值数量 num_activations前向传播过程中每一层计算出的中间输出张量对于每一层、每一个数据点、每一个维度都必须存储激活值。即结果为 B * D * num_layers梯度数量 num_gradients 梯度数量与参数数量相同反向传播为每一个参数计算出一个梯度优化器状态数量 num_optimizer_states 对于AdaGrad存储的状态为每个参数对应一个g2即数量同样与参数数量相同故总内存量假设数据格式为Float32即4个字节total_memory 4 * ( num_parameters num_activations num_gradients num_optimizer_states )FLOPs 6 * B * num_parameters3.6 Train Loopimport torch from torch import nn from torch.nn import Linear from torch.optim import SGD import torch.nn.functional as F class Cruncher(nn.Module): def __init__(self, dim: int, num_layers: int): super().__init__() self.layers nn.ModuleList([ Linear(dim, dim) for i in range(num_layers) ]) self.final Linear(dim, 1) def forward(self, x: torch.Tensor) - torch.Tensor: # Apply linear layers B, D x.size() for layer in self.layers: x layer(x) # Apply final head x self.final(x) assert x.size() torch.Size([B, 1]) # Remove the last dimension x x.squeeze(-1) assert x.size() torch.Size([B]) return x def train(name: str, get_batch, D: int, num_layers: int, B: int, num_train_steps: int, lr: float): device cuda:0 model Cruncher(dimD, num_layers0).to(device) # num_layers 0 说明只有最后的 final 层输出形状 (B, 1) optimizer SGD(model.parameters(), lr0.01) for t in range(num_train_steps): # Get data x, y get_batch(BB) # Forward (compute loss) pred_y model(x) loss F.mse_loss(pred_y, y) # Backward (compute gradients) loss.backward() # Update parameters optimizer.step() optimizer.zero_grad(set_to_noneTrue) if t % 2 0: # 每2步打印一次 print(fStep {t}: loss {loss.item():.6f}) print(Final model weight:) print(model.final.weight) print(True weight:) print(true_w) D 16 true_w torch.arange(D, dtypetorch.float32, devicecuda:0) def get_batch(B: int) - tuple[torch.Tensor, torch.Tensor]: x torch.randn(B, D).to(cuda:0) true_y x true_w return (x, true_y) # 获取数据对x 以及 true_y, 用来与训练出的结果 pred_y 对比 train(simple, get_batch, DD, num_layers0, B4, num_train_steps100, lr0.01)num_train_steps 分别设为 10 和 100 跑了两次结果如下3.7 Checkpointing为了防止训练过程中崩溃训练结果全部丢失。在训练中途就可以保存训练中间结果。model Cruncher(dim64, num_layers3).to(get_device()) optimizer AdaGrad(model.parameters(), lr0.01) checkpoint { model: model.state_dict(), optimizer: optimizer.state_dict(), } torch.save(checkpoint, model_checkpoint.pt) loaded_checkpoint torch.load(model_checkpoint.pt) # 加载3.8 Mix Precision Trainig精度更高则更准确、更稳定但成本更高低精度则相反a. 对第五个前向传播使用较低的精度但对其余部分使用 float32b. PyTorch提供了自动进行混合精度训练automatic mixed precisionAMP的工具