How does fsdp algorithm work?

gozdi: <p>Hi<br>
I’m trying to understand fsdp but can’t find any good resources. I’ve broken it down into steps can someone tell me if I’m right??<br>
So lets say that I have 4 gpu’s and a model that can’t fit into one gpu so it gets split into 2 shards.</p>
<ol>
<li>load weights from shard 1 into each gpu</li>
<li>split batch of data into 4 parts</li>
<li>do forward pass with part 1 on gpu 1, part 2 on gpu 2 etc.</li>
<li>save intermediate outputs for backprop locally. They are different on each gpu because of different data</li>
<li>offload weights from shard 1</li>
<li>load weights from shard 2 into each gpu</li>
<li>do forward pass with outputs from step 4 and save intermediate outputs. Again they are different on each gpu because of different inputs. outputs from step 4 are still saved locally</li>
<li>calculate loss</li>
<li>load optimizer state regarding weights from shard 2</li>
<li>calculate gradients for weights from shard 2. They are different on each gpu</li>
<li>synchronize gradients so they are the same on all gpus</li>
<li>update weights on each gpu locally. after update each gpu has the same weights</li>
<li>offload optimizer state and weights from shard 2</li>
<li>load weights from shard 1 and optimizer state regarding them.</li>
<li>calculate gradients for  weights from shard 1. They are different on each gpu</li>
<li>synchronize gradients so they are the same on all gpus</li>
<li>update weights on each gpu locally. after update each gpu has the same weights</li>
<li>offload optimizer state</li>
<li>load another batch of data and split into 4 parts</li>
<li>repeat steps 3-19 until stopping criteria is met</li>
</ol>
<p>Am I right? Would fsdp work like that in a situation I described?</p>

ptrblck: <p><a href="https://pytorch.org/tutorials/intermediate/FSDP_tutorial.html">This tutorial</a> and <a href="https://www.youtube.com/playlist?list=PL_lsbAsL_o2BT6aerEKgIoufVD_fodnuT">these 10 videos</a> might give you the details about its implementation.</p>

agu: <p>If you still have specific FSDP questions after the resources linked, I can try to help.</p>

gozdi: <p>I do<br>
My main question is about general idea how it works. So let’s say that I have 8 gpus and model that can’t fit fully into any of them each can fit half of it. Will fsdp create 2 shards, load shard  1 on all gpus do a forward pass offload shard 1 and do the same with shard 2 or will it load shard 1 on 4 gpus and shard 2 on the other 4 gpus?</p>

gozdi: <p>At high level FSDP works as follow:</p>
<p><em>In constructor</em></p>
<ul>
<li>Shard model parameters and each rank only keeps its own shard</li>
</ul>
<p><em>In forward path</em></p>
<ul>
<li>Run all_gather to collect all shards from all ranks to recover the full parameter in this FSDP unit</li>
</ul>
<p>What does it mean  to collect all shards from all ranks to recover the full parameter in this FSDP unit?</p>
<p>What is an FSDP unit?</p>
<ul>
<li>Run forward computation</li>
<li>Discard parameter shards it has just collected</li>
</ul>
<p><em>In backward path</em></p>
<ul>
<li>Run all_gather to collect all shards from all ranks to recover the full parameter in this FSDP unit</li>
<li>Run backward computation</li>
<li>Run reduce_scatter to sync gradients</li>
<li>Discard parameters.</li>
</ul>

agu: <p>This looks good overall. A few notes:</p>
<ul>
<li>For the backward pass, FSDP discards parameters before reduce-scattering gradients to free that memory earlier.</li>
</ul>
<blockquote>
<p>What does it mean to collect all shards from all ranks to recover the full parameter in this FSDP unit?</p>
</blockquote>
<p>This means that every rank calls all-gather with its local shard as input. After the all-gather, every rank has the full unsharded parameters for the FSDP unit.</p>
<blockquote>
<p>What is an FSDP unit?</p>
</blockquote>
<p>At a high level, an FSDP unit is a layer or collection of layers whose parameters are communicated together. Practically, an FSDP unit refers to one <code>FullyShardedDataParallel</code> instance. In the current design, one <code>FullyShardedDataParallel</code> instance wraps one <code>nn.Module</code> and constructs a single <code>FlatParameter</code> out of the <code>nn.Module</code>’s parameters (that are not already assigned to a nested <code>FullyShardedDataParallel</code> instance).</p>
<blockquote>
<p>So let’s say that I have 8 gpus and model that can’t fit fully into any of them each can fit half of it. Will fsdp create 2 shards, load shard 1 on all gpus do a forward pass offload shard 1 and do the same with shard 2 or will it load shard 1 on 4 gpus and shard 2 on the other 4 gpus?</p>
</blockquote>
<p>There are two dimensions to consider: (1) how many “units” you partition your model into and (2) how many ranks/workers you shard each “unit” across.</p>
<ul>
<li>For (2), the default <code>ShardingStrategy.FULL_SHARD</code> and the <code>ShardingStrategy.SHARD_GRAD_OP</code> strategies shard across all ranks/workers. In your example, this means sharding over the 8 GPUs.</li>
<li>For (1), this depends on how you apply <code>FullyShardedDataParallel</code> to your model. If you only wrap once at the top-level, then there is only one “unit”. However, generally you want more than one “unit” to achieve communication-computation overlap and decrease peak memory. If you only have one “unit”, then the entire model is materialized for forward or for backward, so the entire model contributes to peak memory.</li>
</ul>
<p>Putting this together, suppose you have a model that you partition into 3 “FSDP units”, denoted <code>unit_0</code>, <code>unit_1</code>, and <code>unit_2</code>, and you have 8 GPUs. The mental model can be like:<br>
For <code>i in {0, 1, 2}</code>:</p>
<ul>
<li><code>unit_i</code> all-gather parameters → all 8 ranks have all parameters of <code>unit_i</code></li>
<li><code>unit_i</code> forward</li>
<li><code>unit_i</code> free parameters → rank <code>i</code> only has its <code>1/8</code> of the parameters of <code>unit_i</code></li>
</ul>
<p>For <code>i in {2, 1, 0}</code>:</p>
<ul>
<li><code>unit_i</code> all-gather parameters → all 8 ranks have all parameters of <code>unit_i</code></li>
<li><code>unit_i</code> backward → all ranks have the full gradient for its local batch (representing a partial result)</li>
<li><code>unit_i</code> free parameters → rank <code>i</code> only has its <code>1/8</code> of the parameters of <code>unit_i</code></li>
<li><code>unit_i</code> reduce-scatter gradients → all ranks have its <code>1/8</code> of the gradient for the global batch</li>
</ul>
<p>In practice, due to the nested structure of <code>nn.Module</code>s, you may not see such a simple pattern order of pre/post-forward/backward logic, but the above gives the general mental model.</p>

mileseverett: <p>Would FSDP be ideal in the circumstance where I have a small parameter size, but large calculations within the layers. Meaning that while I can fit all of my layers into memory, trying to store multiple layers gradients causes OOM? Essentially what I’m wanting to do is have one layer processing on one GPU at once and then remove it once it has been processed.</p>
<p>I have tried activation checkpointing but it leads to the models loss not improving.</p>
<p>Thanks<br>
Miles</p>

mileseverett: <p>So if I understand correctly:</p>
<blockquote>
<p>Putting this together, suppose you have a model that you partition into 3 “FSDP units”, denoted <code>unit_0</code>, <code>unit_1</code>, and <code>unit_2</code>, and you have 8 GPUs. The mental model can be like:<br>
For <code>i in {0, 1, 2}</code>:</p>
<ul>
<li><code>unit_i</code> all-gather parameters → all 8 ranks have all parameters of <code>unit_i</code></li>
<li><code>unit_i</code> forward</li>
<li><code>unit_i</code> free parameters → rank <code>i</code> only has its <code>1/8</code> of the parameters of <code>unit_i</code></li>
</ul>
<p>For <code>i in {2, 1, 0}</code>:</p>
<ul>
<li><code>unit_i</code> all-gather parameters → all 8 ranks have all parameters of <code>unit_i</code></li>
<li><code>unit_i</code> backward → all ranks have the full gradient for its local batch (representing a partial result)</li>
<li><code>unit_i</code> free parameters → rank <code>i</code> only has its <code>1/8</code> of the parameters of <code>unit_i</code></li>
<li><code>unit_i</code> reduce-scatter gradients → all ranks have its <code>1/8</code> of the gradient for the global batch</li>
</ul>
<p>In practice, due to the nested structure of <code>nn.Module</code>s, you may not see such a simple pattern order of pre/post-forward/backward logic, but the above gives the general mental model.</p>
</blockquote>
<p>Each GPU is not processing different parts of the model, but is processing a shard of the larger model. E.g. if I split my 200 layer model into 4 equal shards of 50 layers, GPU 1 and 2 would both perform the forward for layers 1-50 and then 51-100 etc and then in the backwards start with 151-200 backwards</p>
<p>Therefore any multi GPU FSDP training should also work fine with single GPU FSDP training as it is the same thing but without the need to gather/scatter parameters between GPUs?</p>

agu: <p>FSDP would not be the right solution for that case since it sounds like you have large activations. (Gradient memory is the same size as parameter memory, so small parameter size means small gradient size.)</p>
<p>Activation checkpointing or some kind of tensor parallelism both can help, where activation checkpointing may be easier to integrate. If your loss is not improving, then that is most likely a bug somewhere either in model code or in the integration of activation checkpointing. It should not change the numerics.</p>

agu: <p>Yes, for the single GPU case, FSDP devolves to be the same as local training. There is no sharding, and there are no collectives (all-gather / reduce-scatter). As such, there is no reason to use FSDP then.</p>

mileseverett: <p>I think I might be misunderstanding. Does FSDP not split a large model into smaller chunks? Meaning that we can train large models one chunk/shard at a time provided that this chunk/shard fits into memory?</p>

agu: <p>In short, the ‘smaller chunks’ are only relevant if there is sharding involved. FSDP’s semantics ensures that each chunk is restored to look like an equivalent local version (without sharding) before computation. However, if the world size (i.e. number of GPUs) is 1, then there is no sharding.</p>
<p>FSDP is not doing anything special to permit you to only partially train some layers/modules of a model.</p>

mileseverett: <p>Oh okay. Much appreciated for clearing up my misconceptions about FSDP. So what does the CPUOffload actually do?</p>

agu: <p><code>CPUOffload</code> means that the parameters are offloaded to CPU when not used for computation.</p>
<p>One way to think about it is that under FSDP, each parameter has two states: <em>sharded</em> and <em>unsharded</em>. The invariant is that the parameter must be unsharded for computation, but outside of computation, FSDP can return it to the sharded state (for memory savings).</p>
<ul>
<li>In the vanilla form, the unsharded state is the original <code>float32</code> parameter on GPU, and the sharded state is the worker’s <code>1 / world_size</code> fraction of the <code>float32</code> parameter on GPU.</li>
<li>If we enable CPU offloading, then the sharded state changes to being instead on the CPU.</li>
<li>If we enable parameter mixed precision, then the unsharded state changes to being in the low precision (e.g. <code>bfloat16</code>).</li>
</ul>

vanshil_shah: <p>where are the activations(result from the forward pass) of the FSDP unit stored in? Like according to my understanding there each gpu recieves different unit of data so when the parameters are unsharded for forward pass calculation, the results of the forward pass remain on the gpu and each gpu would have different result as each gpu recieve different piece of data ?</p>

agu: <p>Since FSDP is a form of data parallelism, each rank has its own activations for its local batch. Your understanding is correct!</p>

lema-balerion: <p>Could you clarify some implementation detail specifics?:</p>
<ol>
<li>Are parameters in an FSDP unit updated immediately after the reduce_scatter for that unit, or only after reduce_scatter has been run for <strong>all units?</strong></li>
<li>Related, are previous gradients freed after the next unit’s gradients are calculated (given that one could keep a running product of the previous terms of the gradient), or are they kept until after all reduce-scatters?</li>
<li>How does activation checkpointing work here? Params are presumably freed, then when it’s time for the backward call, the output of the precursor unit is loaded and run through that unit?</li>
<li>Does HSDP (HYBRID_SHARD) add additional overhead? I’m running HSDP on a single node with 4 GPUs, and I’m noticing HSDP consistently takes up more memory and operates slower than FSDP, even though hypothetically they should be identical for a single node?</li>
</ol>

agu: <ol>
<li>By “updated”, do you mean the optimizer step? If so, then the optimizer step normally is disjoint from the backward, so parameters are updated after all reduce-scatters.</li>
<li>At a high level, for each FSDP unit: compute unsharded gradients → reduce-scatter unsharded gradients to get sharded gradients → free unsharded gradients. The unsharded gradients are not kept until all reduce-scatters have finished. There is some nuance on exactly when the unsharded gradients are freed, but it should be okay to think that they are freed by before the next reduce-scatter.</li>
<li>We generally want activation checkpointing to only happen within an FSDP unit (i.e. checkpointed region should not span multiple FSDP units). Otherwise, the FSDP all-gather would get recomputed in backward (i.e. run twice for one FSDP unit in backward) despite those parameters being needed soon anyway. Instead, if AC is only within one FSDP unit, then FSDP’s pre-backward hook will be responsible for the all-gather, and the recomputed forward does not need to do any all-gathering.</li>
<li>Interesting. I am not exactly sure why it would use more memory. I guess there is no point to do 4-way sharding and 1-way replication with HSDP though.</li>
</ol>

Kyle_Sargent: <p>Does FSDP promise that parameters will be sharded the same way every training step? I can’t find a record of this, but if not, doesn’t this break the recommended optimizer usage pattern?</p>
<p>For instance, suppose that at training step 0 FSDP decides to shard the first half of a linear layer’s params to device 0 and the second half to device 1, and vice-versa for the next training step. Then the optimizer states will be invalid.</p>

agu: <blockquote>
<p>Does FSDP promise that parameters will be sharded the same way every training step?</p>
</blockquote>
<p>It does. FSDP shards parameters at construction time and keeps those sharded parameters in GPU memory thereafter.</p>