之前有过FSDP1相关的介绍，这次来看一看FSDP2，也是偏源码分析级 不过有一个特殊的点是FSDP2在Github上的文档非常全面，把相关特性的支持，代码结构的设计讲的都很清楚，所以这篇文章主要是来做一个补全。推荐在阅读FSDP2的代码之前，先看看这个文档https://github.com/pytorch/pytorch/issues/114299 还有一个点是因为我个人对torch dynamo相关的不太熟悉，所以FSDP2和编译优化相关的事情就不提了 FSDP2和FSDP1个人认为最主要的区别点有几个： 不…
整理一下FSDP1相关的实现，偏源码级，适合想读代码的同学来看 Core 核心流程如图 图中可以看到，FSDP虽然shard了模型，但是在计算(forward)的过程中，还是需要收集所有的参数做计算。所以他还是被划分成DataParallel Init 主要分为两块： 把模型切分成fsdp unit 把每一个fsdp unit的参数展开成FlatParameter 以一个CausalTransformer为例，模型会被组织成一个树形的结构。每一个节点都会有自己的Parameter，定义自己的forward计算逻辑。…