Issue Station
Sign In
Tools hereToolbox

pytorch - ✅(Solved) Fix [DTensor] greedy redistribution cost favors all_gather over all_to_all [1 pull requests, 3 comments, 3 participants]

Official PRs (…)
ON THIS PAGE

Recommended Tools

×6

Utilities matched from this issue’s tags and category — try them while you read without losing context.

GitHub issue graph ai analysis

Paste a GitHub issue URL. We fetch that issue, discover linked issues from bodies/comments/timeline, collect linked pull requests, and produce a structured English report.

The report is written in English Markdown for sharing and archival.

Helpful · Quick feedback

Loading…
GitHub stats
pytorch/pytorch#177663•Fetched 2026-04-08 00:52:38
View on GitHub
Comments
3
Participants
3
Timeline
46
Reactions
0
Author
wconstabwconstab
Participants
fmassafmassawanchaolwanchaolwconstabwconstab
Timeline (top)
mentioned ×18subscribed ×18labeled ×4commented ×3

Fix Action

Fixed

PR fix notes

PR #178175: [DTensor] Remove artificial +1.0 shard→shard penalty in cost model

Description (problem / solution / changelog)

Stack from ghstack (oldest at bottom):

  • -> #178175
  • #177798

The cost model in _compute_placement_transition_cost added a +1.0 penalty on top of the allgather cost for shard→shard transitions (e.g., S(0)→S(1)), with a comment saying "since we haven't implement it yet". However, alltoall IS implemented via shard_dim_alltoall (C++ in Functional.cpp, calls ncclAllToAll). This artificial penalty caused the cost model to overestimate shard→shard transitions, creating inconsistencies between the full-expansion strategy search and the Dijkstra search (which could find multi-hop paths that bypass the penalty).

This change adds a proper alltoall_cost function and uses it for shard→shard transitions. The formula matches allgather_cost as a safe upper bound — alltoall is likely cheaper on NVSwitch hardware (direct exchange vs ring), but we use the same formula since the existing model is already approximate (hardcoded BW, ring assumption).

With the accurate cost, the planner now correctly prefers alltoall over allgather when alltoall moves less data (alltoall doesn't inflate comm bytes like allgather does). This affects strategy selection for several ops (mm, reductions, scans) and produces shorter redistribution paths.

Fixes #177663

Authored with Claude.

Changed files

  • test/distributed/tensor/debug/test_comm_mode.py (modified, +2/-1)
  • test/distributed/tensor/debug/test_debug_mode.py (modified, +22/-23)
  • test/distributed/tensor/test_math_ops.py (modified, +9/-13)
  • test/distributed/tensor/test_matrix_ops.py (modified, +3/-1)
  • test/distributed/tensor/test_redistribute.py (modified, +4/-6)
  • torch/distributed/tensor/_collective_utils.py (modified, +13/-4)

Code Example

_collective_utils.py:356-360:
  elif current_placement.is_shard() and target_placement.is_shard():
      # should be alltoall comm, since we haven't implement it yet, add 1.0 as penalty
      # to favor allgather instead
      # TODO: add alltoall_cost
      return allgather_cost(comm_bytes_gb, mesh_topo, mesh_dim) + 1.0, comm_bytes_gb
RAW_BUFFERClick to expand / collapse

The code has this TODO

_collective_utils.py:356-360:
  elif current_placement.is_shard() and target_placement.is_shard():
      # should be alltoall comm, since we haven't implement it yet, add 1.0 as penalty
      # to favor allgather instead
      # TODO: add alltoall_cost
      return allgather_cost(comm_bytes_gb, mesh_topo, mesh_dim) + 1.0, comm_bytes_gb

I'd like to remove this +1.0 and the todo. We have had all_to_all support for a while, i'm not sure whether we just forgot to update this code or if there was another reason someone preferred this.

Opening this issue just for folks to have a chance to comment, i'll change it if not.

cc @awgu @wanchaol @fegin @fduwjj @wz337 @d4l3k @pragupta @msaroufim @dcci @aditvenk @xmfan @tianyu-l @XilunWu @SherlockNoMad @ppwwyyxx @zpcore @fmassa @ezyang

extent analysis

Fix Plan

To remove the TODO and the penalty, we need to implement the all_to_all_cost function and use it instead of allgather_cost with a penalty.

Code Changes

We will replace the existing code with the following:

elif current_placement.is_shard() and target_placement.is_shard():
    return all_to_all_cost(comm_bytes_gb, mesh_topo, mesh_dim), comm_bytes_gb

We also need to ensure that the all_to_all_cost function is implemented and returns the correct cost.

Example Implementation

If the all_to_all_cost function is not already implemented, we can add it as follows:

def all_to_all_cost(comm_bytes_gb, mesh_topo, mesh_dim):
    # calculate the cost of all-to-all communication
    # this is a simple example and may need to be adjusted based on the actual cost model
    return comm_bytes_gb * mesh_dim

Verification

To verify that the fix worked, we can test the code with different inputs and check that the correct cost is being returned. We can also compare the performance of the code before and after the change to ensure that it is not introducing any regressions.

Vote matrix · Quick signals

Works
Did the solution work? Tap to confirm.
Easy Fix
Was it a quick fix?
Time Saver
Did it save you time?
Blocking
Was it severely blocking?
Common Issue
Are others likely hitting this too?
Flaky / Intermittent
Is it intermittent?
Verified / Reproducible
Can you reproduce it reliably?
Loading…
#api#ssr#installation#tensor shape#autograd error#embedding generation#cache error#pipeline error#runtime error#dependency conflict

Still need to ship something?

×6

Another batch ranked right after the header list — different links, same matching logic.

Back to top recommendations

TRENDING