Supporting Autograd for Collectives

username: wconstab
cooked: <p><a href="mailto:fmassa@meta.com">Francisco Massa</a> <a href="mailto:lcw@meta.com">Luca Wehrstedt</a> <a href="mailto:kw2501@meta.com">Ke Wen</a> <a href="mailto:whc@meta.com">Will Constable</a></p>
<p>TL;DR: We avoided defining backwards for collectives for too long, because it was not urgent enough to push through some obscure stuff and align on an implementation. Now we’re seeing more usage of collectives and missing backwards support is becoming a serious problem.</p>
<h1><a name="p-469680-why-has-this-taken-so-long-1" class="anchor" href="#p-469680-why-has-this-taken-so-long-1"></a>Why Has This Taken So Long?</h1>
<p>In short, because it’s harder than it sounds and nobody prioritized it enough. Autograd support was part of the ‘goals’ in the first <a href="https://github.com/pytorch/pytorch/issues/93173" rel="noopener nofollow ugc">Functional Collective RFC</a>, but it was never a priority because in so many cases (FSDP, DDP, PP) there is no need for collective backwards, and for DTensor users, gradients are handled by DTensor itself. Worse, the gradient formula for a given collective is context-dependent, calling into question the validity of providing a ‘default’ backwards implementation in the first place (e.g. <a href="https://github.com/pytorch/pytorch/pull/123599#issuecomment-2044899594" rel="noopener nofollow ugc">this comment</a>).</p>
<p>I hope this post will serve first to demystify the issues around collective backwards. Second, I think there are some immediate actions we can take to improve things. And Third, we may need additional discussion on larger changes (e.g. RFCs), and we can refer back to this post as a ground truth.</p>
<h1><a name="p-469680-what-are-the-correct-backwards-formulas-anyway-2" class="anchor" href="#p-469680-what-are-the-correct-backwards-formulas-anyway-2"></a>What Are the Correct Backwards Formulas Anyway?</h1>
<p>Let’s start by writing out the expected formulas from first principles, avoiding the implementation detail of multiple processes or comm libraries.</p>
<h4><a name="p-469680-example-naive_all_reduce-3" class="anchor" href="#p-469680-example-naive_all_reduce-3"></a>Example: naive_all_reduce</h4>
<p>We can implement all_reduce as a function taking many input tensors and returning the reduced tensor multiple times as output.</p>
<pre><code class="lang-auto">def naive_all_reduce(inputs: list[torch.Tensor]) -&gt; list[torch.Tensor]:
   common_output = torch.sum(inputs)
   return [common_output for _ in inputs]
</code></pre>
<p>Because we return multiple copies of the summed output, the gradients that come back from each ‘output path’ will be summed into the .grad field of ‘common_output’, before being passed back to each of the inputs. So this means our backwards pass is essentially another sum allreduce:</p>
<pre><code class="lang-auto">def naive_all_reduce_backward(d_outputs: list[torch.tensor]) -&gt; list[torch.Tensor]:
   accumulated_grad = torch.sum(d_outputs)
   return [accumulated_grad for _ in d_outputs]
</code></pre>
<p>Using this framing, the mathematical formula for each operator should be unambiguous and indeed easy to verify just by using autograd on top of local torch operations.</p>
<p>One thing to call out, which will become important later: we are not making assumptions about our input/output tensors being related to one another in any way! This is the default assumption in pytorch, and indeed it is what you’d get if you ran torch.autograd on <code>naive_all_reduce</code>. But as we will see later, it is not what we actually want for many applications in distributed training.</p>
<h3><a name="p-469680-collectives-4" class="anchor" href="#p-469680-collectives-4"></a>Collectives</h3>
<div class="md-table">
<table>
<thead>
<tr>
<th>Forward</th>
<th>Backward</th>
<th>Notes</th>
</tr>
</thead>
<tbody>
<tr>
<td>gather</td>
<td>scatter</td>
<td></td>
</tr>
<tr>
<td>scatter</td>
<td>gather</td>
<td></td>
</tr>
<tr>
<td>reduce (avg, sum, premul_sum)</td>
<td>broadcast</td>
<td>Bitwise ops not supported for grad (band, bor, bxor)</td>
</tr>
<tr>
<td>reduce (max, min)</td>
<td>Identity (for max/min src) Scaled (for a tie) 0 (for others)</td>
<td></td>
</tr>
<tr>
<td>reduce (product)</td>
<td>fwd_out / fwd_in * dout</td>
<td></td>
</tr>
<tr>
<td>broadcast</td>
<td>reduce(sum)</td>
<td></td>
</tr>
<tr>
<td>all_to_all</td>
<td>all_to_all</td>
<td></td>
</tr>
<tr>
<td>all_reduce (avg, sum, premul_sum)</td>
<td>all_reduce(sum)</td>
<td>Common exception, e.g. megatron TP, see below; Bitwise ops not supported for grad (band, bor, bxor)</td>
</tr>
<tr>
<td>all_reduce (max, min)</td>
<td>all_reduce(sum) (for max/min src) 0 (for others)</td>
<td></td>
</tr>
<tr>
<td>all_reduce (product)</td>
<td>fwd_out / fwd_in * allreduce(sum, dout)</td>
<td></td>
</tr>
<tr>
<td>all_gather</td>
<td>reduce_scatter(sum)</td>
<td></td>
</tr>
<tr>
<td>reduce_scatter</td>
<td>all_gather</td>
<td></td>
</tr>
<tr>
<td>all_to_all</td>
<td>all_to_all</td>
<td></td>
</tr>
</tbody>
</table>
</div><h3><a name="p-469680-non-spmd-cases-5" class="anchor" href="#p-469680-non-spmd-cases-5"></a>Non-SPMD cases</h3>
<p>We probably should not implement backwards automatically for non-spmd communication operations, even though the mathematical formula is trivial. This is because the ordering of operations during the backwards pass is critical to maintain correctness and avoid hangs, but we have no practical way of ensuring that we execute backwards graphs in a particular order. This can technically affect SPMD collectives too, but is particularly bad for Non-SPMD collectives.</p>
<div class="md-table">
<table>
<thead>
<tr>
<th>Forward</th>
<th>Backward</th>
</tr>
</thead>
<tbody>
<tr>
<td>send</td>
<td>recv</td>
</tr>
<tr>
<td>recv</td>
<td>send</td>
</tr>
</tbody>
</table>
</div><p>Now that we have enumerated the backwards formulas, let’s see how people actually use collectives in practice.</p>
<p>These ‘mathematically correct’ backwards formulas for collectives are definitely useful in some cases. At least we have seen several examples recently: <a href="mailto:tristanr@meta.com">Tristan Rice</a> used some for his Context parallelism prototype (<a href="https://github.com/pytorch/pytorch/commit/ddd0ed1b430f3241490d9bc071036188466b0f22" rel="noopener nofollow ugc">code</a>), and <a href="mailto:kw2501@meta.com">Ke Wen</a> used some for his DeepSeek implementation (<a href="https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/deepseek_v3/symm_mem_recipes/triton_on_device_all_to_all_v.py#L163" rel="noopener nofollow ugc">code</a>).</p>
<h1><a name="p-469680-exceptions-that-overrule-the-rule-6" class="anchor" href="#p-469680-exceptions-that-overrule-the-rule-6"></a>Exceptions That Overrule The Rule</h1>
<p>Let’s take a look at Tensor Parallelism as described in the original Megatron paper [<a href="https://arxiv.org/pdf/1909.08053" rel="noopener nofollow ugc">1909.08053</a>]. This is just one example, but the pattern it describes is extremely common in deep learning (moreso than the examples I could list above).</p>
<p><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXe2odWoNerq6OblmvKEz-1AkkBVxVrjH8ch8uuizsMFiUc4ynkLUbvFs28OtGZCuTW9HQAwX24YjfteybrmflwL8dFrYj329TlCPaaO-FKMb8k2gBGwjm2UkGzzbb8wy37LKzW9?key=XZowS_22YGi5PwBhC4WyijiM" alt="" width="398" height="165" role="presentation"></p>
<p>Figure 3. Blocks of Transformer with Model Parallelism. f and g are conjugate. f is an identity operator in the forward pass and all reduce in the backward pass while g is an all reduce in the forward pass and identity in the backward pass.</p>
<p>Communication operations f and g are introduced to perform a particular pattern of communication during forward and backwards that does not match any particular operator in our table above!</p>
<p>The whole point of f and g is to make a semantic leap. f takes a replicated tensor and makes it non-replicated. g merges non-replicated tensors back to replicated ones.</p>
<p>PyTorch DTensor has a similar concept. DTensors automatically track whether a tensor is a ‘Replica’ (e.g. X input to f) or ‘Partial’ (e.g. Z1 input to g), and offers APIs like ‘redistribute()’ to convert between these layouts.</p>
<p>Jax uses typing to tag values as ‘device-variant’ or ‘device-invariant’ [<a href="https://docs.jax.dev/en/latest/jep/17111-shmap-transpose.html#full-solution-statically-track-device-varying-vs-device-invariant-intermediates-plus-new-primitives" rel="noopener nofollow ugc">docs</a>], and defines a set of operators that move values back and forth [<a href="https://docs.jax.dev/en/latest/jep/17111-shmap-transpose.html#tracking-device-invariance-in-avals-a-k-a-avals-with-names-revived" rel="noopener nofollow ugc">docs</a>]. <code>pbroadcast</code> is equivalent to f, while <code>psum2</code> is equivalent to g. Note: <code>psum</code> itself also exists, and is equivalent to ‘all_reduce’, and the whole reason for introducing <code>psum2</code> is to solve the issues outlined above.</p>
<p><img src="https://lh7-rt.googleusercontent.com/docsz/AD_4nXf-AZlCkFbTUzbKs2y8e3LMkhgxuHl1u2xHaGLHk1c5k3uS8a0euEQfVdF4qv1tHMxFej_lg0zcnAMpH_lojZrL1yVixIZnrlJIH5kxMNAkusWY-8Glm85kI_nzSm6JjS4gopwJyQ?key=XZowS_22YGi5PwBhC4WyijiM" alt="" width="595" height="188" role="presentation"></p>
<p>Figure 4. Snippet from the table of operators introduced in JAX to transform between device-invariant and device-variant states.</p>
<h1><a name="p-469680-so-what-should-we-do-7" class="anchor" href="#p-469680-so-what-should-we-do-7"></a>So What Should We Do?</h1>
<p>To recap, we can define mathematically pure gradient formulas for our collectives, but they don’t address some of the most common use cases. And we have a higher-level concept (call it replica vs partial, or device invariant vs variant) that we can make clearer and better supported in PyTorch.</p>
<h3><a name="p-469680-fix-and-improve-collective-apis-8" class="anchor" href="#p-469680-fix-and-improve-collective-apis-8"></a>Fix and Improve Collective APIs</h3>
<ol>
<li>
<p>[P0] We should patch the silent-correctness issue with functional collectives and C10d collectives by throwing an error in backwards (<a href="https://github.com/pytorch/pytorch/issues/152131" class="inline-onebox" rel="noopener nofollow ugc">[C10D] Autograd Support for Collectives · Issue #152131 · pytorch/pytorch · GitHub</a></p>
</li>
<li>
<p>[P1] Implement mathematically pure backwards formulas for functional collectives and C10d collectives. (Note: naive backwards impl would be synchronous.) ((<a href="https://github.com/pytorch/pytorch/issues/152131" class="inline-onebox" rel="noopener nofollow ugc">[C10D] Autograd Support for Collectives · Issue #152131 · pytorch/pytorch · GitHub</a>))</p>
</li>
<li>
<p>Deduplicate funcol, autograd_funcol, and nn.functional</p>
</li>
<li>
<p>[P2 / Optional] Improve eager backwards performance. (Alternative: rely on compile)</p>
</li>
</ol>
<h3><a name="p-469680-explicitly-support-the-device-varying-case-9" class="anchor" href="#p-469680-explicitly-support-the-device-varying-case-9"></a>Explicitly Support the Device-Varying Case</h3>
<p><a href="mailto:lcw@meta.com">Luca Wehrstedt</a> and <a href="mailto:fmassa@meta.com">Francisco Massa</a> point out that the existing DTensor internal helper method <code>redistribute_local_tensor</code> is very nearly what we’d need to offer a set of manual APIs that let users convert plain torch.Tensor between logically ‘replicated’ and ‘partial/device-varying’ states.</p>
<p>For users that don’t want to use DTensor for their whole program, we can still leverage its internals to offer a manual API that takes care of correctly transforming tensors between device-varying and device-invariant forms. In other words, we can offer the equivalent of f and g operations (as well as others described by JAX).</p>
<p>However, there are a few things we might need to harden about DTensor’s semantics. We’ll have to dive a little deeper here.</p>
<ul>
<li>Currently DTensor sharding rules are not as disciplined as they could be: A ‘replica’ tensor can directly be used as input to a matmul with a ‘shard’ input, while in a stricter sense we should force DTensor to convert the replica to a ‘partial’ first</li>
<li>‘redistribute_local_tensor’ has some questionable logic about skipping partial to replica conversion during backwards</li>
</ul>
<p>I plan to look further into the explicit API and DTensor case, and update RFCs on pytorch accordingly. Hopefully this note clarifies the topic for folks, and if there is still something unclear or a mistake in the reasoning, please let me know!</p>

username: wanchaol
cooked: <p>Thanks for writing this down!</p>
<blockquote>
<p>Currently DTensor sharding rules are not as disciplined as they could be: A ‘replica’ tensor can directly be used as input to a matmul with a ‘shard’ input, while in a stricter sense we should force DTensor to convert the replica to a ‘partial’ first</p>
</blockquote>
<p>Can you elaborate this more? I don’t understand the statement here. There’s no partial input should be accepted to a matmul.</p>
<blockquote>
<p>‘redistribute_local_tensor’ has some questionable logic about skipping partial to replica conversion during backwards</p>
</blockquote>
<p>Hmm what is questionable here?</p>
<p>I feel supporting collective autograd formulas is a good direction and could help dedupe and remove the undocumented nn.functional collectives. There’s a concern about backward compatibilty (i.e. how to actually dedup those implementations, are autograd version of collectives stay different API to avoid BC breaking). I feel a concrete proposal on the details would be more helpful <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=14" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>But I think it is questionable whether to expose some manual f/g operations, as the f/g operations actually embedding the concepts of sharding implicitly without any explicit sharding annotations. There is not something general enough that could be exposed, because the input shardings would be different for case to case. i.e. <code>redistribute_local_tensor</code> needs user to explicitly pass src/dst sharding spec in order to deduce the correct operation. Also this API does not implement autograd at all, and if you implement autograd, this is basically the <code>redistribute</code> api DTensor had today</p>
<p>Honestly for users who don’t want to use DTensor in the whole program, there’s already a way for them to construct a DTensor via from_local, do redistribute, and convert back to local tensor. Each operation is very explicitly and general already, so I don’t see a need to introduce some other API to achieve the same thing.</p>
<p>IMO it’s better to separate the DTensor discussion to the autograd collectives posted here as I think they are cleanly two separate layers of abstractions.</p>

username: wconstab
cooked: <aside class="quote no-group" data-username="wanchaol" data-post="2" data-topic="219430">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/wanchaol/48/50343_2.png" class="avatar"> wanchaol:</div>
<blockquote>
<p>But I think it is questionable whether to expose some manual f/g operations, as the f/g operations actually embedding the concepts of sharding implicitly without any explicit sharding annotations.</p>
</blockquote>
</aside>
<p>Totally agree that this is ‘implicit’ and that makes it dangerous.  Point of fact, lots of users do it anyway.  I think its worth offering an API to do it and document what it means.  But I <em>also</em> think its cleaner and safer to just use DTensor.  We should recommend this path for this reason.</p>
<blockquote>
<p>Can you elaborate this more? I don’t understand the statement here. There’s no partial input should be accepted to a matmul.</p>
</blockquote>
<p>Yea I need to spend more time looking into this part. I will try to write something further once I do.</p>
<blockquote>
<p>. <code>redistribute_local_tensor</code> needs user to explicitly pass src/dst sharding spec in order to deduce the correct operation.</p>
</blockquote>
<p>This part seems fine to me.  Either (a) the API is that users pass src/dst sharding spece manualy, or (b) we offer apis like ‘psum2’ in jax, which desugars into this.  Either way, its a way to manipulate these states even though there is no metadata that makes sure you don’t mess up.</p>
<blockquote>
<p>Also this API does not implement autograd at all, and if you implement autograd, this is basically the <code>redistribute</code> api DTensor had today</p>
</blockquote>
<p>Yea, maybe what I <em>actually</em> want is to use <code>redistribute</code> but then remove the DTensor wrapper immediately.  I definitely want autograd supported inside this API.  The idea would be that users can do the ‘replicate → partial’ transformation explicitly during forward (which is really a no-op), and then they would get an all-reduce ‘for free’ during backwards.</p>

username: Chillee
cooked: <p>A couple thoughts:</p>
<blockquote>
<p>Worse, the gradient formula for a given collective is context-dependent, calling into question the validity of providing a ‘default’ backwards implementation in the first place (e.g. <a href="https://github.com/pytorch/pytorch/pull/123599#issuecomment-2044899594" rel="noopener nofollow ugc">this comment</a>).</p>
</blockquote>
<p>I was talking with Yifu, and in some sense, I agree with him that there is indeed a single “correct” backwards implementation of all-gather (or any other collective). The argument is similar to the one given with naive_all_reduce - namely, if you view every output of the collective as just another usage of the tensor, then the backwards pass should clearly just be the sum of the gradients of every output!</p>
<p>For example, say we had</p>
<pre><code class="lang-auto">x1, x2, x3, x4 = y
loss = f(x1, x2, x3, x4)
</code></pre>
<p>clearly, <code>gradX = sum(gradX1, gradX2, gradX3, gradX4)</code>!</p>
<p>However, if we were to talk about this from the perspective of DTensor, what are we saying about the layouts of <code>gradX1/gradX2/...</code>? We are saying that these are all <strong>partial sums</strong> of <code>gradX</code>.</p>
<p>i.e. when I write this and it runs on 4 ranks</p>
<pre><code class="lang-auto">y = allreduce(x)
</code></pre>
<p>in reality, I actually have <code>y_0, y_1, y_2, y_3</code>.</p>
<p>So, using DTensor vocab, we are guaranteeing that the <code>gradY</code> above is <em>Partial</em>. In other words, from the autograd engine’s perspective, <code>gradY</code> only contains the loss gradient wrt <code>y_0</code> and not <code>y_1/y_2/y_3</code>. However, in DTensor, there is another layout it could be - Replicate. And if it were replicated, then that would mean that <code>gradY_0</code> contains <code>sum(gradY_i)</code>, and so if the backwards pass were an allreduce that would be wrong.</p>
<p>So, if we have our differentiable collectives, can we guarantee that we’ll never run into the <code>Replicate</code> case?</p>
<p>Actually, <strong>yes</strong>! It is consistent for us to say that on every single device, <code>gradX</code> always represents the gradient contribution of <code>X</code> to the final loss. Thus, the proposal above is sufficient - allgather ↔ reducescatter, allreduce ↔ allreduce, and we’re good!</p>
<p>However, there were a couple of inconsistencies here that made me feel like this wasn’t a complete picture. For example:</p>
<ol>
<li>If there is only one correct backwards for a collective, why does the original TP formulation Megatron-LM clearly <em>not</em> use those backwards? (i.e. in figure 3, <code>g</code> is an allreduce, but its backward is the identity!). The post says that these are for a semantic leap, but even so, this seems like a weird contradiction to me. The post writes “the mathematical formula for each operator should be unambiguous” but then for Megatron, the backwards pass of an allreduce is the identity??? How can it both be unambiguous but also variable?</li>
<li>Can we implement DDP with these autogradable collectives? I mean, the backwards of DDP is an allreduce, but there’s no corresponding allreduce in the forwards? That’s kinda weird.</li>
</ol>
<p>After thinking about it, my conclusion is that: <strong>The original sin is that we often want parameters (leaf nodes!) to be replicated across machines.</strong> All other inconsistencies flow downstream from this desire.</p>
<p>Basically, this does not make sense from the autograd perspective. Multiple leaf nodes that actually represent the same value? A replicated parameter doesn’t make any sense - the “morally” correct thing to do is to have it live on rank 0 and then be broadcasted out to every rank if we want each rank to use it, which then makes the parameter “sharded”, but it doesn’t matter because the backwards of a broadcast is a “reduce”, which gives us our desired semantics.</p>
<p>So, really, when we have a replicated parameter, before using it, we are essentially inserting a “fake_broadcast(Replicated) → Sharded”. The backwards pass of this is a <code>fake_reduce</code> (implemented with an allreduce). This answers question 2 - if we want to implement DDP then we need to support replicated parameters, and in order to support replicated parameters we need a <code>fake_broadcast</code>, which morally acts as if it broadcasts all values from rank 0, but in reality is a no-op.</p>
<p>Now, let’s look at question 1 - why do the collectives in Megatron-LM’s collectives not follow the given autograd rules?</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/0/00079e220c5851117010639cd4a2c608220cfc59.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/00079e220c5851117010639cd4a2c608220cfc59" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/0/00079e220c5851117010639cd4a2c608220cfc59_2_690x251.jpeg" alt="image" data-base62-sha1="gjToptd0n87LJXyvmYWTWTBOx" width="690" height="251" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/0/00079e220c5851117010639cd4a2c608220cfc59_2_690x251.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/0/0/00079e220c5851117010639cd4a2c608220cfc59_2_1035x376.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/0/00079e220c5851117010639cd4a2c608220cfc59_2_1380x502.jpeg 2x" data-dominant-color="E7EBE7"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1392×508 106 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>First of all, let’s imagine implementing Megatron-LM using the autogradable collectives listed above. To do so, we would skip <code>f</code>, and <code>g</code>’s backwards would be another allreduce. Kinda like this diagram</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/2/e2bf69b2b84865267b2340d15befb720c007e142.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/e2bf69b2b84865267b2340d15befb720c007e142" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/2/e2bf69b2b84865267b2340d15befb720c007e142_2_517x244.jpeg" alt="image" data-base62-sha1="wlU0wNyqYJGtV4ymu2xtXzo7z2i" width="517" height="244" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/2/e2bf69b2b84865267b2340d15befb720c007e142_2_517x244.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/e/2/e2bf69b2b84865267b2340d15befb720c007e142_2_775x366.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/e/2/e2bf69b2b84865267b2340d15befb720c007e142.jpeg 2x" data-dominant-color="CDD9CE"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1032×488 79.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Ignoring anything outside of the block (like norms) for a second, note that this is totally correct! The transformer block goes back to back, and so if we remove <code>f</code> from the equation, it immediately goes into the previous transformer block’s <code>g</code>, thus performing the allreduce that would have been done by <code>f</code>! So… if they could use these “autogradable” collectives, why did Megatron-LM formulate it differently? And what’s the difference?</p>
<p>As it turns out, there is a difference! And it all comes back around to the original sin - replicated parameters. Inbetween the transformer blocks we also have layernorms, which have parameters. In Megatron-LM, they replicate the layernorm parameters.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/0/804d05b9bcd8a6b745be1dafa71f24bdc62723c8.png" data-download-href="https://discuss.pytorch.org/uploads/default/804d05b9bcd8a6b745be1dafa71f24bdc62723c8" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/0/804d05b9bcd8a6b745be1dafa71f24bdc62723c8_2_345x103.png" alt="image" data-base62-sha1="ij08qiDx4IlLTvjoOoEv96OFcA8" width="345" height="103" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/0/804d05b9bcd8a6b745be1dafa71f24bdc62723c8_2_345x103.png, https://discuss.pytorch.org/uploads/default/optimized/3X/8/0/804d05b9bcd8a6b745be1dafa71f24bdc62723c8_2_517x154.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/0/804d05b9bcd8a6b745be1dafa71f24bdc62723c8_2_690x206.png 2x" data-dominant-color="D7D7D7"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">970×292 25.5 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>So, let’s use the above view that “replicated” actually means “lives on rank 0”. So, in our scheme, what would we do?</p>
<p>To write it out, in our scheme using autogradable collectives, we would have:</p>
<pre data-code-wrap="python"><code class="lang-python">X: Sharded # beginning
.... # bunch of stuff
Z: Sharded = all_reduce(X) # g forwards
norm_weight: Rank0
broadcast_norm_weight: Sharded = fake_broadcast(norm_weight) # no-op
Z_normed: Sharded = layer_norm(Z: Sharded, broadcast_norm: Sharded)
X2: Sharded = Z_normed # starts next layer
</code></pre>
<p>and thus this in the backwards</p>
<pre data-code-wrap="python"><code class="lang-python">grad_X2: Sharded
grad_Z_normed: Sharded = grad_X2
grad_broadcast_norm_weight: Sharded = layer_norm_backward(grad_Z_normed: Sharded)
grad_norm_weight: Rank0 = fake_reduce(grad_broadcast_norm_weight) # all-reduce
gradX: Sharded = all_reduce(grad_Z) # g backwards
</code></pre>
<p>Note that here, we need to allreduce <code>norm_weight</code>, since its gradient computation is sharded.</p>
<p>However, the Megatron-LM strategy looks like this:</p>
<pre data-code-wrap="python"><code class="lang-python">X: Sharded # beginning
.... # bunch of stuff
Z: Rank0 = fake_reduce(X) # g forwards
norm_weight: Rank0
Z_normed: Rank0 = layer_norm(Z: Rank0, broadcast_norm: Rank0)
X2: Sharded = fake_broadcast(Z_normed)  # f forwards
</code></pre>
<p>with the corresponding backwards being</p>
<pre data-code-wrap="python"><code class="lang-python">grad_X2: Sharded
grad_Z_normed: Rank0 = fake_reduce(grad_X2)  # f backwards
grad_norm_weight: Rank0 = layer_norm_backward(grad_Z_normed: Rank0)
gradX: Sharded = broadcast(grad_Z) # g backwards
</code></pre>
<p>Basically, because <code>norm_weight</code> is replicated (i.e. on “Rank0”), there “should” be an allreduce in the backwards pass. However, if you ensure that all computation performed with it is also replicated (i.e. on “Rank0”), then you can avoid the <code>fake_broadcast</code>, and thus avoid the allreduce.</p>
<p>Note: The general insights here bears similarities to Jax’s page on “device-varying” vs “device-invariant” values (although on their page it’s kind of a weird special case for the interactions between shmap output and pjit). In their terminology, <code>fake_broadcast</code> == <code>pbroadcast</code> and goes from “invariant” to “device-varying”. However, because shmap is fundamentally a function wrapper, their doc doesn’t touch on “replicated parameters”. I also prefer the conceptual framing that they morally only live on one device (i.e. rank0).</p>
<h2><a name="p-469704-conclusion-1" class="anchor" href="#p-469704-conclusion-1"></a>Conclusion</h2>
<p>This was way too long haha. Arguably, this is a lot of words to not really disagree with anything in the original post. But I hope the explanation of why the “Megatron-LM exception” is not really an exception is useful.</p>
<p>But, my overall thoughts on the plans/what to do next.</p>
<ol>
<li>
<p>I think it makes sense to provide the “autograd-able collectives” with the defined autograd formulas, as suggested above. However, we should emphasize that these assume that no inputs are “replicated”/“device-invariant”/“morally live on rank0”. For example, if you were to implement the megatron-style TP without the additional <code>fake_broadcast</code>, we would end up with different gradients for every rank’s <code>norm_weight</code>, essentially treating that as if it were sharded and not replicated.</p>
</li>
<li>
<p>I think we should consider adding the collective pairs that map from “replicated” to “sharded” and vice-versa. Specifically, I think there’s “fake_broadcast” (broadcasts from rank0, rank0 → sharded, no-op ↔ allreduce), “fake_reduce” (reduces down to rank 0, sharded → rank0, allreduce ↔ noop), “fake_scatter” (scatters from rank0, rank0 → sharded, split ↔ all-gather), and “fake_gather” (gathers to rank0, sharded → rank0, all-gather ↔ split). We wouldn’t be able to guarantee that users are necessarily using them correctly, but I think this would allow users to basically implement most fundamental collective patterns via autograd, albeit not necessarily efficiently (we could also add a “slow” mode that does a check to ensure that “replicated” values are actually replicated).</p>
</li>
<li>
<p>I think we need to think carefully about DTensor (i.e. global SPMD) interops with default PyTorch (i.e. local SPMD).</p>
</li>
</ol>

username: wanchaol
cooked: <blockquote>
<p>maybe what I <em>actually</em> want is to use <code>redistribute</code> but then remove the DTensor wrapper immediately. I definitely want autograd supported inside this API. The idea would be that users can do the ‘replicate → partial’ transformation explicitly during forward (which is really a no-op)</p>
</blockquote>
<p>Hmmm but “replica → partial” transformation is not really a no-op in the forward…</p>
<p>It seems like you mainly want to expose some function that achieves the f/g operation in the megatron-lm style, and you want to reuse some of the DTensor components. IMO exposing sth like those with the dtensor concepts are not really easy and might bring more confusions to users (i.e. even you are confused in those cases).</p>
<p>I think the underlying problem here is: DTensor <strong>Does NOT</strong> have a one-to-one mapping to the Megatron <code>f/g</code> operations. The way that DTensor performs megatron TP achieves “functionally equivalence” with f/g, but DTensor does not simply formalizes the execution as two autograd function, let’s take the first column parallel linear as an example:</p>
<ul>
<li>Megatron <code>f</code> autograd function:
<ul>
<li>forward → no-op</li>
<li>backward → allreduce</li>
</ul>
</li>
<li>DTensor:
<ul>
<li>forward → replicated input @ column sharded weight: no communication</li>
<li>backward → columned sharded grad weight @ row sharded grad in: partial grad out, need allreduce</li>
</ul>
</li>
</ul>
<p>There is no single autograd function that can encapsulate the shardings rationally without explicitly specifying all of those shardings, specifically if you look at the details, why in the DTensor case the forward is no-op? It is because the operation it’s performing is a matmul and the input have the right shardings, so the <strong>no-op ties specifically to the matmul</strong> operation! Forcing it to be a single autograd function by specifying all of those shardings does not make sense, as you don’t know the following compute operation to know whether it should be a no-op or not. If we look closer into the MegatronLM paper, it actually formalizes the matmul sharding first then introduce those autograd functions, which also means f/g tights to the specific matmul sharding strategy.</p>
<p>This is why I think DTensor and collectives are two clearly separated level of abstractions, so I still don’t think the “redistribute but then remove the DTensor wrapper” is the right thing to do, at least not with the dtensor concepts. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=14" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

username: wanchaol
cooked: <p>Great insights about the implications of replicated parameters! One minor correction maybe:<br>
the last line of megatron strategy backwards should be</p>
<pre><code class="lang-auto">gradX: Sharded = fake_broadcast(grad_Z) # g backwards
</code></pre>
<p>It looks to me that the <code>fake_broadcast</code> and <code>fake_reduce</code> basically mimics the f/g operation, but it make things clear about why autograd behaves in that way using the per-device semantic (or in your term, local SPMD) mode, this is nice! One thing I wonder is that how to rationalize the “fake” in a way user could understand and whether/how to deal with rank0 broadcast/reduction as they could necessary bring down the efficiency?</p>
<p>I also agree that we should be careful about the behaviors between DTensor (implements “global SPMD”) and manual collectives with PyTorch (local SPMD). I do think as a higher level abstraction like DTensor, it hides certain details so exposing collective layer is important to users.</p>
<p>For distributed algorithms with PyTorch, what we have seen so far is mostly a state between a global SPMD DTensor and local SPMD autograd collectives: users actually implemented “global SPMD algorithms” using the PyTorch manual collectives (like Megatron TP). So for the next steps we are discussing, I think we should evaluate it together and answer the following questions:</p>
<ul>
<li>Which option is better? I personally perfer the third option
<ul>
<li>implement the proposed autograd formula for the current collectives (i.e. allreduce) by default and document the per-device semantics clearly</li>
<li>offer those autograd enabled collectives as separate APIs, still cleanly document the behaviors so that user can choose to use them as needed</li>
<li>keep the current state (don’t implement the autograd formulas for those collectives), but rather provide collective pairs like jax pbroadcast (the second proposal <a class="mention" href="/u/chillee">@Chillee</a> mentioned)</li>
<li>Do two or all of the above three</li>
</ul>
</li>
<li>Whether/How to encode sharded/replicated concepts to any of the proposal above? IMO directly exposing dtensor concepts might be confusing to users as I replied above.</li>
</ul>

username: tianyu
cooked: <p>sharing some of my thoughts as a note here</p>
<h2><a name="p-470228-autogradable-all-reduce-1" class="anchor" href="#p-470228-autogradable-all-reduce-1"></a>Autogradable all-reduce</h2>
<p>For simplicity of discussion, I’m restricting the scope of collectives in this note to all-reduce only. Other collectives should follow similar reasoning.</p>
<p><strong>In general, the backward of all-reduce is still all-reduce, but the bwd reduction op depends.</strong> Recall that the definition of gradient of an op is about “if input changes delta, how much output would change?”.</p>
<p>Let’s decompose all-reduce into mathematically equivalent reduce → broadcast, where the reduction op for reduce could be e.g. SUM, AVG, etc.</p>
<ul>
<li>The backward of reduce is
<ul>
<li>IDENTITY if the reduction op is SUM.</li>
<li>“MULTIPLY 1/world_size” if the reduction op if AVG.</li>
</ul>
</li>
<li>The backward of broadcast depends on how the fwd output is consumed / what loss function we use. E.g.
<ul>
<li>If we do .sum().backward() on the output of broadcast, the bwd should be SUM, or equivalently “MULTIPLY world_size”.</li>
</ul>
</li>
</ul>
<h2><a name="p-470228-tensor-parallel-2" class="anchor" href="#p-470228-tensor-parallel-2"></a>Tensor Parallel</h2>
<p>In a world where the only existing autogradable comm op is the general all-reduce, i.e. all-reduce(fwd) + all-reduce(bwd), assuming some flexibility in choosing the fwd/bwd reduction ops, we would perform TP as follows.</p>
<ul>
<li>fwd: → all-reduce with AVG (f) → TP → all-reduce with SUM (g) → Norm → </li>
<li>bwd: ← all-reduce with SUM (f) ← TP ← all-reduce with AVG (g) ← Norm ← </li>
</ul>
<p>The reason we need op g is because we need the activation to be Replicate (from Partial, in DTensor terms) to perform (Layer)Norm, not because Norm parameters are Replicate (cc Horace). Likewise, the reason we need op f is because we need the gradient to be Replicate (from Partial).</p>
<p>In Megatron TP, the forward of f and backward of g are no-ops. In fact, they essentially are. Because of the replicated computation of Norm, the occurrences of “all-reduce with AVG” can be omitted. Therefore, from such perspective, the f and g in Megatron TP are special cases of general autogradable all-reduce collectives. <strong>The specialty comes from the device-invariant region, without which there is no correctness guarantee for the no-op parts of Megatron-style f and g.</strong> In other words, the autogradable collectives of f and g are only valid with context / metadata.</p>
<h2><a name="p-470228-how-dtensor-does-it-3" class="anchor" href="#p-470228-how-dtensor-does-it-3"></a>How DTensor does it</h2>
<p>As one can see, we need some sort of special treatment to achieve Megatron TP, to avoid the wasteful general autogradable all-reduce. In the case of DTensor, it’s NOT done via autogradable ops; it is done via a metadata-enhanced tensor subclass. <strong>In DTensor, communications always happen as a means to change placement (an annotation of global distribution of data), e.g. Partial → Replicate, whether they are explicit via DTensor.redistribute, or implicit via sharding propagation strategies/rules.</strong> As Wanchao mentioned in <a href="https://discuss.pytorch.org/t/supporting-autograd-for-collectives/219430/5" class="inline-onebox">Supporting Autograd for Collectives - #5 by wanchaol</a>, this is functionally equivalent to the Megatron-style f and g, but with a different “philosophy”.</p>
<h2><a name="p-470228-to-summarize-4" class="anchor" href="#p-470228-to-summarize-4"></a>To summarize</h2>
<p><strong>DTensor cares primarily about data placement, and derives what collectives to perform based on source and target placements.</strong> Users don’t have explicit control over what collectives would be performed, and users probably don’t even care about them. DTensor engine should make it efficient, not the users themselves.</p>
<p><strong>Autogradable collectives, on the other hand, should focus on performing the desired collectives, ideally with correctness guarantee.</strong> E.g., whenever we use Megatron-style f in fwd, we should make sure it enters a device-invariant state. Depending on what users are used to, we may / may not need to support the derivations of correct DTensor placements after an arbitrary collective op.</p>
<p>In a nutshell, I think both are important</p>
<ul>
<li>DTensor provides ergonomic ways to deal with communication details, which could be particularly useful for users who are not experts in distributed.</li>
<li>Autogradable collectives provide low-level control which DTensor lacks support (and hopefully also correctness guarantee).</li>
</ul>
<p>But given how different and complementary they are, we probably should provide different levels of abstractions. I haven’t thought much on the interaction between the two approaches – I tentatively think they should be used disjointly.</p>