Add support for minimal all_reduce for Llama shapes #17792

avoraTT · 2025-02-10T17:32:55Z

Ticket

Problem description

Provide context for the problem.

What's changed

Describe the approach used to solve the problem.
Summarize the changes made and its impact.

Checklist

All post commit CI passes
Blackhole Post commit CI passes (if applicable)
Model regression CI passes (if applicable)
Device performance regression CI passes (if applicable)
(For models and ops writers) Full new models tests CI passes (if applicable)
New/Existing tests provide coverage for changes

…kernel (dataflow + compute)

…_wait in reduction

…opy into the output cb (temporary). Next: add reduction compute kernel.

… tensor.

…el error Cannot add semaphore on core (x=0,y=0). Max number of semaphores (8) reached./build_metal.sh --debug Link=2 hangs at reduction signal wait only for second worker

…t of e2e trace perf

caixunshiren added 2 commits February 10, 2025 16:43

added ccl async llama perf and correctness measurement tests to CI

503ca1c

updated ccl perf target for llama all gather async

ccd3f98

avoraTT self-assigned this Feb 10, 2025

avoraTT and others added 22 commits February 10, 2025 17:48

WIP adding infra for the new all reduce.

0c9ccd7

interleaved all gather works with good PCC. Next step: add reduction …

c2a11de

…kernel (dataflow + compute)

wip reduction stuff.

4b67c01

#0: added noc semaphore multicast in writer, seeing a hang on noc_sem…

ce9e52d

…_wait in reduction

Add fix for reduction worker hang.

bf9cb47

Add reduction and output cb. Currently, the reduction kernel does a c…

600352e

…opy into the output cb (temporary). Next: add reduction compute kernel.

All-reduce for FF1/FF3 works

5d4d879

Add support for reshard. TODO: add support to drop padding from input…

27414ef

… tensor.

Add support for unpadded shapes.

31119ef

Remove dprints.

24e65bb

Fix bug in mcast bbox. Fix QKV output num cores.

7251015

#0: multi-link support added for 3 all_reduce. Link=3 fails with kern…

b67ea6d

…el error Cannot add semaphore on core (x=0,y=0). Max number of semaphores (8) reached./build_metal.sh --debug Link=2 hangs at reduction signal wait only for second worker

#0: multi-link=3 works

21f1779

Add cleanup for multi-link.

e4a7b09

Rebase and fix/cleanup stuff.

e0ce950

Clean up pytest and enable trace.

ea2edb7

Adding gsem fix for multi-iter.

b6277d9

#0: added api to subtract corerangesets

4520362

#0: updated choose_worker_cores function to omit reserved_cores

203e2e8

#0: fix placement of link worker cores

694ef4d

Add support for input shard not divisble by output shard.

a131897

added all reduce into llama ccl perf test and added proper measuremen…

bcf48a6

…t of e2e trace perf

caixunshiren force-pushed the avora/all_reduce branch from 088cb34 to bcf48a6 Compare February 10, 2025 22:24

Extend llama sharded all gather for LN

b779894

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add support for minimal all_reduce for Llama shapes #17792

Add support for minimal all_reduce for Llama shapes #17792

avoraTT commented Feb 10, 2025

Add support for minimal all_reduce for Llama shapes #17792

Are you sure you want to change the base?

Add support for minimal all_reduce for Llama shapes #17792

Conversation

avoraTT commented Feb 10, 2025

Ticket

Problem description

What's changed

Checklist