TP/FSDP + sync_module_states/cpu_offload

Corey_Lowman: <p>Hi, I’m working on writing a distributed training guide.</p>
<p>One of the chapters I’m working on currently is a chapter on 2d parallelism (tensor parallel + FSDP). The  full code is in the PR here: <a href="https://github.com/LambdaLabsML/distributed-training-guide/pull/39" class="inline-onebox" rel="noopener nofollow ugc">[WIP] 2d parallelism chapter by corey-lambda · Pull Request #39 · LambdaLabsML/distributed-training-guide · GitHub</a></p>
<p>I have a couple of questions, but first here’s some background:</p>
<ol>
<li>I have a 64 GPU cluster (8 nodes, each with 8 H100 GPUs)</li>
<li>I’m training Llama 3.1 405B</li>
<li>I load the pretrained weights on rank 0 on the CPU, and every other rank uses the meta device</li>
<li>When just using fsdp (you can see the full code for this here: <a href="https://github.com/LambdaLabsML/distributed-training-guide/blob/main/06-training-llama-405b/train_llm.py" class="inline-onebox" rel="noopener nofollow ugc">distributed-training-guide/06-training-llama-405b/train_llm.py at main · LambdaLabsML/distributed-training-guide · GitHub</a>), I used <code>sync_module_states=True</code> and <code>CPUOffload(offload_params=True)</code></li>
</ol>
<p>For 2d parallelism:</p>
<ol>
<li>I’m using the pytorch example code from the TP/FSDP example as a reference (can’t include link for some reason, but I can share in a comment)</li>
<li>I have a 2d device mesh (size 8 in dp and size 8 in tp dimension)</li>
<li>I’m applying tensor parallel to all the weights of the model along the TP dimension of the 2d mesh</li>
<li>I’m applying FSDP along the DP dimension of the 2d mesh</li>
</ol>
<p>Now for my questions:</p>
<ol>
<li>I keep getting OOM errors when FSDP starts materializing the tensors. I’m wondering if CPU Offloading works with DTensors from tensor parallel (maybe when combined with meta device?)</li>
<li>How does sync_module_states=True work with FSDP &amp; device_mesh? Since the weights are not really shared between the ranks anymore (they are sharded in device mesh groups), do I need to do something else to ensure that all the ranks have the correct weights?</li>
</ol>

Corey_Lowman: <p>Another update:</p>
<p>So since I moved to using TP, I removed the auto wrapping policy from the FSDP constructor. That means the only module that is wrapped by FSDP is the top level module.</p>
<p>Interestingly, parameters during initialization are batch moved between devices instead of 1 by 1. So since there is only 1 FSDP here, all 405b params (split into 8 TP shards) are moved directly to GPU device when calling param_init_fn: <a href="https://github.com/pytorch/pytorch/blob/v2.4.1/torch/distributed/fsdp/_init_utils.py#L561" class="inline-onebox" rel="noopener nofollow ugc">pytorch/torch/distributed/fsdp/_init_utils.py at v2.4.1 · pytorch/pytorch · GitHub</a>.</p>
<p>Since I’m using 8 H100 nodes I have a total of 640 GB of GPU memory on a node, but 405b bf16 params take ~800 GB of gpu memory, so that’s why I need CPU offloading or staged device transfer.</p>
<p>I changed my param_init_fn to initialize on CPU device instead, but am getting another error now about sync module states. I may need to bring back auto_wrap_policy just for staged cpu offloading.</p>
<p>I also think I need to disable sync_module_states. With TP, rank 0 only has a shard of the model parameters, not the full set, so it can’t actually share with all the ranks.</p>
<p>Instead I think I need to have all the nodes load the weights into CPU at first and not use the meta device. Or figure out some way for all node 0 (rank 0-8) to broadcast weights to the other nodes (i.e. rank 1 would broadcast to rank 9/17/etc).</p>

agu: <p>Do you think it would be possible to use FSDP2? The TP support should be simpler and better. <a href="https://github.com/pytorch/torchtitan" class="inline-onebox" rel="noopener nofollow ugc">GitHub - pytorch/torchtitan: A native PyTorch Library for large model training</a></p>

Corey_Lowman: <p>Okay I will check it out - I’m learning that TP support in pytorch is still evolving rapidly from reading through some of the stuff in torchtitan. Thanks for linking that.</p>