Repository: pytorch/torchtitan
Pull request #941: [DeepSeek] all-to-all-v kernel writes out `output_splits`
Author: login: kwen2501
Labels: name: CLA Signed
description: This label is managed by the Meta Open Source bot.
Created: 2025-03-07T00:51:02Z | Merged: 2025-03-20T06:07:13Z

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #956
* #989
* #958
* #954
* #952
* __->__ #941

This PR contains three changes:
    
    1. all-to-all-v kernel writes out `output_splits` instead of total received
    (so that we can use the split info as `input_splits` in the second shuffle or backward)
    
    2. Add backward for all_to_all_v (WIP)
      
    3. Move parts of the model out of `no_grad()`

ghstack-source-id: 8e0dda40b4dc9206343a020c42e8cd8c5741f860
Pull Request resolved: https://github.com/pytorch/torchtitan/pull/941
    
Test
```
$ torchrun --standalone --nproc-per-node 4 run.py
...
logits: torch.Size([2, 128, 102400])
logits: torch.Size([2, 128, 102400])
loss=tensor(510.3908, device='cuda:2', grad_fn=<MeanBackward0>)
loss=tensor(506.3376, device='cuda:3', grad_fn=<MeanBackward0>)
Backward done
Backward done
Backward done
Backward done
```

cc @yifuwang @ngimel


Diff:
diff --git a/torchtitan/experiments/deepseek_v3/model.py b/torchtitan/experiments/deepseek_v3/model.py
index 3ceb4bfa36..145cfd46e3 100644
--- a/torchtitan/experiments/deepseek_v3/model.py
+++ b/torchtitan/experiments/deepseek_v3/model.py
@@ -525,32 +525,43 @@ def forward(self, hidden_states):
         # for each token, select top-k experts, and compute the weight for each expert
         topk_idx, topk_weight = self.gate(hidden_states)
         hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-        y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape)
+        y = self.moe_forward(hidden_states, topk_idx, topk_weight).view(*orig_shape)
         if self.config.n_shared_experts is not None:
             y = y + self.shared_experts(identity)
         return y
 
-    @torch.no_grad()
-    def moe_infer(self, x, topk_ids, topk_weight):
+    def moe_forward(self, x, topk_ids, topk_weight):
         if self.ep_size > 1 and (not self.has_symm_mem):
             # Set up symmetric memory for the first time, then reuse it
             self.setup_symm_mem(x.dtype, x.device)
 
-        # [seq_len, n_routed_experts]
-        cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
-        # Fill 1 to the selected experts
-        cnts.scatter_(1, topk_ids, 1)
-        tokens_per_expert = cnts.sum(dim=0)
-        # Token indices for each expert
-        idxs = topk_ids.view(-1).argsort()
-        sorted_tokens_shape = idxs.shape + x.shape[1:]
+        # This part sorts the token indices so that tokens routed to the same expert reside consecutively.
+        # An implication is that tokens to the same "expert group" (i.e., device) are also consecutive.
+        # Since this is an "aritificial" index creation (final outcome being
+        # `idxs`), we don't need gradients here.
+        with torch.no_grad():
+            # [seq_len, n_routed_experts]
+            cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
+            # Fill 1 to the selected experts
+            cnts.scatter_(1, topk_ids, 1)
+            tokens_per_expert = cnts.sum(dim=0)
+            # Token indices for each expert
+            idxs = topk_ids.view(-1).argsort()
+            sorted_tokens_shape = idxs.shape + x.shape[1:]
+
         # Take necessary space from the `token_send_buf` symm mem
         sorted_tokens = self.token_send_buf[: idxs.shape[0]]
+        sorted_tokens.copy_(x[idxs // topk_ids.shape[1]])
+        # TODO: I tried to use `out=` to avoid copy, but it is not
+        # differentiable
+        # torch.index_select(x, 0, idxs // topk_ids.shape[1], out=sorted_tokens)
         assert sorted_tokens.shape == sorted_tokens_shape
-        # Use `out=` to avoid copy, equivalent to:
-        # `sorted_tokens = x[idxs // topk_ids.shape[1]]`
-        torch.index_select(x, 0, idxs // topk_ids.shape[1], out=sorted_tokens)
-        if self.ep_size > 1:
+
+        # This part exchange the information about the number of tokens send and
+        # received by each expert. We can understand this information as "side
+        # band", which is not part of the actual data. Thus no gradient is
+        # needed.
+        with torch.no_grad():
             # Sum the tokens over local experts, then we get tokens per EP rank,
             # which is the input splits
             torch.sum(
@@ -562,39 +573,44 @@ def moe_infer(self, x, topk_ids, topk_weight):
             dist.all_to_all_single(
                 tokens_per_expert_group, tokens_per_expert, group=self.ep_group
             )
-            torch.sum(
-                tokens_per_expert_group.view(self.ep_size, -1),
-                dim=1,
-                out=self.output_splits,
-            )
-            # print(
-            #     f"EP rank {self.ep_rank} receives {self.output_splits} tokens"
-            #     f"EP rank {self.ep_rank} sends {self.input_splits} tokens"
-            # )
-            # Total received tokens of current rank
-            received = torch.empty(1, dtype=torch.int64, device=x.device)
-            # DP to EP token shuffle
-            on_device_all_to_all_v(
-                self.token_gather_buf,
-                received,
-                sorted_tokens,
-                self.input_splits,
-                group=self.ep_group,
-            )
+
+        # DP to EP token shuffle. This part needs gradient.
+        on_device_all_to_all_v(
+            self.token_gather_buf,
+            self.output_splits,
+            sorted_tokens,
+            self.input_splits,
+            self.ep_group,
+        )
+
+        with torch.no_grad():
+            # Output splits sanity check
+            # expected_splits = tokens_per_expert_group.view(self.ep_size, -1).sum(dim=1)
+            # torch.testing.assert_close(self.output_splits, expected_splits)
             # Received tokens from all other ranks. TODO: use mask instead
-            gathered_tokens = self.token_gather_buf[:received]
+            received = self.output_splits.sum()
+
+        # TODO: don't use `received`
+        gathered_tokens = self.token_gather_buf[:received]
+
+        # This part prepares the token indices for each expert, because we
+        # receive them in a whole chunk for the entire expert group (i.e.,
+        # device). This part doesn't need gradient.
+        with torch.no_grad():
             tokens_per_expert_post_gather = tokens_per_expert_group.view(
                 self.ep_size, self.experts_per_rank
             ).sum(dim=0)
             gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0],), dtype=np.int32)
             s = 0
-            for i, k in enumerate(tokens_per_expert_group.cpu().numpy()):
+            # TODO: remove `tolist()`
+            for i, k in enumerate(tokens_per_expert_group.tolist()):
                 gatherd_idxs[s : s + k] = i % self.experts_per_rank
                 s += k
             gatherd_idxs = gatherd_idxs.argsort()
-            sorted_tokens = gathered_tokens[gatherd_idxs]
-            tokens_per_expert = tokens_per_expert_post_gather
-        tokens_per_expert = tokens_per_expert.cpu().numpy()
+            # TODO: remove `tolist()`
+            tokens_per_expert = tokens_per_expert_post_gather.tolist()
+
+        sorted_tokens = gathered_tokens[gatherd_idxs]
 
         outputs = []
         start_idx = 0
@@ -612,25 +628,24 @@ def moe_infer(self, x, topk_ids, topk_weight):
         # `sorted_tokens` would have shape [0, hidden_dim], we use it so that
         # `outs` is an empty tensor with shape [0, hidden_dim]
         outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens
-        if self.ep_size > 1:
-            # Take necessary space from `token_gather_buf` symm mem
-            new_x = self.token_gather_buf[: outs.shape[0]]
-            new_x[gatherd_idxs] = outs
-            gathered_tokens = new_x.new_empty(*sorted_tokens_shape)
-            # EP to DP token shuffle
-            on_device_all_to_all_v(
-                gathered_tokens,
-                received,  # unused
-                new_x,
-                self.output_splits,
-                group=self.ep_group,
-            )
-            outs = gathered_tokens
+        # Take necessary space from `token_gather_buf` symm mem
+        new_x = self.token_gather_buf[: outs.shape[0]]
+        new_x[gatherd_idxs] = outs
+        gathered_tokens = new_x.new_empty(*sorted_tokens_shape)
+        received_splits = torch.empty_like(self.output_splits)  # unused
+        # EP to DP token shuffle
+        on_device_all_to_all_v(
+            gathered_tokens,
+            received_splits,  # unused
+            new_x,
+            self.output_splits,
+            self.ep_group,
+        )
 
-        new_x = torch.empty_like(outs)
-        new_x[idxs] = outs
+        returned_tokens = torch.empty_like(gathered_tokens)
+        returned_tokens[idxs] = gathered_tokens
         final_out = (
-            new_x.view(*topk_ids.shape, -1)
+            returned_tokens.view(*topk_ids.shape, -1)
             .type(topk_weight.dtype)
             .mul_(topk_weight.unsqueeze(dim=-1))
             .sum(dim=1)
diff --git a/torchtitan/experiments/deepseek_v3/symm_mem_recipes/triton_on_device_all_to_all_v.py b/torchtitan/experiments/deepseek_v3/symm_mem_recipes/triton_on_device_all_to_all_v.py
index 5174ea85eb..ba58a656d7 100644
--- a/torchtitan/experiments/deepseek_v3/symm_mem_recipes/triton_on_device_all_to_all_v.py
+++ b/torchtitan/experiments/deepseek_v3/symm_mem_recipes/triton_on_device_all_to_all_v.py
@@ -55,9 +55,9 @@ def _exchange_row_offsets(
 @triton.jit
 def on_device_all_to_all_v_kernel(
     output_ptr,
-    received_ptr,
+    output_splits_ptr,
     input_ptrs,
-    split_sizes_ptrs,
+    input_splits_ptr,
     signal_pad_ptrs,
     dim: tl.constexpr,  # Separate dim for easier vectorization
     rank: tl.constexpr,
@@ -73,11 +73,13 @@ def on_device_all_to_all_v_kernel(
     block_offset = tl.program_id(0) % BLOCKS_PER_REMOTE_RANK
 
     input_row_offset, output_row_offset, num_rows = _exchange_row_offsets(
-        split_sizes_ptrs, rank, world_size, BLOCKS_PER_REMOTE_RANK
+        input_splits_ptr, rank, world_size, BLOCKS_PER_REMOTE_RANK
     )
 
-    if tl.program_id(0) == tl.num_programs(0) - 1:
-        tl.store(received_ptr, output_row_offset + num_rows)
+    output_splits_ptr = output_splits_ptr.to(tl.pointer_type(tl.uint64))
+    if block_offset == 0:
+        # Update output_splits
+        tl.store(output_splits_ptr + remote_rank, num_rows)
 
     input_ptr = (
         tl.load(input_ptrs.to(tl.pointer_type(tl.uint64)) + remote_rank).to(
@@ -121,15 +123,15 @@ def on_device_all_to_all_v_kernel(
     return
 
 
-def on_device_all_to_all_v(
+def _on_device_all_to_all_v(
     output: torch.Tensor,
-    received: torch.Tensor,
+    output_splits: torch.Tensor,
     input: torch.Tensor,
-    split_sizes: torch.Tensor,
+    input_splits: torch.Tensor,
+    group: dist.ProcessGroup = dist.group.WORLD,
     BLOCKS_PER_REMOTE_RANK=8,
     UNROLL_FACTOR: int = 8,
     BLOCK_SIZE: int = 16384,
-    group: dist.ProcessGroup = dist.group.WORLD,
 ):
     assert output.dim() == 2, f"{output.shape}"
     assert input.dim() == 2, f"{input.shape}"
@@ -137,14 +139,14 @@ def on_device_all_to_all_v(
 
     dim = output.shape[1]
     input_hdl = symm_mem.rendezvous(input, group=group)
-    split_sizes_hdl = symm_mem.rendezvous(split_sizes, group=group)
+    input_splits_hdl = symm_mem.rendezvous(input_splits, group=group)
 
     num_blocks = input_hdl.world_size * BLOCKS_PER_REMOTE_RANK
     kernel = on_device_all_to_all_v_kernel[(num_blocks, 1, 1)](
         output,
-        received,
+        output_splits,
         input_hdl.buffer_ptrs_dev,
-        split_sizes_hdl.buffer_ptrs_dev,
+        input_splits_hdl.buffer_ptrs_dev,
         input_hdl.signal_pad_ptrs_dev,
         dim=dim,
         rank=input_hdl.rank,
@@ -156,3 +158,44 @@ def on_device_all_to_all_v(
     )
     # log_triton_kernel(kernel)
     return output
+
+
+class OnDeviceAllToAllV(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        output: torch.Tensor,
+        output_splits: torch.Tensor,
+        input: torch.Tensor,
+        input_splits: torch.Tensor,
+        group: dist.ProcessGroup = dist.group.WORLD,
+    ):
+        _on_device_all_to_all_v(output, output_splits, input, input_splits, group=group)
+        ctx.save_for_backward(output_splits)
+        ctx.group = group
+        ctx.input_shape = input.shape
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        # TODO: autograd requires tensors not be modified a second time, this
+        # conflicts with our wish of sharing the symm mem across layers and/or
+        # PP microbatches.
+        return NotImplementedError(
+            "OnDeviceAllToAllV backward is not ready, please use it for inference only"
+        )
+        grad_output_splits = ctx.saved_tensors
+        grad_input_splits = torch.empty_like(grad_output_splits)
+        grad_input = grad_output.new_empty(*ctx.input_shape)
+        _on_device_all_to_all_v(
+            grad_input,
+            grad_input_splits,
+            grad_output,
+            grad_output_splits,
+            group=ctx.group,
+        )
+        return None, None, grad_input, None, None
+
+
+# Alias
+on_device_all_to_all_v = OnDeviceAllToAllV.apply


Reviews:
user: login: ngimel
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-07T02:10:57Z
user: login: kwen2501
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-07T02:21:34Z
user: login: ngimel
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-07T17:38:34Z
user: login: wconstab
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T01:03:17Z
user: login: wconstab
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T01:07:42Z
user: login: wconstab
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T01:18:17Z
user: login: wconstab
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T01:21:20Z
user: login: wconstab
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T01:30:31Z
user: login: wconstab
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T01:33:36Z
user: login: EugenHotaj
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T04:07:03Z
user: login: wconstab
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T23:47:42Z
user: login: wconstab
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T23:51:58Z
user: login: wconstab
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T00:00:39Z
user: login: wconstab
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T00:02:57Z
user: login: kwen2501
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-20T00:54:49Z
user: login: kwen2501
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-20T02:20:46Z
user: login: kwen2501
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-20T02:21:36Z
user: login: wconstab
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-20T03:38:49Z
user: login: wconstab
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-20T03:39:36Z
user: login: ngimel
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-20T16:17:31Z
user: login: kwen2501
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-20T17:20:03Z

Review comments:
diff_hunk: @@ -562,26 +562,22 @@ def moe_infer(self, x, topk_ids, topk_weight):
             dist.all_to_all_single(
                 tokens_per_expert_group, tokens_per_expert, group=self.ep_group
             )
-            torch.sum(
-                tokens_per_expert_group.view(self.ep_size, -1),
-                dim=1,
-                out=self.output_splits,
-            )
-            # print(
-            #     f"EP rank {self.ep_rank} receives {self.output_splits} tokens"
-            #     f"EP rank {self.ep_rank} sends {self.input_splits} tokens"
-            # )
-            # Total received tokens of current rank
-            received = torch.empty(1, dtype=torch.int64, device=x.device)
             # DP to EP token shuffle
             on_device_all_to_all_v(
                 self.token_gather_buf,
-                received,
+                self.output_splits,
                 sorted_tokens,
                 self.input_splits,
                 group=self.ep_group,
             )
+            # Output splits sanity check
+            # expected_splits = torch.sum(
+            #     tokens_per_expert_group.view(self.ep_size, -1),
+            #     dim=1,
+            # )
+            # torch.testing.assert_close(self.output_splits, expected_splits)
             # Received tokens from all other ranks. TODO: use mask instead
+            received = self.output_splits.sum().item()
path: torchtitan/experiments/deepseek_v3/model.py
user: login: ngimel
body: this is an h2d sync?
created_at: 2025-03-07T02:10:57Z
author_association: NONE
diff_hunk: @@ -562,26 +562,22 @@ def moe_infer(self, x, topk_ids, topk_weight):
             dist.all_to_all_single(
                 tokens_per_expert_group, tokens_per_expert, group=self.ep_group
             )
-            torch.sum(
-                tokens_per_expert_group.view(self.ep_size, -1),
-                dim=1,
-                out=self.output_splits,
-            )
-            # print(
-            #     f"EP rank {self.ep_rank} receives {self.output_splits} tokens"
-            #     f"EP rank {self.ep_rank} sends {self.input_splits} tokens"
-            # )
-            # Total received tokens of current rank
-            received = torch.empty(1, dtype=torch.int64, device=x.device)
             # DP to EP token shuffle
             on_device_all_to_all_v(
                 self.token_gather_buf,
-                received,
+                self.output_splits,
                 sorted_tokens,
                 self.input_splits,
                 group=self.ep_group,
             )
+            # Output splits sanity check
+            # expected_splits = torch.sum(
+            #     tokens_per_expert_group.view(self.ep_size, -1),
+            #     dim=1,
+            # )
+            # torch.testing.assert_close(self.output_splits, expected_splits)
             # Received tokens from all other ranks. TODO: use mask instead
+            received = self.output_splits.sum().item()
path: torchtitan/experiments/deepseek_v3/model.py
user: login: kwen2501
body: It is. I am looking for an mechanism that can consume the `output_splits` directly.
Would your PR be related? https://github.com/pytorch/pytorch/pull/148531
Without that mechanism, I am putting a sync here for now.
created_at: 2025-03-07T02:21:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -562,26 +562,22 @@ def moe_infer(self, x, topk_ids, topk_weight):
             dist.all_to_all_single(
                 tokens_per_expert_group, tokens_per_expert, group=self.ep_group
             )
-            torch.sum(
-                tokens_per_expert_group.view(self.ep_size, -1),
-                dim=1,
-                out=self.output_splits,
-            )
-            # print(
-            #     f"EP rank {self.ep_rank} receives {self.output_splits} tokens"
-            #     f"EP rank {self.ep_rank} sends {self.input_splits} tokens"
-            # )
-            # Total received tokens of current rank
-            received = torch.empty(1, dtype=torch.int64, device=x.device)
             # DP to EP token shuffle
             on_device_all_to_all_v(
                 self.token_gather_buf,
-                received,
+                self.output_splits,
                 sorted_tokens,
                 self.input_splits,
                 group=self.ep_group,
             )
+            # Output splits sanity check
+            # expected_splits = torch.sum(
+            #     tokens_per_expert_group.view(self.ep_size, -1),
+            #     dim=1,
+            # )
+            # torch.testing.assert_close(self.output_splits, expected_splits)
             # Received tokens from all other ranks. TODO: use mask instead
+            received = self.output_splits.sum().item()
path: torchtitan/experiments/deepseek_v3/model.py
user: login: ngimel
body: It's related but it comes with "all splits should be divisible by 16" limitation, for inference maybe this grouped gemm https://github.com/pytorch/FBGEMM/blob/98e9e680c664421a6157d55079e58c7938452798/fbgemm_gpu/experimental/gemm/triton_gemm/grouped_gemm.py#L131 is better. 
created_at: 2025-03-07T17:38:34Z
author_association: NONE
diff_hunk: @@ -525,32 +525,43 @@ def forward(self, hidden_states):
         # for each token, select top-k experts, and compute the weight for each expert
         topk_idx, topk_weight = self.gate(hidden_states)
         hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-        y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape)
+        y = self.moe_forward(hidden_states, topk_idx, topk_weight).view(*orig_shape)
         if self.config.n_shared_experts is not None:
             y = y + self.shared_experts(identity)
         return y
 
-    @torch.no_grad()
-    def moe_infer(self, x, topk_ids, topk_weight):
+    def moe_forward(self, x, topk_ids, topk_weight):
path: torchtitan/experiments/deepseek_v3/model.py
user: login: wconstab
body: would be nice to add docstring to moe_forward (well, probably to other methods too) that document the shapes of the input tensors, e.g. `topk_ids: [seq_len, k]`
created_at: 2025-03-18T01:03:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -525,32 +525,43 @@ def forward(self, hidden_states):
         # for each token, select top-k experts, and compute the weight for each expert
         topk_idx, topk_weight = self.gate(hidden_states)
         hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-        y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape)
+        y = self.moe_forward(hidden_states, topk_idx, topk_weight).view(*orig_shape)
         if self.config.n_shared_experts is not None:
             y = y + self.shared_experts(identity)
         return y
 
-    @torch.no_grad()
-    def moe_infer(self, x, topk_ids, topk_weight):
+    def moe_forward(self, x, topk_ids, topk_weight):
         if self.ep_size > 1 and (not self.has_symm_mem):
             # Set up symmetric memory for the first time, then reuse it
             self.setup_symm_mem(x.dtype, x.device)
 
-        # [seq_len, n_routed_experts]
-        cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
-        # Fill 1 to the selected experts
-        cnts.scatter_(1, topk_ids, 1)
-        tokens_per_expert = cnts.sum(dim=0)
-        # Token indices for each expert
-        idxs = topk_ids.view(-1).argsort()
-        sorted_tokens_shape = idxs.shape + x.shape[1:]
+        # This part sorts the token indices so that tokens routed to the same expert reside consecutively.
+        # An implication is that tokens to the same "expert group" (i.e., device) are also consecutive.
+        # Since this is an "aritificial" index creation (final outcome being
+        # `idxs`), we don't need gradients here.
+        with torch.no_grad():
+            # [seq_len, n_routed_experts]
path: torchtitan/experiments/deepseek_v3/model.py
user: login: wconstab
body: maybe collaps the 3 lines (cnts=, cnts.scatter_, tokens_per_expert) into a group of 3 and give one comment above them that says 
`# form a one_hot encoding of shape (seq_len, experts) and then sum to get #tokens per expert`
and maybe rename `tokens_per_expert` to `num_tokens_per_expert` or something?
created_at: 2025-03-18T01:07:42Z
author_association: CONTRIBUTOR
diff_hunk: @@ -525,32 +525,43 @@ def forward(self, hidden_states):
         # for each token, select top-k experts, and compute the weight for each expert
         topk_idx, topk_weight = self.gate(hidden_states)
         hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-        y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape)
+        y = self.moe_forward(hidden_states, topk_idx, topk_weight).view(*orig_shape)
         if self.config.n_shared_experts is not None:
             y = y + self.shared_experts(identity)
         return y
 
-    @torch.no_grad()
-    def moe_infer(self, x, topk_ids, topk_weight):
+    def moe_forward(self, x, topk_ids, topk_weight):
         if self.ep_size > 1 and (not self.has_symm_mem):
             # Set up symmetric memory for the first time, then reuse it
             self.setup_symm_mem(x.dtype, x.device)
 
-        # [seq_len, n_routed_experts]
-        cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
-        # Fill 1 to the selected experts
-        cnts.scatter_(1, topk_ids, 1)
-        tokens_per_expert = cnts.sum(dim=0)
-        # Token indices for each expert
-        idxs = topk_ids.view(-1).argsort()
-        sorted_tokens_shape = idxs.shape + x.shape[1:]
+        # This part sorts the token indices so that tokens routed to the same expert reside consecutively.
+        # An implication is that tokens to the same "expert group" (i.e., device) are also consecutive.
+        # Since this is an "aritificial" index creation (final outcome being
+        # `idxs`), we don't need gradients here.
+        with torch.no_grad():
+            # [seq_len, n_routed_experts]
+            cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
+            # Fill 1 to the selected experts
+            cnts.scatter_(1, topk_ids, 1)
+            tokens_per_expert = cnts.sum(dim=0)
+            # Token indices for each expert
+            idxs = topk_ids.view(-1).argsort()
+            sorted_tokens_shape = idxs.shape + x.shape[1:]
path: torchtitan/experiments/deepseek_v3/model.py
user: login: wconstab
body: so we flatten the idxs (each sequence is appended to one long row), then we shuffle them so that tokens from different sequences that are going to the same expert are next to each other in memory.

and x.shape is what, (b*s,h)? if so, then i think it makes sense that sorted_tokens_shape would be (b*s, h)
created_at: 2025-03-18T01:18:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -525,32 +525,43 @@ def forward(self, hidden_states):
         # for each token, select top-k experts, and compute the weight for each expert
         topk_idx, topk_weight = self.gate(hidden_states)
         hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-        y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape)
+        y = self.moe_forward(hidden_states, topk_idx, topk_weight).view(*orig_shape)
         if self.config.n_shared_experts is not None:
             y = y + self.shared_experts(identity)
         return y
 
-    @torch.no_grad()
-    def moe_infer(self, x, topk_ids, topk_weight):
+    def moe_forward(self, x, topk_ids, topk_weight):
         if self.ep_size > 1 and (not self.has_symm_mem):
             # Set up symmetric memory for the first time, then reuse it
             self.setup_symm_mem(x.dtype, x.device)
 
-        # [seq_len, n_routed_experts]
-        cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
-        # Fill 1 to the selected experts
-        cnts.scatter_(1, topk_ids, 1)
-        tokens_per_expert = cnts.sum(dim=0)
-        # Token indices for each expert
-        idxs = topk_ids.view(-1).argsort()
-        sorted_tokens_shape = idxs.shape + x.shape[1:]
+        # This part sorts the token indices so that tokens routed to the same expert reside consecutively.
+        # An implication is that tokens to the same "expert group" (i.e., device) are also consecutive.
+        # Since this is an "aritificial" index creation (final outcome being
+        # `idxs`), we don't need gradients here.
+        with torch.no_grad():
+            # [seq_len, n_routed_experts]
+            cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
+            # Fill 1 to the selected experts
+            cnts.scatter_(1, topk_ids, 1)
+            tokens_per_expert = cnts.sum(dim=0)
+            # Token indices for each expert
+            idxs = topk_ids.view(-1).argsort()
+            sorted_tokens_shape = idxs.shape + x.shape[1:]
+
         # Take necessary space from the `token_send_buf` symm mem
         sorted_tokens = self.token_send_buf[: idxs.shape[0]]
path: torchtitan/experiments/deepseek_v3/model.py
user: login: wconstab
body: nit: is `[: idxs.shape[0]]` broken formatting? i thought it would be :idxs... (no space) - otherwise looked like it was 2 dim and missing comma.
created_at: 2025-03-18T01:21:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -525,32 +525,43 @@ def forward(self, hidden_states):
         # for each token, select top-k experts, and compute the weight for each expert
         topk_idx, topk_weight = self.gate(hidden_states)
         hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-        y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape)
+        y = self.moe_forward(hidden_states, topk_idx, topk_weight).view(*orig_shape)
         if self.config.n_shared_experts is not None:
             y = y + self.shared_experts(identity)
         return y
 
-    @torch.no_grad()
-    def moe_infer(self, x, topk_ids, topk_weight):
+    def moe_forward(self, x, topk_ids, topk_weight):
         if self.ep_size > 1 and (not self.has_symm_mem):
             # Set up symmetric memory for the first time, then reuse it
             self.setup_symm_mem(x.dtype, x.device)
 
-        # [seq_len, n_routed_experts]
-        cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
-        # Fill 1 to the selected experts
-        cnts.scatter_(1, topk_ids, 1)
-        tokens_per_expert = cnts.sum(dim=0)
-        # Token indices for each expert
-        idxs = topk_ids.view(-1).argsort()
-        sorted_tokens_shape = idxs.shape + x.shape[1:]
+        # This part sorts the token indices so that tokens routed to the same expert reside consecutively.
+        # An implication is that tokens to the same "expert group" (i.e., device) are also consecutive.
+        # Since this is an "aritificial" index creation (final outcome being
+        # `idxs`), we don't need gradients here.
+        with torch.no_grad():
+            # [seq_len, n_routed_experts]
+            cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
+            # Fill 1 to the selected experts
+            cnts.scatter_(1, topk_ids, 1)
+            tokens_per_expert = cnts.sum(dim=0)
+            # Token indices for each expert
+            idxs = topk_ids.view(-1).argsort()
+            sorted_tokens_shape = idxs.shape + x.shape[1:]
+
         # Take necessary space from the `token_send_buf` symm mem
         sorted_tokens = self.token_send_buf[: idxs.shape[0]]
+        sorted_tokens.copy_(x[idxs // topk_ids.shape[1]])
path: torchtitan/experiments/deepseek_v3/model.py
user: login: wconstab
body: x: [b*s, h]?
idxs: [b*s]?
topk_ids: [s, k]

idxs represents "the indices of tokens respective to b*s flattening, such that the tokens for expert0 come first, etc", right?

i'm confused about what `idxs // topk_ids.shape[1]` is doing. 
created_at: 2025-03-18T01:30:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -156,3 +158,34 @@ def on_device_all_to_all_v(
     )
     # log_triton_kernel(kernel)
     return output
+
+
+class OnDeviceAllToAllV(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        output: torch.Tensor,
+        output_splits: torch.Tensor,
+        input: torch.Tensor,
+        input_splits: torch.Tensor,
+        group: dist.ProcessGroup = dist.group.WORLD,
path: torchtitan/experiments/deepseek_v3/symm_mem_recipes/triton_on_device_all_to_all_v.py
user: login: wconstab
body: should we default this? it feels like usually you need to know what group you're doing this over. i could see leaving it as 'world' or unit testing purposes but for any real usage in models it seems unlikely the default would be sufficient?
created_at: 2025-03-18T01:33:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -156,3 +158,34 @@ def on_device_all_to_all_v(
     )
     # log_triton_kernel(kernel)
     return output
+
+
+class OnDeviceAllToAllV(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        output: torch.Tensor,
+        output_splits: torch.Tensor,
+        input: torch.Tensor,
+        input_splits: torch.Tensor,
+        group: dist.ProcessGroup = dist.group.WORLD,
+    ):
+        _on_device_all_to_all_v(output, output_splits, input, input_splits, group=group)
+        ctx.save_for_backward(output_splits)
+        ctx.group = group
+        ctx.input_shape = input.shape
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_output_splits = ctx.saved_tensors
+        grad_input_splits = torch.empty_like(grad_output_splits)
+        grad_input = grad_output.new_empty(*ctx.input_shape)
+        _on_device_all_to_all_v(
+            grad_input, grad_input_splits, grad_output, grad_output_splits, group=group
path: torchtitan/experiments/deepseek_v3/symm_mem_recipes/triton_on_device_all_to_all_v.py
user: login: EugenHotaj
body: I might be missing something but I don't think this `backward()` function is ever being called -- otherwise `group` would fail here since it should be `ctx.group`
created_at: 2025-03-18T04:07:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -525,32 +525,43 @@ def forward(self, hidden_states):
         # for each token, select top-k experts, and compute the weight for each expert
         topk_idx, topk_weight = self.gate(hidden_states)
         hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-        y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape)
+        y = self.moe_forward(hidden_states, topk_idx, topk_weight).view(*orig_shape)
         if self.config.n_shared_experts is not None:
             y = y + self.shared_experts(identity)
         return y
 
-    @torch.no_grad()
-    def moe_infer(self, x, topk_ids, topk_weight):
+    def moe_forward(self, x, topk_ids, topk_weight):
         if self.ep_size > 1 and (not self.has_symm_mem):
             # Set up symmetric memory for the first time, then reuse it
             self.setup_symm_mem(x.dtype, x.device)
 
-        # [seq_len, n_routed_experts]
-        cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
-        # Fill 1 to the selected experts
-        cnts.scatter_(1, topk_ids, 1)
-        tokens_per_expert = cnts.sum(dim=0)
-        # Token indices for each expert
-        idxs = topk_ids.view(-1).argsort()
-        sorted_tokens_shape = idxs.shape + x.shape[1:]
+        # This part sorts the token indices so that tokens routed to the same expert reside consecutively.
+        # An implication is that tokens to the same "expert group" (i.e., device) are also consecutive.
+        # Since this is an "aritificial" index creation (final outcome being
+        # `idxs`), we don't need gradients here.
+        with torch.no_grad():
+            # [seq_len, n_routed_experts]
+            cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
+            # Fill 1 to the selected experts
+            cnts.scatter_(1, topk_ids, 1)
+            tokens_per_expert = cnts.sum(dim=0)
+            # Token indices for each expert
+            idxs = topk_ids.view(-1).argsort()
+            sorted_tokens_shape = idxs.shape + x.shape[1:]
+
         # Take necessary space from the `token_send_buf` symm mem
         sorted_tokens = self.token_send_buf[: idxs.shape[0]]
+        sorted_tokens.copy_(x[idxs // topk_ids.shape[1]])
path: torchtitan/experiments/deepseek_v3/model.py
user: login: wconstab
body: ok- the explanation is
1) idxs contains the indexes of tokens in a sequence that is K times longer than the original input sequence, due to the view(-1) above before the argsort
2) the index values logically refer to tokens in the original input sequence (x) so to make them useful again we floor-didide by `k` (the value used for topk).
created_at: 2025-03-18T23:47:42Z
author_association: CONTRIBUTOR
diff_hunk: @@ -562,39 +573,44 @@ def moe_infer(self, x, topk_ids, topk_weight):
             dist.all_to_all_single(
                 tokens_per_expert_group, tokens_per_expert, group=self.ep_group
             )
-            torch.sum(
-                tokens_per_expert_group.view(self.ep_size, -1),
-                dim=1,
-                out=self.output_splits,
-            )
-            # print(
-            #     f"EP rank {self.ep_rank} receives {self.output_splits} tokens"
-            #     f"EP rank {self.ep_rank} sends {self.input_splits} tokens"
-            # )
-            # Total received tokens of current rank
-            received = torch.empty(1, dtype=torch.int64, device=x.device)
-            # DP to EP token shuffle
-            on_device_all_to_all_v(
-                self.token_gather_buf,
-                received,
-                sorted_tokens,
-                self.input_splits,
-                group=self.ep_group,
-            )
+
+        # DP to EP token shuffle. This part needs gradient.
+        on_device_all_to_all_v(
+            self.token_gather_buf,
+            self.output_splits,
+            sorted_tokens,
+            self.input_splits,
+            self.ep_group,
+        )
+
+        with torch.no_grad():
+            # Output splits sanity check
+            # expected_splits = tokens_per_expert_group.view(self.ep_size, -1).sum(dim=1)
+            # torch.testing.assert_close(self.output_splits, expected_splits)
             # Received tokens from all other ranks. TODO: use mask instead
-            gathered_tokens = self.token_gather_buf[:received]
+            received = self.output_splits.sum()
+
+        # TODO: don't use `received`
+        gathered_tokens = self.token_gather_buf[:received]
+
+        # This part prepares the token indices for each expert, because we
+        # receive them in a whole chunk for the entire expert group (i.e.,
+        # device). This part doesn't need gradient.
+        with torch.no_grad():
             tokens_per_expert_post_gather = tokens_per_expert_group.view(
                 self.ep_size, self.experts_per_rank
             ).sum(dim=0)
             gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0],), dtype=np.int32)
             s = 0
-            for i, k in enumerate(tokens_per_expert_group.cpu().numpy()):
+            # TODO: remove `tolist()`
+            for i, k in enumerate(tokens_per_expert_group.tolist()):
path: torchtitan/experiments/deepseek_v3/model.py
user: login: wconstab
body: this is still a cuda-cpu sync point, right? can we avoid it by doing a scan on device or something?
created_at: 2025-03-18T23:51:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -612,25 +628,24 @@ def moe_infer(self, x, topk_ids, topk_weight):
         # `sorted_tokens` would have shape [0, hidden_dim], we use it so that
         # `outs` is an empty tensor with shape [0, hidden_dim]
         outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens
-        if self.ep_size > 1:
-            # Take necessary space from `token_gather_buf` symm mem
-            new_x = self.token_gather_buf[: outs.shape[0]]
-            new_x[gatherd_idxs] = outs
-            gathered_tokens = new_x.new_empty(*sorted_tokens_shape)
-            # EP to DP token shuffle
-            on_device_all_to_all_v(
-                gathered_tokens,
-                received,  # unused
-                new_x,
-                self.output_splits,
-                group=self.ep_group,
-            )
-            outs = gathered_tokens
+        # Take necessary space from `token_gather_buf` symm mem
+        new_x = self.token_gather_buf[: outs.shape[0]]
+        new_x[gatherd_idxs] = outs
+        gathered_tokens = new_x.new_empty(*sorted_tokens_shape)
+        received_splits = torch.empty_like(self.output_splits)  # unused
+        # EP to DP token shuffle
+        on_device_all_to_all_v(
+            gathered_tokens,
+            received_splits,  # unused
+            new_x,
+            self.output_splits,
+            self.ep_group,
+        )
 
-        new_x = torch.empty_like(outs)
-        new_x[idxs] = outs
+        returned_tokens = torch.empty_like(gathered_tokens)
+        returned_tokens[idxs] = gathered_tokens
         final_out = (
-            new_x.view(*topk_ids.shape, -1)
+            returned_tokens.view(*topk_ids.shape, -1)
path: torchtitan/experiments/deepseek_v3/model.py
user: login: wconstab
body: i'd also like to see `topk_weight` documented at the top of the function.  I am assuming that this helps to implement a weighted sum for combining the expert outputs into one output, but its a little hard to follow the code.

created_at: 2025-03-19T00:00:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -156,3 +158,34 @@ def on_device_all_to_all_v(
     )
     # log_triton_kernel(kernel)
     return output
+
+
+class OnDeviceAllToAllV(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        output: torch.Tensor,
+        output_splits: torch.Tensor,
+        input: torch.Tensor,
+        input_splits: torch.Tensor,
+        group: dist.ProcessGroup = dist.group.WORLD,
+    ):
+        _on_device_all_to_all_v(output, output_splits, input, input_splits, group=group)
+        ctx.save_for_backward(output_splits)
+        ctx.group = group
+        ctx.input_shape = input.shape
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_output_splits = ctx.saved_tensors
+        grad_input_splits = torch.empty_like(grad_output_splits)
+        grad_input = grad_output.new_empty(*ctx.input_shape)
+        _on_device_all_to_all_v(
+            grad_input, grad_input_splits, grad_output, grad_output_splits, group=group
path: torchtitan/experiments/deepseek_v3/symm_mem_recipes/triton_on_device_all_to_all_v.py
user: login: wconstab
body: yea good point. but it does look like the on_device all2allv is called in forward part of moe layer in a part that has gradients enabled.  (it is also called under 'with no_grad()` but not exclusively)
created_at: 2025-03-19T00:02:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -156,3 +158,34 @@ def on_device_all_to_all_v(
     )
     # log_triton_kernel(kernel)
     return output
+
+
+class OnDeviceAllToAllV(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        output: torch.Tensor,
+        output_splits: torch.Tensor,
+        input: torch.Tensor,
+        input_splits: torch.Tensor,
+        group: dist.ProcessGroup = dist.group.WORLD,
+    ):
+        _on_device_all_to_all_v(output, output_splits, input, input_splits, group=group)
+        ctx.save_for_backward(output_splits)
+        ctx.group = group
+        ctx.input_shape = input.shape
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_output_splits = ctx.saved_tensors
+        grad_input_splits = torch.empty_like(grad_output_splits)
+        grad_input = grad_output.new_empty(*ctx.input_shape)
+        _on_device_all_to_all_v(
+            grad_input, grad_input_splits, grad_output, grad_output_splits, group=group
path: torchtitan/experiments/deepseek_v3/symm_mem_recipes/triton_on_device_all_to_all_v.py
user: login: kwen2501
body: @EugenHotaj @wconstab Thanks. You are right. 
I am backing out the backward function for `on_device_all_to_all_v` from this PR.
The reason is that autograd does not allow touching the buffer a second time between fwd and bwd. This conflicts with sharing the symm mem buffer across layers (to save memory) as in PR #958.

Instead, I am adding backward support via the traditional `all_to_all_single_autograd` (Yes, there is a CPU sync.) But the backward is more stable. PR is at #989.

`on_device_all_to_all_v` is now limited to inference / generation.
created_at: 2025-03-20T00:54:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -525,32 +525,43 @@ def forward(self, hidden_states):
         # for each token, select top-k experts, and compute the weight for each expert
         topk_idx, topk_weight = self.gate(hidden_states)
         hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-        y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape)
+        y = self.moe_forward(hidden_states, topk_idx, topk_weight).view(*orig_shape)
         if self.config.n_shared_experts is not None:
             y = y + self.shared_experts(identity)
         return y
 
-    @torch.no_grad()
-    def moe_infer(self, x, topk_ids, topk_weight):
+    def moe_forward(self, x, topk_ids, topk_weight):
         if self.ep_size > 1 and (not self.has_symm_mem):
             # Set up symmetric memory for the first time, then reuse it
             self.setup_symm_mem(x.dtype, x.device)
 
-        # [seq_len, n_routed_experts]
-        cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
-        # Fill 1 to the selected experts
-        cnts.scatter_(1, topk_ids, 1)
-        tokens_per_expert = cnts.sum(dim=0)
-        # Token indices for each expert
-        idxs = topk_ids.view(-1).argsort()
-        sorted_tokens_shape = idxs.shape + x.shape[1:]
+        # This part sorts the token indices so that tokens routed to the same expert reside consecutively.
+        # An implication is that tokens to the same "expert group" (i.e., device) are also consecutive.
+        # Since this is an "aritificial" index creation (final outcome being
+        # `idxs`), we don't need gradients here.
+        with torch.no_grad():
+            # [seq_len, n_routed_experts]
+            cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
+            # Fill 1 to the selected experts
+            cnts.scatter_(1, topk_ids, 1)
+            tokens_per_expert = cnts.sum(dim=0)
+            # Token indices for each expert
+            idxs = topk_ids.view(-1).argsort()
+            sorted_tokens_shape = idxs.shape + x.shape[1:]
+
         # Take necessary space from the `token_send_buf` symm mem
         sorted_tokens = self.token_send_buf[: idxs.shape[0]]
path: torchtitan/experiments/deepseek_v3/model.py
user: login: kwen2501
body: I remember the space is required by formatter too.
created_at: 2025-03-20T02:20:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -156,3 +158,34 @@ def on_device_all_to_all_v(
     )
     # log_triton_kernel(kernel)
     return output
+
+
+class OnDeviceAllToAllV(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        output: torch.Tensor,
+        output_splits: torch.Tensor,
+        input: torch.Tensor,
+        input_splits: torch.Tensor,
+        group: dist.ProcessGroup = dist.group.WORLD,
path: torchtitan/experiments/deepseek_v3/symm_mem_recipes/triton_on_device_all_to_all_v.py
user: login: kwen2501
body: The default here is similar to torch.dist defaulting group to `None`
created_at: 2025-03-20T02:21:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -156,3 +158,34 @@ def on_device_all_to_all_v(
     )
     # log_triton_kernel(kernel)
     return output
+
+
+class OnDeviceAllToAllV(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        output: torch.Tensor,
+        output_splits: torch.Tensor,
+        input: torch.Tensor,
+        input_splits: torch.Tensor,
+        group: dist.ProcessGroup = dist.group.WORLD,
+    ):
+        _on_device_all_to_all_v(output, output_splits, input, input_splits, group=group)
+        ctx.save_for_backward(output_splits)
+        ctx.group = group
+        ctx.input_shape = input.shape
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_output_splits = ctx.saved_tensors
+        grad_input_splits = torch.empty_like(grad_output_splits)
+        grad_input = grad_output.new_empty(*ctx.input_shape)
+        _on_device_all_to_all_v(
+            grad_input, grad_input_splits, grad_output, grad_output_splits, group=group
path: torchtitan/experiments/deepseek_v3/symm_mem_recipes/triton_on_device_all_to_all_v.py
user: login: wconstab
body: sounds good for short term.  longer term we'd want to fix the backward+symm memm issue though right?  can we mark the tensor as a 'buffer' to make autograd not care about it, or something?  
created_at: 2025-03-20T03:38:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -156,3 +158,34 @@ def on_device_all_to_all_v(
     )
     # log_triton_kernel(kernel)
     return output
+
+
+class OnDeviceAllToAllV(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        output: torch.Tensor,
+        output_splits: torch.Tensor,
+        input: torch.Tensor,
+        input_splits: torch.Tensor,
+        group: dist.ProcessGroup = dist.group.WORLD,
+    ):
+        _on_device_all_to_all_v(output, output_splits, input, input_splits, group=group)
+        ctx.save_for_backward(output_splits)
+        ctx.group = group
+        ctx.input_shape = input.shape
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_output_splits = ctx.saved_tensors
+        grad_input_splits = torch.empty_like(grad_output_splits)
+        grad_input = grad_output.new_empty(*ctx.input_shape)
+        _on_device_all_to_all_v(
+            grad_input, grad_input_splits, grad_output, grad_output_splits, group=group
path: torchtitan/experiments/deepseek_v3/symm_mem_recipes/triton_on_device_all_to_all_v.py
user: login: ngimel
body: Autograd does not allow touching the buffer that's going to be used by autograd, is that the case here? If you still need all this memory live, then registering multiple copies with symmetric memory shouldn't be overall memory hit. 
created_at: 2025-03-20T16:17:30Z
author_association: NONE
diff_hunk: @@ -156,3 +158,34 @@ def on_device_all_to_all_v(
     )
     # log_triton_kernel(kernel)
     return output
+
+
+class OnDeviceAllToAllV(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        output: torch.Tensor,
+        output_splits: torch.Tensor,
+        input: torch.Tensor,
+        input_splits: torch.Tensor,
+        group: dist.ProcessGroup = dist.group.WORLD,
+    ):
+        _on_device_all_to_all_v(output, output_splits, input, input_splits, group=group)
+        ctx.save_for_backward(output_splits)
+        ctx.group = group
+        ctx.input_shape = input.shape
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_output_splits = ctx.saved_tensors
+        grad_input_splits = torch.empty_like(grad_output_splits)
+        grad_input = grad_output.new_empty(*ctx.input_shape)
+        _on_device_all_to_all_v(
+            grad_input, grad_input_splits, grad_output, grad_output_splits, group=group
path: torchtitan/experiments/deepseek_v3/symm_mem_recipes/triton_on_device_all_to_all_v.py
user: login: kwen2501
body: @wconstab @ngimel yes. 
@EugenHotaj and @yifuwang suggested something similar -- we can try moving the symm mem inside an autograd function, so that it is not visible to the autograd engine.  
created_at: 2025-03-20T17:20:03Z
author_association: CONTRIBUTOR

PR comments:
user: login: yifuwang
created_at: 2025-03-07T22:12:25Z
body: Looks good. Though you probably want to return both `received` and `output_split_sizes`. This doesn't avoid the CPU sync if you need `received` on the host, but avoids a `sum` kernel launch.
author_association: NONE
user: login: kwen2501
created_at: 2025-03-10T17:33:57Z
body: @yifuwang Good idea. I hope I can find a way to **not** used the summed `received` entirely in some later PRs. So just keeping things as is for now.
author_association: CONTRIBUTOR
user: login: kwen2501
created_at: 2025-03-20T03:08:14Z
body: @wconstab thanks for all the requests for comments. I will put a PR on top of the current stack -- too hard to modify this one without resolving conflicts with other PRs on top.
author_association: CONTRIBUTOR