#2065: Updated all reduce code to handle 0 or 1 mesh cluster axis #2215
Conversation
tapspatel
requested review from
jnie-TT,
kmabeeTT,
AleksKnezevic,
pilkicTT,
sdjordjevicTT,
svuckovicTT,
mtopalovicTT,
jserbedzijaTT,
azecevicTT,
nsmithtt and
mrakitaTT
as code owners
|
@tapspatel Can you wait for this PR to be merged? #2149 I had to fix some of the code for importing all_reduce and it's conflicting now. I believe you can apply this change on top of the pr. |
tapspatel
force-pushed
the
tpatel/issue-2065
branch
2 times, most recently
from
3002451 to
a846738
Compare
tapspatel
force-pushed
the
tpatel/issue-2065
branch
from
a846738 to
26fddac
Compare
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
tapspatel
force-pushed
the
tpatel/issue-2065
branch
from
26fddac to
ac11675
Compare
|
Perhaps the commit title should become the commit message and the title could be something a bit more concise? |
tapspatel
changed the title
| auto firstElementIt = replicaGroups.begin(); | ||
| auto secondElementIt = firstElementIt + 1; | ||
|
|
||
| clusterAxis = (((*firstElementIt) + 1) == *secondElementIt); |
There was a problem hiding this comment.
Doesn't this assume that device IDs are consecutive and in a particular topology?
There was a problem hiding this comment.
From docs: https://docs.jax.dev/en/latest/_autosummary/jax.make_mesh.html
Essentially, the mesh device ordering is determined by the TPU structure internally. It will optimize the ordering based on the mesh shape. However, if you simulate CPUs, it orders by consecutive numbers. To simplify this, it is an assumption that the mesh provided by tt-xla will be monotonically increasing, because we don't have a way to propagate an efficient mesh mapping based on hardware structure yet in compiler (but this can be done).
I can add this assumption as a comment for now.
There was a problem hiding this comment.
For topologies, we are only supporting 2d grid topologies in compiler from conversation with Wooseok.
There was a problem hiding this comment.
I have observed the ascending device ids from 0 to N-1 so far as well, so I guess it's fine. In the future, we may leverage the device IDs to best distribute data for our hardware config, but didn't get there yet.
tapspatel
force-pushed
the
tpatel/issue-2065
branch
from
ac11675 to
d5d9ebd
Compare
| auto firstElementIt = replicaGroups.begin(); | ||
| auto secondElementIt = firstElementIt + 1; | ||
|
|
||
| clusterAxis = (((*firstElementIt) + 1) == *secondElementIt); |
There was a problem hiding this comment.
I have observed the ascending device ids from 0 to N-1 so far as well, so I guess it's fine. In the future, we may leverage the device IDs to best distribute data for our hardware config, but didn't get there yet.
tapspatel
force-pushed
the
tpatel/issue-2065
branch
from
d5d9ebd to
959ef8e
Compare
tapspatel
force-pushed
the
tpatel/issue-2065
branch
from
959ef8e to
ca39af3
Compare
|
@wooseokTT let me know if the responses to all the comments are sufficient. |
tapspatel
force-pushed
the
tpatel/issue-2065
branch
from
ca39af3 to
0cec5f6
Compare
…ned up dialect representations of all reduce in ttir and ttnn. Update algorithms for calculating gather and scatter dimensions
tapspatel
force-pushed
the
tpatel/issue-2065
branch
from
0cec5f6 to
f02b2fa
Compare
) Updated all reduce code to handle 0 or 1 cluster axis and cleaned up dialect representations of all reduce in ttir and ttnn. Update algorithms for calculating gather and scatter dimensions. Migrated all workaround code into TTNN workaround pass such that we don't clog up ttir or ttnn definitions of all_reduce.
#2065: Updated all reduce code to handle 0 or 1 cluster axis and cleaned up dialect representations of all reduce in ttir and ttnn. Update algorithms for calculating gather and scatter dimensions. Migrated all workaround code into TTNN workaround pass such that we don't clog up ttir or ttnn definitions of all_reduce.