PyTorch

复习 RNN 的 backward 时，我关注的是参数共享：同一份 weight 在多个 timestep 被使用，各处产生的梯度最后怎么汇到它上面？把时间展开成计算图之后，这件事就可以用普通的 chain rule 来算。 training loop 那篇重新理过 forward / backward / optimizer 的职责划分，RNN state transition 那篇理过 hidden state、output、prediction 的区别和 weight sharing，next-token prediction 那篇解释了为什么每个 timestep 都能产生一个 loss。这篇把这几块拼起来，回答一个问题：

起因 # 最近在做 CS336（Language Modeling from Scratch）的 Transformer 实现，代码里反复出现这种写法： 1rearrange(X, "... seq (heads d) -> ... heads seq d", heads=self.num_heads) 第一次看到时，它甚至比 x.view(...).transpose(...) 更加陌生——至少后者你知道它是 PyTorch 原生 API。

起因 # 我们的线上推理系统中存在大量模型——MLP、GNN、MoE，以及它们的各种组合。这个 workload 有几个重要特点： 单张 GPU 上可能同时存在多个模型； 很多模型的 architecture 相同或者高度相似； 模型会随着业务时间轴不断上线、下线和切换； 模型切换过程对 latency 比较敏感。 随着模型数量增加以及动态上下线越来越频繁，我们逐渐观察到两个问题：GPU memory fragmentation，以及 model loading / switching latency。

起因：GPU 没吃满，但不是算力问题 # 最近组里有同学训练模型时，最直接的感受是：GPU utilization 上不去，step time 里总有一段在等数据。

起因 # 最近在学 Stanford CS336（Language Modeling from Scratch），做 tokenizer 作业时，第一次在测试代码里见到了 pytest.mark.xfail。 具体是这段： 1@pytest.mark.xfail( 2 reason="Tokenizer.encode is expected to take more memory than allotted (1MB)." 3) 4def test_encode_memory_usage(): 5 """ 6 We expect this test to fail, since Tokenizer.encode is not expected to be memory efficient. 7 """ 8 tokenizer = get_tokenizer_from_vocab_merges_path( 9 vocab_path=VOCAB_PATH, 10 merges_path=MERGES_PATH, 11 ) 12 with open(FIXTURES_PATH / "tinystories_sample_5M.txt") as f: 13 contents = f.read() 14 _ = _encode(tokenizer, contents) 我当时第一反应是困惑：既然作者明确知道 encode 会超过 1MB，为什么不直接写成 assert memory > 1MB？或者干脆删掉？为什么要提交一个「预期会失败」的测试？

背景 # 公司内部的基于torch的toolbox在某个版本之后，结果发生了偏移。通过一系列排查，发现当导入cupy和torch的顺序不同时，计算结果会有所差异。 也就是说，如下两段代码会导致模型训练等环节得到不同的计算结果。

记录一些常用的…总去查文档也是有点麻烦 * tensor.view 的作用是reshape 比如 a = torch.range(1, 16) 得到一个tensor that has 16 elements from 1 to 16. 在a=a.view(4,4)就得到了一个44的tensor。 需要注意reshape之后元素的个数不能改变(16==44) 参数-1的作用是，我懒得算这一维度应该是多少，（由于元素个数不能改变）所以希望自动被计算。**需要注意的是，只有一个维度可以写-1。 **不过view和reshape有些区别：reshape always copies memory. view never copies memory # * torch.squeeze 将输入张量形状中的1 去除并返回。 如果输入是形如(A×1×B×1×C×1×D)，那么输出形状就为： (A×B×C×D)当给定dim时，那么挤压操作只在给定维度上。例如，输入形状为： (A×1×B), squeeze(input, 0) 将会保持张量不变，只有用 squeeze(input, 1)，形状会变成 (A×B)。注意： 返回张量与输入张量共享内存，所以改变其中一个的内容会改变另一个。 # * torch.unsqueeze 返回一个新的张量，对输入的指定位置插入维度 1 注意： 返回张量与输入张量共享内存，所以改变其中一个的内容会改变另一个。如果dim为负，则将会被转化dim+input.dim()+1 # 1>>> x = torch.Tensor([1, 2, 3, 4]) 2>>> torch.unsqueeze(x, 0) 3 1 2 3 4 4[torch.FloatTensor of size 1x4] 5>>> torch.unsqueeze(x, 1) 6 1 7 2 8 3 9 4 10[torch.FloatTensor of size 4x1] 11 12 13 14 * tensor.expand(size) 扩展tensor.可以保持维度数目不变，每一维度的size增加(比如AB变到C*D，其中C>=A,D>=B).-1参数表示某一个维度的size不发生改变 . 有可以扩展tensor到更多的维度，新增加的维度会默认放在最前面，并且不能以-1作为参数。 # * tensor.contiguous 将一个tensor变成连续的。（一些ops如expand/expand_as会让tensor 不连续） #