应该是一个比较常见的手段了，之前在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就行 这里有两个…
简单看看tensor/stoage相关的实现，为后面的autograd相关的做一下准备 StorageImpl StorageImpl = 一块连续内存 buffer 的所有者。它只关心三件事:内存在哪(data_ptr_)、有多大(size_bytes_)、谁来分配/释放(allocator_) stride，shape，dtype相关的则是由TensorImpl来负责 注释中： // NB: storage is supposed to uniquely own a data pointer; e.g., //…
boxing相关的概念在这篇博客里也讲的比较详细：https://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/ 在具体看boxing/unboxing实现之前，需要先看一下为什么会有这两种调用方式，这两种调用方式的示例，以及为什么需要支持相互调用 这里有一个概念要分清，就是博客中提到的API和kernel，我们会有box/unbox api，以及box/unbox kernel Unbox api -> unbox kernel就是性…
先看这篇博客理解一下dispatch的高层设计 https://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/ DispatchKey Dispatch key是64bit中的某一个bit。bit位越高优先级就越高。 这块很多的设计都可以从DispatchKey.h中看到设计 有一个关键的设计就是dispatch key被拆分成了两块，分别是 48位的functionality 16位（目前）的backend，包括CPU，CUDA，XL…
https://github.com/ysj1173886760/repo_summary ai帮忙写了一个repo summary的工具 针对指定的仓库，周级别的拉PR，然后总结，用来看这个项目这一周的进展，看看有没有自己比较关注的PR 目前我这里默认关注了megatromlm, pytorch和torchtitan 用的ds v4的模型，差不多一个月几块钱 跑一段时间先看看
MemPool 每个MemPool都靠一个(u64, u64)的二元组来唯一定位。分别是uid和uuid。 在aten/src/Aten/cuda/MemPool.cpp中有详细的解释 uid每次用户申请一个新的mem pool会自增 CUDAGraph每次创建mem pool会增加uuid {0, 0}表示默认的pool 在CUDACachingAllocator中，通过PrivatePool来管理所有的内存池 有一个use_count/cudaMalloc_count，相当于引用计数。当变成0的…
机制 数据类型 何时记录 何时启用 谁消费 Stat 统计 DeviceStats 每次 alloc/free/split/merge 永远开（O(1) 计数） memory_stats() / memory_allocated() 等 Python API \ Trace 历史 TraceEntry 同上 + 段级事件 默认关，需要 _record_memory_history() 打开 memory_snapshot() / _dump_snapshot() Profiler 上报 reportMemoryUs…