Skip to content

Conversation

@IvanKobzarev
Copy link
Contributor

@IvanKobzarev IvanKobzarev commented Sep 19, 2025

  1. Fork of torch/_inductor/fx_passes/micro_pipeline_tp.py, torch/distributed/_symmetric_memory/__init__.py for fast experimentation.

PRs with changes on the top of base version to PyTorch repo:

pytorch/pytorch#162794

pytorch/pytorch#163068

pytorch/pytorch#163069

  1. Changes to Solver
  • Communication Cost reduction -10% (heuristic based on AsyncTP post in discussion , needs remeasure):

2.1 matmul + reduce_scatter (Partial -> Shard(dim), dim is not the last of matmul
(the last dim also supported, but has additional restride .contiguous() inside)

2.2 ag + matmul Shard(dim) -> R for argument_A of matmul, dim is not the last dim that will be reduced)

@meta-cla meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Sep 19, 2025
@IvanKobzarev IvanKobzarev changed the title [asynctp] Knobs to enable asynctp; Adding constraints to solver to incentivize asynctp agmm, mmrs [WIP][asynctp] Knobs to enable asynctp; Adding constraints to solver to incentivize asynctp agmm, mmrs Sep 19, 2025
@IvanKobzarev IvanKobzarev force-pushed the async_tp branch 2 times, most recently from 5d54740 to 1d4a56d Compare September 22, 2025 08:18
@IvanKobzarev IvanKobzarev changed the title [WIP][asynctp] Knobs to enable asynctp; Adding constraints to solver to incentivize asynctp agmm, mmrs [asynctp] Async_tp pass and ops fork + changes; Solver addition to incentivize async_tp fusable redistributions Sep 22, 2025
@IvanKobzarev IvanKobzarev force-pushed the async_tp branch 5 times, most recently from 8f8ddd5 to ef3a9a9 Compare September 22, 2025 15:57
@IvanKobzarev IvanKobzarev force-pushed the async_tp branch 2 times, most recently from b981eec to c46a0f6 Compare September 22, 2025 18:05
@fmassa
Copy link
Contributor

fmassa commented Sep 23, 2025

Shared this on chat.

Let me go over this carefully once I'm back from PTO.

IMO having "async-tp" as an option inside AutoParallel seems counterintuitive to me, but let me look into this carefully next week.

@fmassa
Copy link
Contributor

fmassa commented Sep 30, 2025

Subsumed by #167 and #168

@fmassa fmassa closed this Sep 30, 2025
@fmassa fmassa deleted the async_tp branch September 30, 2025 07:06
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants