应该是一个比较常见的手段了，之前在Megatron MoE的paper中他们有系统的讲。今天在看recompute相关的东西的时候突然想到了这个，就来整理一下。 逻辑是这样的，在MoE fc2之后，我们有一步根据probs加权的操作。 考虑到，对于一个matmul（fc2）来说，他的backward只依赖他的输入，不依赖他的输出。 也就是说如果没有最后那个加权，FC2的输出是不需要save的。 这里就结合到矩阵乘法的特性，g * (h @ W2) = (g * h) @ W2。只要满足g是在行这个维…
背景是之前看torchtitan中用了很多这个SAC，然后看实现比较复杂。最近把dispatch/autograd差不多读了读有了基本的概念，这里来看一下SAC的实现，希望可以更好的整合对dispatch/autograd的理解 实现主要在torch/utils/checkpoint.py中。 这里有一段claude帮忙整理的基本思路，可以对SAC的实现思路有一个了解: 普通(vanilla)非重入 checkpoint 的做法是:forward 时用 saved_tensors_hooks 把所有本该被 auto…
这篇文章来简单整理一下autograd相关的hook，这些hook在优化模型训练方便有比较多的贡献，比如： fsdp所依赖的各种hook，做pre-backward, post-backward。 forward相关的hook是在python的module层，不在autograd save_tensor hook，做activation offload，activation checkpoint都依赖这些hook的实现 Node level Node的定义上有这四种hook。区别是执行的时机和顺序不同。代码…
https://zhuanlan.zhihu.com/p/321449610，介绍了autograd的各种能力 https://dev-discuss.pytorch.org/t/how-to-read-the-autograd-codebase/383 torch论坛的一个autograd的介绍 https://blog.ezyang.com/2019/05/pytorch-internals/: 中间这块确实太过复杂，有基本的概念就快速跳过，有需要再来单独读。 从python怎么写到kernel就行 这里有两个…