Capture training graph with collectives via TorchTitan

amodab01: <p>Hello<br>
I am currently training a llama1b with Torchtitan and hoping to capture the training graph(s) with the various collectives when I work with different parallelism combinations. Is there a way to capture those? I tried to intercept a training step and use the aot_module in functorch.compile to capture it, but I think the Fake Tensor propagation is not working with it. Moreover, is there a better way to capture the graph and all the collectives that get inserted in it as a result of the various degrees of parallelisms?</p>

tianyu: <p>Does <a href="https://github.com/pytorch/torchtitan/tree/main/torchtitan/experiments/simple_fsdp" class="inline-onebox" rel="noopener nofollow ugc">torchtitan/torchtitan/experiments/simple_fsdp at main · pytorch/torchtitan · GitHub</a> provide what you need?</p>
<p>You get the full graph with all collectives in it, although there is no communication optimizations for FSDP in the graph you get, which means comms like all-gathers and reduce-scatters for FSDP are exposed.</p>

amodab01: <p>Thanks <a class="mention" href="/u/tianyu">@tianyu</a><br>
This looks helpful. Will this work with other parallelism strategies, such as TP, CP, and PP?</p>
<p>Also, when I tried to run this simply with <code>TORCH_TRACE="outputs/compile_trace" CONFIG_FILE="./torchtitan/models/llama3/train_configs/llama3_1B.toml" NGPU=4 ./run_train.sh --model.name llama3_simple_fsdp --training.compile</code></p>
<p>I get the following error: `torch._dynamo.exc.Unsupported: TypeError when making fake tensor call<br>
Explanation:</p>
<pre><code>Developer debug context: TypeError &lt;function TensorVariable.method_redistribute.&lt;locals&gt;.redistribute_fn_with_prim_types at 0x7f3cb7d1f9c0&gt;: DTensor.redistribute() got an unexpected keyword argument 'forward_dtype'`
</code></pre>
<pre><code class="lang-auto">  from user code:
     File "/home/amodab01/kyojin/torchtitan/models/llama3/model.py", line 507, in forward
      h = self.tok_embeddings(tokens) if self.tok_embeddings else tokens
    File "/home/amodab01/anaconda3/envs/ml_training/lib/python3.11/site-packages/torch/nn/modules/sparse.py", line 192, in forward
      self.weight,
    File "/home/amodab01/anaconda3/envs/ml_training/lib/python3.11/site-packages/torch/nn/utils/parametrize.py", line 406, in get_parametrized
      return parametrization()
    File "/home/amodab01/anaconda3/envs/ml_training/lib/python3.11/site-packages/torch/nn/utils/parametrize.py", line 299, in forward
      x = self[0](self.original)
    File "/home/amodab01/kyojin/torchtitan/experiments/simple_fsdp/simple_fsdp.py", line 213, in forward
      output = self.replicate_compute(x)
    File "/home/amodab01/kyojin/torchtitan/experiments/simple_fsdp/simple_fsdp.py", line 189, in replicate_compute
      output = x.redistribute(
</code></pre>
<p>I can provide the log files for more detail</p>

tianyu: <p>Yes it should work with all other parallelisms, except for HSDP+TP which we are still working on.</p>
<p>For the error you saw, it’s because SimpleFSDP + mixed precision training relies on a feature we recently added. So you may need to use pytorch nightly, instead of a stable release.</p>

amodab01: <p>Thanks <a class="mention" href="/u/tianyu">@tianyu</a><br>
A quick naive follow up, does the graph automatically get captured? I tried turning on a few torch inductor configs but wasn’t able to find it. I tried these few configs</p>
<pre><code class="lang-auto"> import torch._inductor.config
 torch._inductor.config.trace.debug_log = True
 torch._inductor.config.trace.enabled = True
 torch._inductor.config.trace.graph_diagram = True
 torch._inductor.config.trace.draw_orig_fx_graph = True
</code></pre>

amodab01: <p>Also, is there any way I can intercept the NCCL calls in the graph?</p>

ruisizhang123: <aside class="quote no-group" data-username="amodab01" data-post="5" data-topic="220885">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/amodab01/48/69696_2.png" class="avatar"> amodab01:</div>
<blockquote>
<p>aph automatically get captured? I tried turning on a few torch inductor configs but wasn’t able to fin</p>
</blockquote>
</aside>
<p>You could add a <code>TORCH_COMPILE_DEBUG=1</code> command when running the code, which will dump generated triton code in your local torch_compile_debug folder. Inside <code>torch_compile_debug/*/torchinductor/*/fx_graph_runnable.py</code>, you will find communication operators (e.g., all_gather_into_tensor, reduce_scatter_tensor) along with other compute operations.</p>

amodab01: <p>Hello <a class="mention" href="/u/ruisizhang123">@ruisizhang123</a> <a class="mention" href="/u/tianyu">@tianyu</a><br>
I had a follow-up question on this. I have access to only a few GPUs, but I wanted to replicate how a distributed graph would look like across let’s say 64 devices. Using just torch distributed for a simple model, I was able to achieve that as such</p>
<pre data-code-wrap="python"><code class="lang-python">from torch.testing._internal.distributed.fake_pg import FakeStore
store = FakeStore()
dist.init_process_group(backend=“fake”, rank=rank, world_size=world_size, store=store)
</code></pre>
<pre data-code-wrap="markdown"><code class="lang-markdown">Create DeviceMesh
</code></pre>
<pre data-code-wrap="python"><code class="lang-python">from torch.distributed.device_mesh import DeviceMesh
from torch.distributed.tensor import distribute_tensor, Shard, Replicate
</code></pre>
<pre><code class="lang-auto">mesh = DeviceMesh(“cpu”, list(range(world_size)))
</code></pre>
<pre data-code-wrap="python"><code class="lang-python">class SimpleBlock(nn.Module):def init(self, embed_dim):
    super().init()
    self.linear1 = nn.Linear(embed_dim, embed_dim, bias=False)
    self.linear2 = nn.Linear(embed_dim, embed_dim, bias=False)
</code></pre>
<pre data-code-wrap="python"><code class="lang-python">def forward(self, x):
    residual = x
    x = self.linear1(x)
    x = self.linear2(x)
    return x + residual
</code></pre>
<pre data-code-wrap="python"><code class="lang-python">model = SimpleBlock(embed_dim=512).to(“meta”)
</code></pre>
<pre data-code-wrap="markdown"><code class="lang-markdown">Shard linear1 weight along dim=1 (columns)
</code></pre>
<pre data-code-wrap="python"><code class="lang-python">model.linear1.weight = torch.nn.Parameter(distribute_tensor(model.linear1.weight, mesh, placements=[Shard(1)]))
</code></pre>
<pre data-code-wrap="markdown"><code class="lang-markdown">Replicate linear2 weight
</code></pre>
<pre data-code-wrap="python"><code class="lang-python">model.linear2.weight = torch.nn.Parameter(distribute_tensor(model.linear2.weight, mesh, placements=[Replicate()]))
</code></pre>
<pre data-code-wrap="markdown"><code class="lang-markdown">Sharded input DTensor
</code></pre>
<pre data-code-wrap="python"><code class="lang-python">meta_input = torch.empty(64, 512, device=“meta”, requires_grad=True)
dt_input = distribute_tensor(meta_input, mesh, placements=[Shard(0)])
</code></pre>
<pre data-code-wrap="python"><code class="lang-python">compiled_model = torch.compile(model)
</code></pre>
<p>How can I do something similar with SimpleFSDP? Is that possible?</p>

ruisizhang123: <p>For replicate a distributed graph, you mean doing sth like DDP/HSDP instead of FSDP? yess, it’s possible. In these lines of code: <a href="https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/simple_fsdp/simple_fsdp.py#L275-L284" class="inline-onebox" rel="noopener nofollow ugc">torchtitan/torchtitan/experiments/simple_fsdp/simple_fsdp.py at main · pytorch/torchtitan · GitHub</a> ; mode == “replicate” is for DDP and mode == “hybrid_shard” is for HSDP.</p>