[NVIDIA GPU] Support multi-operand collective-permute #18838

terryysun · 2024-10-29T20:12:55Z

For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because

it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion;
fewer collectives make it easier for LHS to make better decision.

In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands.

frgossen

Looks good overall. Have a few comments and questions.

Can you check to use the XLA casts?
What is the inplace change needed for?

frgossen · 2024-11-01T16:02:12Z

xla/hlo/analysis/hlo_dataflow_analysis.cc

          // data are optional.
          define_value_at(/*index=*/{});
          define_value_at(/*index=*/{1});
          for (int i = 2; i < instruction->shape().tuple_shapes_size(); ++i) {
            define_value_at(/*index=*/{i});
          }

-          if (instruction->operand_count() > 1) {
+          if (static_cast<HloCollectivePermuteInstruction*>(instruction)


Use Xla's casts

frgossen · 2024-11-01T16:06:12Z

xla/service/hlo_verifier.cc

@@ -681,9 +681,11 @@ absl::Status CheckBufferOffset(const Shape& buffer_shape,
 }

 absl::Status CheckInplaceCollectivePermute(HloInstruction* collective_permute) {
-  if (collective_permute->operand_count() == 1) {
+  if (!static_cast<HloCollectivePermuteInstruction*>(collective_permute)


Please use XLA cast

frgossen · 2024-11-01T16:08:30Z

xla/hlo/ir/hlo_instructions.h

@@ -1008,6 +1012,7 @@ class HloCollectivePermuteInstruction : public HloChannelInstruction {

  const std::vector<std::pair<int64_t, int64_t>> source_target_pairs_;
  const std::vector<std::vector<int64_t>> slice_sizes_;
+  bool inplace_;


What is this needed for? Is this necessary for combined collective-permutes?

There are a couple of places in the codebase where in-place or not is determined by the number of operands (for example https://github.com/openxla/xla/pull/18838/files#diff-dad808ebc0b02889a21c167165b2c2c67dd7b691b99a8d1d9f77563c3a9edd3cL97 and https://github.com/openxla/xla/pull/18838/files#diff-b97e25fdfa46e451787dfc87463abf3af674798e57b45b3450f0bb3aaac875f2L2588), basically any cp with > 1 operands will be treated as an in-place cp, which should have 4 operands. With the multi-operand support, this assumption no longer holds and we need to let them rely on this new flag.

frgossen

Thanks!

penpornk · 2024-11-06T19:43:30Z

xla/python/ops.cc

  ops.def("CollectivePermute", &CollectivePermute, nb::arg("operand"),
-          nb::arg("source_target_pairs"), nb::arg("channel_id") = std::nullopt);
+          nb::arg("source_target_pairs"), nb::arg("channel_id") = std::nullopt,
+          nb::arg("inplace") = false);
+  ops.def("MultiCollectivePermute", &MultiCollectivePermute,
+          nb::arg("operands"), nb::arg("source_target_pairs"),
+          nb::arg("channel_id") = std::nullopt, nb::arg("inplace") = false);


Could you please update their corresponding definitions in //xla/python/xla_extension/ops.pyi as well?
E.g.,

xla/xla/python/xla_extension/ops.pyi

Lines 135 to 138 in 0f6331b

def CollectivePermute(

operand: XlaOp,

source_target_pairs: Sequence[tuple[int, int]],

channel_id: Optional[_ChannelHandle] = ...) -> XlaOp: ...

Thank you! I think the update needs some more changes though. Please see my other comments on the file ops.pyi.

penpornk · 2024-11-07T08:01:48Z

xla/python/xla_extension/ops.pyi

@@ -136,6 +136,10 @@ def CollectivePermute(
    operand: XlaOp,
    source_target_pairs: Sequence[tuple[int, int]],
    channel_id: Optional[_ChannelHandle] = ...) -> XlaOp: ...
+def MultiCollectivePermute(


The list is alphabetically sorted, so this should be lower in the file (between Map and NextAfter).

thanks! updated

penpornk · 2024-11-07T08:02:30Z

xla/python/xla_extension/ops.pyi

@@ -136,6 +136,10 @@ def CollectivePermute(
    operand: XlaOp,
    source_target_pairs: Sequence[tuple[int, int]],
    channel_id: Optional[_ChannelHandle] = ...) -> XlaOp: ...


Please update the parameter list (add inplace).

penpornk · 2024-11-07T08:03:17Z

xla/python/xla_extension/ops.pyi

+def MultiCollectivePermute(
+    operands: Sequence[XlaOp],
+    source_target_pairs: Sequence[tuple[int, int]],
+    channel_id: Optional[_ChannelHandle] = ...) -> XlaOp: ...


Please add inplace to the parameter list.

penpornk · 2024-11-07T08:03:47Z

xla/python/ops.cc

  ops.def("CollectivePermute", &CollectivePermute, nb::arg("operand"),
-          nb::arg("source_target_pairs"), nb::arg("channel_id") = std::nullopt);
+          nb::arg("source_target_pairs"), nb::arg("channel_id") = std::nullopt,
+          nb::arg("inplace") = false);
+  ops.def("MultiCollectivePermute", &MultiCollectivePermute,
+          nb::arg("operands"), nb::arg("source_target_pairs"),
+          nb::arg("channel_id") = std::nullopt, nb::arg("inplace") = false);


Thank you! I think the update needs some more changes though. Please see my other comments on the file ops.pyi.

penpornk

Thank you for the changes!

terryysun · 2024-11-08T21:20:04Z

hi @penpornk could you help check why this pr is not merging? thanks!

Imported from GitHub PR #18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba by Terry Sun <[email protected]>: support multi-operand cp -- 170fead by Terry Sun <[email protected]>: minor refactoring -- 0d85070 by Terry Sun <[email protected]>: update python interface -- 9812a10 by Terry Sun <[email protected]>: polish python interface Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=#18838 from terryysun:terryysun/grouped_cp 9812a10 PiperOrigin-RevId: 696044196

Imported from GitHub PR openxla/xla#18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba5b8f6ae66d1071a1894a87987b6a5bceb by Terry Sun <[email protected]>: support multi-operand cp -- 170fead3de942f5e14f4936df1d76bf7e5e319d4 by Terry Sun <[email protected]>: minor refactoring -- 0d85070baee3f26075f0b3660c4674d7b414c861 by Terry Sun <[email protected]>: update python interface -- 9812a104822ea479d29fef0531b9e10d5c2a831d by Terry Sun <[email protected]>: polish python interface Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#18838 from terryysun:terryysun/grouped_cp 9812a104822ea479d29fef0531b9e10d5c2a831d PiperOrigin-RevId: 696044196

Imported from GitHub PR #18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba by Terry Sun <[email protected]>: support multi-operand cp -- 170fead by Terry Sun <[email protected]>: minor refactoring -- 0d85070 by Terry Sun <[email protected]>: update python interface -- 9812a10 by Terry Sun <[email protected]>: polish python interface Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=#18838 from terryysun:terryysun/grouped_cp 9812a10 PiperOrigin-RevId: 696044196

Imported from GitHub PR openxla/xla#18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba5b8f6ae66d1071a1894a87987b6a5bceb by Terry Sun <[email protected]>: support multi-operand cp -- 170fead3de942f5e14f4936df1d76bf7e5e319d4 by Terry Sun <[email protected]>: minor refactoring -- 0d85070baee3f26075f0b3660c4674d7b414c861 by Terry Sun <[email protected]>: update python interface -- 9812a104822ea479d29fef0531b9e10d5c2a831d by Terry Sun <[email protected]>: polish python interface Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#18838 from terryysun:terryysun/grouped_cp 9812a104822ea479d29fef0531b9e10d5c2a831d PiperOrigin-RevId: 696044196

Imported from GitHub PR #18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba by Terry Sun <[email protected]>: support multi-operand cp -- 170fead by Terry Sun <[email protected]>: minor refactoring -- 0d85070 by Terry Sun <[email protected]>: update python interface -- 9812a10 by Terry Sun <[email protected]>: polish python interface Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=#18838 from terryysun:terryysun/grouped_cp 9812a10 PiperOrigin-RevId: 696044196

Imported from GitHub PR openxla/xla#18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba5b8f6ae66d1071a1894a87987b6a5bceb by Terry Sun <[email protected]>: support multi-operand cp -- 170fead3de942f5e14f4936df1d76bf7e5e319d4 by Terry Sun <[email protected]>: minor refactoring -- 0d85070baee3f26075f0b3660c4674d7b414c861 by Terry Sun <[email protected]>: update python interface -- 9812a104822ea479d29fef0531b9e10d5c2a831d by Terry Sun <[email protected]>: polish python interface Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#18838 from terryysun:terryysun/grouped_cp 9812a104822ea479d29fef0531b9e10d5c2a831d PiperOrigin-RevId: 696044196

penpornk · 2024-11-13T19:34:16Z

@terryysun Could you please help apply clang-format?
Here are the formatting errors:
https://github.com/openxla/xla/actions/runs/11729067298/job/32900476498?pr=18838

penpornk

Thank you for the change!

Imported from GitHub PR #18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba by Terry Sun <[email protected]>: support multi-operand cp -- 170fead by Terry Sun <[email protected]>: minor refactoring -- 0d85070 by Terry Sun <[email protected]>: update python interface -- 9812a10 by Terry Sun <[email protected]>: polish python interface -- 3a1552c by Terry Sun <[email protected]>: formatting -- d3657f8 by Terry Sun <[email protected]>: formatting -- 17e640b by Terry Sun <[email protected]>: fix minor issues -- ed899a9 by Terry Sun <[email protected]>: minor fix Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=#18838 from terryysun:terryysun/grouped_cp ed899a9 PiperOrigin-RevId: 693728463

Imported from GitHub PR #18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba by Terry Sun <[email protected]>: support multi-operand cp -- 170fead by Terry Sun <[email protected]>: minor refactoring -- 0d85070 by Terry Sun <[email protected]>: update python interface -- 9812a10 by Terry Sun <[email protected]>: polish python interface -- 3a1552c by Terry Sun <[email protected]>: formatting -- d3657f8 by Terry Sun <[email protected]>: formatting -- c9202fa by Terry Sun <[email protected]>: fix minor issues Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=#18838 from terryysun:terryysun/grouped_cp c9202fa PiperOrigin-RevId: 693728463

Imported from GitHub PR openxla/xla#18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba5b8f6ae66d1071a1894a87987b6a5bceb by Terry Sun <[email protected]>: support multi-operand cp -- 170fead3de942f5e14f4936df1d76bf7e5e319d4 by Terry Sun <[email protected]>: minor refactoring -- 0d85070baee3f26075f0b3660c4674d7b414c861 by Terry Sun <[email protected]>: update python interface -- 9812a104822ea479d29fef0531b9e10d5c2a831d by Terry Sun <[email protected]>: polish python interface -- 3a1552cbcd2e26f814373e0e01adbe8eceb3be9f by Terry Sun <[email protected]>: formatting -- d3657f81ac57dc1de86561b3449d051d178e0f75 by Terry Sun <[email protected]>: formatting -- c9202facad49a7608ae633130f887f9dc0706191 by Terry Sun <[email protected]>: fix minor issues Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#18838 from terryysun:terryysun/grouped_cp c9202facad49a7608ae633130f887f9dc0706191 PiperOrigin-RevId: 693728463

dimitar-asenov · 2024-11-19T10:05:16Z

Not sure what the diff was before, but it's identical now. Trying to submit it.

dimitar-asenov · 2024-11-19T10:07:13Z

Is it possible that we enforce the same standard for OSS as Google internal build?

We definitely try to make the checks as uniform as possible. As we keep changing things, some divergence is inevitable.

@ddunl Would it be possible to make sure the issue in #18838 (comment) are also detected in OSS?

Imported from GitHub PR #18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba by Terry Sun <[email protected]>: support multi-operand cp -- 170fead by Terry Sun <[email protected]>: minor refactoring -- 0d85070 by Terry Sun <[email protected]>: update python interface -- 9812a10 by Terry Sun <[email protected]>: polish python interface -- 3a1552c by Terry Sun <[email protected]>: formatting -- d3657f8 by Terry Sun <[email protected]>: formatting -- 17e640b by Terry Sun <[email protected]>: fix minor issues Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=#18838 from terryysun:terryysun/grouped_cp 17e640b PiperOrigin-RevId: 697553650

Imported from GitHub PR openxla/xla#18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba5b8f6ae66d1071a1894a87987b6a5bceb by Terry Sun <[email protected]>: support multi-operand cp -- 170fead3de942f5e14f4936df1d76bf7e5e319d4 by Terry Sun <[email protected]>: minor refactoring -- 0d85070baee3f26075f0b3660c4674d7b414c861 by Terry Sun <[email protected]>: update python interface -- 9812a104822ea479d29fef0531b9e10d5c2a831d by Terry Sun <[email protected]>: polish python interface -- 3a1552cbcd2e26f814373e0e01adbe8eceb3be9f by Terry Sun <[email protected]>: formatting -- d3657f81ac57dc1de86561b3449d051d178e0f75 by Terry Sun <[email protected]>: formatting -- 17e640b47719c46aef051ef7a27b88a245591074 by Terry Sun <[email protected]>: fix minor issues Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#18838 from terryysun:terryysun/grouped_cp 17e640b47719c46aef051ef7a27b88a245591074 PiperOrigin-RevId: 697553650

dimitar-asenov · 2024-11-19T10:22:29Z

Ah, sorry, we hit yet more internal issues. Similarly to my previous comment, please make sure that InferCollectivePermuteStartShape also has one overload that works consistently with both one or more shapes. Right now, there are two overloads and they don't work consistently (e.g. there is a difference in behavior between {shape} and just shape, and that's confusing. The additional overload also causes some internal code to not compile.

Imported from GitHub PR #18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba by Terry Sun <[email protected]>: support multi-operand cp -- 170fead by Terry Sun <[email protected]>: minor refactoring -- 0d85070 by Terry Sun <[email protected]>: update python interface -- 9812a10 by Terry Sun <[email protected]>: polish python interface -- 3a1552c by Terry Sun <[email protected]>: formatting -- d3657f8 by Terry Sun <[email protected]>: formatting -- c9202fa by Terry Sun <[email protected]>: fix minor issues Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=#18838 from terryysun:terryysun/grouped_cp c9202fa PiperOrigin-RevId: 693728463

Imported from GitHub PR openxla/xla#18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba5b8f6ae66d1071a1894a87987b6a5bceb by Terry Sun <[email protected]>: support multi-operand cp -- 170fead3de942f5e14f4936df1d76bf7e5e319d4 by Terry Sun <[email protected]>: minor refactoring -- 0d85070baee3f26075f0b3660c4674d7b414c861 by Terry Sun <[email protected]>: update python interface -- 9812a104822ea479d29fef0531b9e10d5c2a831d by Terry Sun <[email protected]>: polish python interface -- 3a1552cbcd2e26f814373e0e01adbe8eceb3be9f by Terry Sun <[email protected]>: formatting -- d3657f81ac57dc1de86561b3449d051d178e0f75 by Terry Sun <[email protected]>: formatting -- c9202facad49a7608ae633130f887f9dc0706191 by Terry Sun <[email protected]>: fix minor issues Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#18838 from terryysun:terryysun/grouped_cp c9202facad49a7608ae633130f887f9dc0706191 PiperOrigin-RevId: 693728463

ddunl · 2024-11-19T18:47:56Z

Is it possible that we enforce the same standard for OSS as Google internal build?

We definitely try to make the checks as uniform as possible. As we keep changing things, some divergence is inevitable.

@ddunl Would it be possible to make sure the issue in #18838 (comment) are also detected in OSS?

I'll have to see what triggered that error internally. If it's a METADATA check or a publicly available clang-tidy check I probably can, if not it's trickier.

Imported from GitHub PR #18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba by Terry Sun <[email protected]>: support multi-operand cp -- 170fead by Terry Sun <[email protected]>: minor refactoring -- 0d85070 by Terry Sun <[email protected]>: update python interface -- 9812a10 by Terry Sun <[email protected]>: polish python interface -- 3a1552c by Terry Sun <[email protected]>: formatting -- d3657f8 by Terry Sun <[email protected]>: formatting -- 65c3484 by Terry Sun <[email protected]>: refactor overloading Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=#18838 from terryysun:terryysun/grouped_cp 65c3484 PiperOrigin-RevId: 693728463

Imported from GitHub PR openxla/xla#18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba5b8f6ae66d1071a1894a87987b6a5bceb by Terry Sun <[email protected]>: support multi-operand cp -- 170fead3de942f5e14f4936df1d76bf7e5e319d4 by Terry Sun <[email protected]>: minor refactoring -- 0d85070baee3f26075f0b3660c4674d7b414c861 by Terry Sun <[email protected]>: update python interface -- 9812a104822ea479d29fef0531b9e10d5c2a831d by Terry Sun <[email protected]>: polish python interface -- 3a1552cbcd2e26f814373e0e01adbe8eceb3be9f by Terry Sun <[email protected]>: formatting -- d3657f81ac57dc1de86561b3449d051d178e0f75 by Terry Sun <[email protected]>: formatting -- 65c3484b0a5face53a5e6980d2b74fb00b99f5bd by Terry Sun <[email protected]>: refactor overloading Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#18838 from terryysun:terryysun/grouped_cp 65c3484b0a5face53a5e6980d2b74fb00b99f5bd PiperOrigin-RevId: 693728463

Imported from GitHub PR #18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba by Terry Sun <[email protected]>: support multi-operand cp -- 170fead by Terry Sun <[email protected]>: minor refactoring -- 0d85070 by Terry Sun <[email protected]>: update python interface -- 9812a10 by Terry Sun <[email protected]>: polish python interface -- 3a1552c by Terry Sun <[email protected]>: formatting -- d3657f8 by Terry Sun <[email protected]>: formatting -- 65c3484 by Terry Sun <[email protected]>: refactor overloading Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=#18838 from terryysun:terryysun/grouped_cp 65c3484 PiperOrigin-RevId: 693728463

Imported from GitHub PR openxla/xla#18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba5b8f6ae66d1071a1894a87987b6a5bceb by Terry Sun <[email protected]>: support multi-operand cp -- 170fead3de942f5e14f4936df1d76bf7e5e319d4 by Terry Sun <[email protected]>: minor refactoring -- 0d85070baee3f26075f0b3660c4674d7b414c861 by Terry Sun <[email protected]>: update python interface -- 9812a104822ea479d29fef0531b9e10d5c2a831d by Terry Sun <[email protected]>: polish python interface -- 3a1552cbcd2e26f814373e0e01adbe8eceb3be9f by Terry Sun <[email protected]>: formatting -- d3657f81ac57dc1de86561b3449d051d178e0f75 by Terry Sun <[email protected]>: formatting -- 65c3484b0a5face53a5e6980d2b74fb00b99f5bd by Terry Sun <[email protected]>: refactor overloading Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#18838 from terryysun:terryysun/grouped_cp 65c3484b0a5face53a5e6980d2b74fb00b99f5bd PiperOrigin-RevId: 693728463

Imported from GitHub PR #18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba by Terry Sun <[email protected]>: support multi-operand cp -- 170fead by Terry Sun <[email protected]>: minor refactoring -- 0d85070 by Terry Sun <[email protected]>: update python interface -- 9812a10 by Terry Sun <[email protected]>: polish python interface -- 3a1552c by Terry Sun <[email protected]>: formatting -- d3657f8 by Terry Sun <[email protected]>: formatting -- 65c3484 by Terry Sun <[email protected]>: refactor overloading Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=#18838 from terryysun:terryysun/grouped_cp 65c3484 PiperOrigin-RevId: 693728463

Imported from GitHub PR openxla/xla#18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba5b8f6ae66d1071a1894a87987b6a5bceb by Terry Sun <[email protected]>: support multi-operand cp -- 170fead3de942f5e14f4936df1d76bf7e5e319d4 by Terry Sun <[email protected]>: minor refactoring -- 0d85070baee3f26075f0b3660c4674d7b414c861 by Terry Sun <[email protected]>: update python interface -- 9812a104822ea479d29fef0531b9e10d5c2a831d by Terry Sun <[email protected]>: polish python interface -- 3a1552cbcd2e26f814373e0e01adbe8eceb3be9f by Terry Sun <[email protected]>: formatting -- d3657f81ac57dc1de86561b3449d051d178e0f75 by Terry Sun <[email protected]>: formatting -- 65c3484b0a5face53a5e6980d2b74fb00b99f5bd by Terry Sun <[email protected]>: refactor overloading Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#18838 from terryysun:terryysun/grouped_cp 65c3484b0a5face53a5e6980d2b74fb00b99f5bd PiperOrigin-RevId: 693728463

Imported from GitHub PR #18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba by Terry Sun <[email protected]>: support multi-operand cp -- 170fead by Terry Sun <[email protected]>: minor refactoring -- 0d85070 by Terry Sun <[email protected]>: update python interface -- 9812a10 by Terry Sun <[email protected]>: polish python interface -- 3a1552c by Terry Sun <[email protected]>: formatting -- d3657f8 by Terry Sun <[email protected]>: formatting -- 65c3484 by Terry Sun <[email protected]>: refactor overloading Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=#18838 from terryysun:terryysun/grouped_cp 65c3484 PiperOrigin-RevId: 697553650

Imported from GitHub PR openxla/xla#18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba5b8f6ae66d1071a1894a87987b6a5bceb by Terry Sun <[email protected]>: support multi-operand cp -- 170fead3de942f5e14f4936df1d76bf7e5e319d4 by Terry Sun <[email protected]>: minor refactoring -- 0d85070baee3f26075f0b3660c4674d7b414c861 by Terry Sun <[email protected]>: update python interface -- 9812a104822ea479d29fef0531b9e10d5c2a831d by Terry Sun <[email protected]>: polish python interface -- 3a1552cbcd2e26f814373e0e01adbe8eceb3be9f by Terry Sun <[email protected]>: formatting -- d3657f81ac57dc1de86561b3449d051d178e0f75 by Terry Sun <[email protected]>: formatting -- 65c3484b0a5face53a5e6980d2b74fb00b99f5bd by Terry Sun <[email protected]>: refactor overloading Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#18838 from terryysun:terryysun/grouped_cp 65c3484b0a5face53a5e6980d2b74fb00b99f5bd PiperOrigin-RevId: 697553650

Imported from GitHub PR #18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba by Terry Sun <[email protected]>: support multi-operand cp -- 170fead by Terry Sun <[email protected]>: minor refactoring -- 0d85070 by Terry Sun <[email protected]>: update python interface -- 9812a10 by Terry Sun <[email protected]>: polish python interface -- 3a1552c by Terry Sun <[email protected]>: formatting -- d3657f8 by Terry Sun <[email protected]>: formatting -- 65c3484 by Terry Sun <[email protected]>: refactor overloading Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=#18838 from terryysun:terryysun/grouped_cp 65c3484 PiperOrigin-RevId: 697553650

Imported from GitHub PR openxla/xla#18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba5b8f6ae66d1071a1894a87987b6a5bceb by Terry Sun <[email protected]>: support multi-operand cp -- 170fead3de942f5e14f4936df1d76bf7e5e319d4 by Terry Sun <[email protected]>: minor refactoring -- 0d85070baee3f26075f0b3660c4674d7b414c861 by Terry Sun <[email protected]>: update python interface -- 9812a104822ea479d29fef0531b9e10d5c2a831d by Terry Sun <[email protected]>: polish python interface -- 3a1552cbcd2e26f814373e0e01adbe8eceb3be9f by Terry Sun <[email protected]>: formatting -- d3657f81ac57dc1de86561b3449d051d178e0f75 by Terry Sun <[email protected]>: formatting -- 65c3484b0a5face53a5e6980d2b74fb00b99f5bd by Terry Sun <[email protected]>: refactor overloading Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#18838 from terryysun:terryysun/grouped_cp 65c3484b0a5face53a5e6980d2b74fb00b99f5bd PiperOrigin-RevId: 697553650

Imported from GitHub PR #18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba by Terry Sun <[email protected]>: support multi-operand cp -- 170fead by Terry Sun <[email protected]>: minor refactoring -- 0d85070 by Terry Sun <[email protected]>: update python interface -- 9812a10 by Terry Sun <[email protected]>: polish python interface -- 3a1552c by Terry Sun <[email protected]>: formatting -- d3657f8 by Terry Sun <[email protected]>: formatting -- 65c3484 by Terry Sun <[email protected]>: refactor overloading Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=#18838 from terryysun:terryysun/grouped_cp 65c3484 PiperOrigin-RevId: 697553650

Imported from GitHub PR openxla/xla#18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba5b8f6ae66d1071a1894a87987b6a5bceb by Terry Sun <[email protected]>: support multi-operand cp -- 170fead3de942f5e14f4936df1d76bf7e5e319d4 by Terry Sun <[email protected]>: minor refactoring -- 0d85070baee3f26075f0b3660c4674d7b414c861 by Terry Sun <[email protected]>: update python interface -- 9812a104822ea479d29fef0531b9e10d5c2a831d by Terry Sun <[email protected]>: polish python interface -- 3a1552cbcd2e26f814373e0e01adbe8eceb3be9f by Terry Sun <[email protected]>: formatting -- d3657f81ac57dc1de86561b3449d051d178e0f75 by Terry Sun <[email protected]>: formatting -- 65c3484b0a5face53a5e6980d2b74fb00b99f5bd by Terry Sun <[email protected]>: refactor overloading Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#18838 from terryysun:terryysun/grouped_cp 65c3484b0a5face53a5e6980d2b74fb00b99f5bd PiperOrigin-RevId: 697553650

dimitar-asenov · 2024-11-20T15:14:40Z

@terryysun The PR fails this test:

xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc

Line 6950 in 66073ea

TEST_F(MemorySpaceAssignmentTest, AsyncOpShortLiveRange) {

Could oyu please have a look?

support multi-operand cp

5e10aba

penpornk requested a review from frgossen October 30, 2024 10:47

frgossen suggested changes Nov 1, 2024

View reviewed changes

terryysun requested a review from frgossen November 2, 2024 03:14

minor refactoring

170fead

frgossen approved these changes Nov 6, 2024

View reviewed changes

penpornk suggested changes Nov 6, 2024

View reviewed changes

terryysun requested a review from penpornk November 7, 2024 03:04

update python interface

0d85070

penpornk suggested changes Nov 7, 2024

View reviewed changes

polish python interface

9812a10

terryysun requested a review from penpornk November 7, 2024 18:42

penpornk approved these changes Nov 7, 2024

View reviewed changes

copybara-service bot mentioned this pull request Nov 13, 2024

PR #18838: [NVIDIA GPU] Support multi-operand collective-permute #19292

Open

copybara-service bot mentioned this pull request Nov 13, 2024

PR #18838: [NVIDIA GPU] Support multi-operand collective-permute tensorflow/tensorflow#79946

Draft

formatting

3a1552c

penpornk approved these changes Nov 13, 2024

View reviewed changes

frgossen approved these changes Nov 13, 2024

View reviewed changes

terryysun added 2 commits November 13, 2024 22:13

Merge branch 'main' into terryysun/grouped_cp

681e38f

formatting

d3657f8

penpornk approved these changes Nov 14, 2024

View reviewed changes

copybara-service bot mentioned this pull request Nov 19, 2024

PR #18838: [NVIDIA GPU] Support multi-operand collective-permute #19473

Open

copybara-service bot mentioned this pull request Nov 19, 2024

PR #18838: [NVIDIA GPU] Support multi-operand collective-permute tensorflow/tensorflow#80245

Draft

terryysun force-pushed the terryysun/grouped_cp branch from c9202fa to 65c3484 Compare November 19, 2024 18:34

refactor overloading

65c3484

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[NVIDIA GPU] Support multi-operand collective-permute #18838

[NVIDIA GPU] Support multi-operand collective-permute #18838

terryysun commented Oct 29, 2024 •

edited

Loading

frgossen left a comment

frgossen Nov 1, 2024

terryysun Nov 2, 2024

frgossen Nov 1, 2024

terryysun Nov 2, 2024

frgossen Nov 1, 2024

terryysun Nov 2, 2024

frgossen left a comment

penpornk Nov 6, 2024

terryysun Nov 7, 2024

penpornk Nov 7, 2024

penpornk Nov 7, 2024

terryysun Nov 7, 2024

penpornk Nov 7, 2024

terryysun Nov 7, 2024

penpornk Nov 7, 2024

terryysun Nov 7, 2024

penpornk Nov 7, 2024

penpornk left a comment

terryysun commented Nov 8, 2024

penpornk commented Nov 13, 2024

penpornk left a comment

dimitar-asenov commented Nov 19, 2024

dimitar-asenov commented Nov 19, 2024

dimitar-asenov commented Nov 19, 2024

ddunl commented Nov 19, 2024

dimitar-asenov commented Nov 20, 2024

	def CollectivePermute(
	operand: XlaOp,
	source_target_pairs: Sequence[tuple[int, int]],
	channel_id: Optional[_ChannelHandle] = ...) -> XlaOp: ...

[NVIDIA GPU] Support multi-operand collective-permute #18838

Are you sure you want to change the base?

[NVIDIA GPU] Support multi-operand collective-permute #18838

Conversation

terryysun commented Oct 29, 2024 • edited Loading

frgossen left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

frgossen left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

penpornk left a comment

Choose a reason for hiding this comment

terryysun commented Nov 8, 2024

penpornk commented Nov 13, 2024

penpornk left a comment

Choose a reason for hiding this comment

dimitar-asenov commented Nov 19, 2024

dimitar-asenov commented Nov 19, 2024

dimitar-asenov commented Nov 19, 2024

ddunl commented Nov 19, 2024

dimitar-asenov commented Nov 20, 2024

terryysun commented Oct 29, 2024 •

edited

Loading