Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[torch.compile] register allreduce operations as custom ops #8526

Merged
merged 13 commits into from
Sep 17, 2024
Merged

[torch.compile] register allreduce operations as custom ops #8526

Changes from 1 commit
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
a1f1931
init version
youkaichao Sep 13, 2024
7c8fccb
Merge branch 'main' into register_allreduce2
youkaichao Sep 16, 2024
bb2615b
remove spaces
youkaichao Sep 16, 2024
bbdfd6b
move should_custom_ar to python
youkaichao Sep 16, 2024
9b4b802
remove should_custom_ar from c
youkaichao Sep 16, 2024
8376f35
slight changes
youkaichao Sep 16, 2024
82e9581
add tests
youkaichao Sep 16, 2024
d25c8a4
add tests
youkaichao Sep 16, 2024
244bbd4
type annotation
youkaichao Sep 16, 2024
2da8f06
fix type
youkaichao Sep 16, 2024
7da3dc8
add one more condition
youkaichao Sep 17, 2024
9f9359c
update tests
youkaichao Sep 17, 2024
e32ee56
add package
youkaichao Sep 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
add one more condition
  • Loading branch information
@youkaichao
youkaichao committed Sep 17, 2024
commit 7da3dc87264bfbde6c9c9eabdd1aac9c12fb8046
2 changes: 2 additions & 0 deletions vllm/distributed/device_communicators/custom_all_reduce.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -230,6 +230,8 @@ def register_graph_buffers(self):
ops.register_graph_buffers(self._ptr, handles, offsets)

def should_custom_ar(self, inp: torch.Tensor):
if self.disabled:
return False
inp_size = inp.numel() * inp.element_size()
# custom allreduce requires input byte size to be multiples of 16
if inp_size % 16 != 0:
Expand Down
Loading