Local batch size in FSDP

dcaffo98: <p>I’m struggling to understand how a local batch of data is processed by a single GPU in a Fully Sharded Data Parallel training. As far as I know, a GPU only holds a shard of the model params, it processes its local batch of data, and it belongs to a certain FSDP unit. In general, there could be multiple FSDP units (<a href="https://pytorch.org/tutorials/_images/fsdp_workflow.png" rel="noopener nofollow ugc">see this diagram</a>), each holding a different shard of the model. During the forward pass, the <em>all_gather</em> will collect all the params from the given FSDP unit into each GPU belonging to it. After the <em>all_gather</em>, those GPUs <strong>still have only a fraction of the whole model</strong>, as the <em>all_gather</em> operates inside a single FSDP unit alone. Consequently, local batches of data on those GPUs will be processed by that fraction of the model. When does those batches will be processed by the full model params?</p>

agu: <blockquote>
<p>After the <em>all_gather</em> , those GPUs <strong>still have only a fraction of the whole model</strong> ,</p>
</blockquote>
<p>The GPUs have a fraction of the whole model, but it is along a different axis. You can think of it like pipelining. For example, if we had a sequence of FSDP units (e.g. transformer blocks), FSDP would all-gather one FSDP unit (one transformer block), run forward with the local batch, free the FSDP, all-gather the next FSDP unit, run forward with the local batch, etc.</p>
<p>FSDP units are chosen on the unsharded model, partitioning it into units/groups without sharding the parameters yet. After they are chosen, the parameters within that unit are fully sharded across GPUs. The all-gather operation undoes the sharding of parameters, but not the partitioning of the model into groups. To run forward/backward, only the activate unit needs to be all-gathered.</p>

dcaffo98: <p>Thank a lot for the reply. I’m still confused a bit.</p>
<blockquote>
<p>For example, if we had a sequence of FSDP units (e.g. transformer blocks), FSDP would all-gather one FSDP unit (one transformer block), run forward with the local batch, free the FSDP, all-gather the next FSDP unit, run forward with the local batch, etc.</p>
</blockquote>
<p>Ok, once all FSDP units are done, we have a complete forward pass of the model on a <strong>local</strong> batch. Right? Then is the process repeated to account for every other local batch on other GPUs?</p>

agu: <p>Yes, since FSDP (fully sharded data parallel) is a form of data parallelism, each worker is processing its own local batch.</p>
<p>E.g., if you had 2 GPUs, then GPU0 and GPU1 all-gather the 0th FSDP unit, each compute forward through their respective local batches, and then proceed to the 1st FSDP unit, etc.</p>
<blockquote>
<p>Then is the process repeated to account for every other local batch on other GPUs?</p>
</blockquote>
<p>In other words, this is happening in parallel on all GPUs.</p>

dcaffo98: <aside class="quote no-group" data-username="agu" data-post="4" data-topic="207454">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/agu/48/49913_2.png" class="avatar"> agu:</div>
<blockquote>
<p>E.g., if you had 2 GPUs, then GPU0 and GPU1 all-gather the 0th FSDP unit, each compute forward through their respective local batches, and then proceed to the 1st FSDP unit, etc.</p>
</blockquote>
</aside>
<p>GPUs 0 and 1 do all-gather of 1st FSDP unit even though they belong to the 0th unit? And then they all-gather the 2nd and so on?</p>

agu: <blockquote>
<p>GPUs 0 and 1 do all-gather of 1st FSDP unit even though they belong to the 0th unit?</p>
</blockquote>
<p>I am not sure I follow this. Every GPU has a shard of every FSDP unit.</p>
<p>Said another way (if the FSDP unit terminology is confusing): you can think of each FSDP unit as a module/layer. Suppose your model is a sequence of layers. You can partition your model layer by layer. For each layer, all GPUs all-gather, run forward, and free. Each GPU has a shard of each layer always; only when they all-gather for computation, all GPUs have the entire layer.</p>

dcaffo98: <aside class="quote no-group" data-username="agu" data-post="6" data-topic="207454">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/agu/48/49913_2.png" class="avatar"> agu:</div>
<blockquote>
<p>Every GPU has a shard of every FSDP unit.</p>
</blockquote>
</aside>
<aside class="quote no-group" data-username="agu" data-post="6" data-topic="207454">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/agu/48/49913_2.png" class="avatar"> agu:</div>
<blockquote>
<p>Each GPU has a shard of each layer always</p>
</blockquote>
</aside>
<p>There it is. I thought every GPU only had a single shard, corresponding to a piece of the model in <strong>its</strong> FSDP unit. There is no such thing as <em><strong>its</strong> FSDP unit</em>. Every GPUs has a piece (i.e. a shard) of every FSDP unit. The FSDP unit partitioning is like a logical partitioning of the model. There is no assignment between a given GPU and a given FSDP unit.</p>
<p>As you initially said</p>
<blockquote>
<p>The GPUs have a fraction of the whole model, but it is along a different axis.</p>
</blockquote>
<p>Assuming we see a model as a stack of layers, we can think about the number of layers in the stack as the vertical axis of the model, while the layer width is the horizontal axis. A GPU has full view of the vertical axis (i.e. a shard of each FSDP unit), but only a partial view along the horizontal axis. After running all-gather, the same GPU temporarily gains the full view along the horizontal axis for a <strong>single</strong> FSDP unit (i.e. a full layer of the model), whereas it still has only a shard for each of the other FSDP units.</p>
<p>Many thanks for your patience.</p>

agu: <p>This sounds right to me!</p>