You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "farshidsp (via GitHub)" <gi...@apache.org> on 2023/09/15 23:47:02 UTC
[GitHub] [tvm] farshidsp opened a new pull request, #15764: [Disco] Add AllGather
farshidsp opened a new pull request, #15764:
URL: https://github.com/apache/tvm/pull/15764
This pr introduce `allgather`, which collects data from multiple processes or workers and distribute the gathered data to all processes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #15764: [Disco] Add AllGather
Posted by "junrushao (via GitHub)" <gi...@apache.org>.
junrushao commented on code in PR #15764:
URL: https://github.com/apache/tvm/pull/15764#discussion_r1333398046
##########
src/runtime/disco/nccl/nccl.cc:
##########
@@ -67,7 +67,6 @@ inline void SetDevice(int device_id) { CUDA_CALL(cudaSetDevice(device_id)); }
inline void StreamSynchronize(deviceStream_t stream) { CUDA_CALL(cudaStreamSynchronize(stream)); }
inline void StreamCreate(deviceStream_t* stream) { CUDA_CALL(cudaStreamCreate(stream)); }
inline void StreamDestroy(deviceStream_t stream) { CUDA_CALL(cudaStreamDestroy(stream)); }
-inline void SetDevice(int device_id) { CUDA_CALL(cudaSetDevice(device_id)); }
Review Comment:
ooops, thanks for catching this!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg merged pull request #15764: [Disco] Add AllGather
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg merged PR #15764:
URL: https://github.com/apache/tvm/pull/15764
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] farshidsp commented on pull request #15764: [Disco] Add AllGather
Posted by "farshidsp (via GitHub)" <gi...@apache.org>.
farshidsp commented on PR #15764:
URL: https://github.com/apache/tvm/pull/15764#issuecomment-1734393526
Thanks @Lunderberg for reviewing. Your changes have been addressed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on pull request #15764: [Disco] Add AllGather
Posted by "junrushao (via GitHub)" <gi...@apache.org>.
junrushao commented on PR #15764:
URL: https://github.com/apache/tvm/pull/15764#issuecomment-1722072728
CC @LeshengJin
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] farshidsp commented on a diff in pull request #15764: [Disco] Add AllGather
Posted by "farshidsp (via GitHub)" <gi...@apache.org>.
farshidsp commented on code in PR #15764:
URL: https://github.com/apache/tvm/pull/15764#discussion_r1333385279
##########
python/tvm/relax/op/ccl/ccl.py:
##########
@@ -44,6 +44,25 @@ def allreduce(x, op_type: str = "sum"): # pylint: disable=invalid-name
return _ffi_api.allreduce(x, op_type) # type: ignore # pylint: disable=no-member
+def allgather(x, num_workers: int): # pylint: disable=invalid-name
Review Comment:
That make sense Eric. Thanks for your suggestion and detailed explanation!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] LeshengJin commented on a diff in pull request #15764: [Disco] Add AllGather
Posted by "LeshengJin (via GitHub)" <gi...@apache.org>.
LeshengJin commented on code in PR #15764:
URL: https://github.com/apache/tvm/pull/15764#discussion_r1336456333
##########
python/tvm/relax/op/ccl/ccl.py:
##########
@@ -44,6 +47,27 @@ def allreduce(x, op_type: str = "sum"): # pylint: disable=invalid-name
return _ffi_api.allreduce(x, op_type) # type: ignore # pylint: disable=no-member
+def allgather(x, num_workers: Union[int, PrimExpr, PrimValue]): # pylint: disable=invalid-name
+ """AllGather operator
+
+ Parameters
+ ----------
+ x : relax.Expr
+ The input tensor.
+
+ num_worker : Union[int, PrimExpr, PrimValue]
+ The number of workers, i.e. the number of parts the given buffer should be chunked into.
Review Comment:
```suggestion
The number of workers to gather data from.
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on a diff in pull request #15764: [Disco] Add AllGather
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #15764:
URL: https://github.com/apache/tvm/pull/15764#discussion_r1335901030
##########
include/tvm/relax/attrs/ccl.h:
##########
@@ -40,6 +40,14 @@ struct AllReduceAttrs : public tvm::AttrsNode<AllReduceAttrs> {
}
}; // struct AllReduceAttrs
+/*! \brief Attributes used in allgather operators */
Review Comment:
Since the `AllGatherAttrs` is no longer used, let's remove its definition as well.
##########
python/tvm/relax/transform/legalize_ops/ccl.py:
##########
@@ -46,6 +46,30 @@ def _allreduce(_bb: BlockBuilder, call: Call) -> Expr:
)
+@register_legalize("relax.ccl.allgather")
+def _allgather(_bb: BlockBuilder, call: Call) -> Expr:
+ output_shape = []
+ assert isinstance(
+ call.args[0].struct_info, TensorStructInfo
+ ), "The input struct info of allgather should be TensorStructInfo."
+ assert isinstance(call.args[0].struct_info.shape.struct_info, ShapeStructInfo)
+ arg_shape = call.args[0].struct_info.shape.struct_info
+ for i, shape_value in enumerate(arg_shape.values):
+ if i == 0:
+ output_shape.append(tir.multiply(shape_value, call.args[1].value))
Review Comment:
Nit: The explicit call to `tir.multiply` can be replaced by `shape_value * call.args[1].value`, using operator overloading to delegate to `PrimExpr.__mul__`. (Whether this is good or bad practice is a topic of a great many flamewars, so I tend to stick with whatever an existing codebase uses.)
##########
python/tvm/relax/transform/legalize_ops/ccl.py:
##########
@@ -46,6 +46,30 @@ def _allreduce(_bb: BlockBuilder, call: Call) -> Expr:
)
+@register_legalize("relax.ccl.allgather")
+def _allgather(_bb: BlockBuilder, call: Call) -> Expr:
+ output_shape = []
+ assert isinstance(
+ call.args[0].struct_info, TensorStructInfo
Review Comment:
Since this expression occurs many times throughout the function, can we define `arg_sinfo = call.args[0].struct_info`?
##########
src/relax/op/ccl/ccl.cc:
##########
@@ -50,6 +50,44 @@ TVM_REGISTER_OP("relax.ccl.allreduce")
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<Bool>("FPurity", Bool(true));
+/* relax.ccl.allgather */
+TVM_REGISTER_NODE_TYPE(AllGatherAttrs);
Review Comment:
As `AllGatherAttrs` is no longer used, this `TVM_REGISTER_NODE_TYPE` declaration can be removed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] farshidsp commented on a diff in pull request #15764: [Disco] Add AllGather
Posted by "farshidsp (via GitHub)" <gi...@apache.org>.
farshidsp commented on code in PR #15764:
URL: https://github.com/apache/tvm/pull/15764#discussion_r1336337450
##########
python/tvm/relax/transform/legalize_ops/ccl.py:
##########
@@ -46,6 +46,30 @@ def _allreduce(_bb: BlockBuilder, call: Call) -> Expr:
)
+@register_legalize("relax.ccl.allgather")
+def _allgather(_bb: BlockBuilder, call: Call) -> Expr:
+ output_shape = []
+ assert isinstance(
+ call.args[0].struct_info, TensorStructInfo
+ ), "The input struct info of allgather should be TensorStructInfo."
+ assert isinstance(call.args[0].struct_info.shape.struct_info, ShapeStructInfo)
+ arg_shape = call.args[0].struct_info.shape.struct_info
+ for i, shape_value in enumerate(arg_shape.values):
+ if i == 0:
+ output_shape.append(tir.multiply(shape_value, call.args[1].value))
Review Comment:
That make sense. Resolved.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] farshidsp commented on a diff in pull request #15764: [Disco] Add AllGather
Posted by "farshidsp (via GitHub)" <gi...@apache.org>.
farshidsp commented on code in PR #15764:
URL: https://github.com/apache/tvm/pull/15764#discussion_r1333393638
##########
src/runtime/disco/builtin.cc:
##########
@@ -84,6 +84,8 @@ void AllReduce(NDArray send, ReduceKind reduce_kind, NDArray recv) {
GetCCLFunc("allreduce")(send, static_cast<int>(reduce_kind), recv);
}
+void AllGathere(NDArray send, NDArray recv) { GetCCLFunc("allgather")(send, recv); }
Review Comment:
Thanks for catching this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on a diff in pull request #15764: [Disco] Add AllGather
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #15764:
URL: https://github.com/apache/tvm/pull/15764#discussion_r1332182751
##########
src/relax/op/ccl/ccl.cc:
##########
@@ -50,6 +50,44 @@ TVM_REGISTER_OP("relax.ccl.allreduce")
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<Bool>("FPurity", Bool(true));
+/* relax.ccl.allgather */
+TVM_REGISTER_NODE_TYPE(AllGatherAttrs);
+
+Expr allgather(Expr x, int num_workers) {
Review Comment:
Here, this would be `Expr allgather(Expr x, Expr num_workers)`.
##########
src/runtime/disco/builtin.h:
##########
@@ -52,6 +52,12 @@ NDArray DiscoEmptyNDArray(ShapeTuple shape, DataType dtype, Device device);
* \return The outcome of allreduce
*/
void AllReduce(NDArray send, ReduceKind reduce_kind, NDArray recv);
+/*!
+ * \brief Perform an allgather operation using the underlying communication library
+ * \param send The array send to perform allgather on
+ * \return The outcome of allgather
+ */
+void AllGather(NDArray send, Optional<NDArray> recv);
Review Comment:
The signature here doesn't match the signature of the implementation. In the implementation, the second parameter is `NDArray recv` and not `Optional<NDArray> recv`.
##########
src/relax/op/ccl/ccl.cc:
##########
@@ -50,6 +50,44 @@ TVM_REGISTER_OP("relax.ccl.allreduce")
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<Bool>("FPurity", Bool(true));
+/* relax.ccl.allgather */
+TVM_REGISTER_NODE_TYPE(AllGatherAttrs);
+
+Expr allgather(Expr x, int num_workers) {
+ ObjectPtr<AllGatherAttrs> attrs = make_object<AllGatherAttrs>();
+ attrs->num_workers = std::move(num_workers);
+ static const Op& op = Op::Get("relax.ccl.allgather");
+ return Call(op, {std::move(x)}, Attrs{attrs}, {});
+}
+
+TVM_REGISTER_GLOBAL("relax.op.ccl.allgather").set_body_typed(allgather);
+
+StructInfo InferStructInfoAllGather(const Call& call, const BlockBuilder& ctx) {
+ TensorStructInfo input_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
+ DataType output_dtype = input_sinfo->dtype;
+
+ const auto* attrs = call->attrs.as<AllGatherAttrs>();
+ int num_workers = attrs->num_workers;
+ auto input_shape = input_sinfo->GetShape();
+ if (!input_shape.defined()) {
+ return input_sinfo;
+ }
+ Array<PrimExpr> output_shape = input_shape.value();
+ output_shape.Set(0, floor(output_shape[0] * num_workers));
+ if (input_sinfo->vdevice.defined()) {
Review Comment:
Instead of duplicating the entire `return TensorStructInfo(...)` call, we can make a local variable. Otherwise, readers would need to inspect the
```c++
VDevice vdevice;
if(input_sinfo->vdevice.defined()) {
vdevice = input_sinfo->vdevice.value();
}
return TensorStructInfo(ShapeExpr(output_shape), output_dtype, vdevice);
```
##########
python/tvm/relax/op/ccl/ccl.py:
##########
@@ -44,6 +44,25 @@ def allreduce(x, op_type: str = "sum"): # pylint: disable=invalid-name
return _ffi_api.allreduce(x, op_type) # type: ignore # pylint: disable=no-member
+def allgather(x, num_workers: int): # pylint: disable=invalid-name
Review Comment:
I like the inclusion of the `num_workers` here, as it allows shape propagation across the `allgather`. Instead of an integer, can we make it be a `relax.PrimValue`? That would (1) allow the `num_workers` to be a symbolic variable for earlier stages of lowering, (2) play nicer with the relax pattern matcher, and (3) remove the need for the `AllGatherAttrs`.
```python
def allgather(x, num_workers: Union[int, PrimExpr, relax.PrimValue]):
if not isinstance(num_workers, relax.PrimValue):
num_workers = relax.PrimValue(num_workers)
...
```
##########
src/relax/op/ccl/ccl.cc:
##########
@@ -50,6 +50,44 @@ TVM_REGISTER_OP("relax.ccl.allreduce")
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutUnaryEwise)
.set_attr<Bool>("FPurity", Bool(true));
+/* relax.ccl.allgather */
+TVM_REGISTER_NODE_TYPE(AllGatherAttrs);
+
+Expr allgather(Expr x, int num_workers) {
+ ObjectPtr<AllGatherAttrs> attrs = make_object<AllGatherAttrs>();
+ attrs->num_workers = std::move(num_workers);
+ static const Op& op = Op::Get("relax.ccl.allgather");
+ return Call(op, {std::move(x)}, Attrs{attrs}, {});
+}
+
+TVM_REGISTER_GLOBAL("relax.op.ccl.allgather").set_body_typed(allgather);
+
+StructInfo InferStructInfoAllGather(const Call& call, const BlockBuilder& ctx) {
+ TensorStructInfo input_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
Review Comment:
With `num_workers` as a `PrimValue`, this would instead be
```c++
CHECK_EQ(call->args.size(), 2);
auto input_sinfo = Downcast<TensorStructInfo>(call->args[0]->struct_info_);
auto num_workers_sinfo = Downcast<PrimStructInfo>(call->args[1]->struct_info_);
auto num_workers = num_workers_sinfo->value;
```
##########
src/runtime/disco/builtin.cc:
##########
@@ -84,6 +84,8 @@ void AllReduce(NDArray send, ReduceKind reduce_kind, NDArray recv) {
GetCCLFunc("allreduce")(send, static_cast<int>(reduce_kind), recv);
}
+void AllGathere(NDArray send, NDArray recv) { GetCCLFunc("allgather")(send, recv); }
Review Comment:
Typo: This should be `AllGather` instead of `Allgathere`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on pull request #15764: [Disco] Add AllGather
Posted by "junrushao (via GitHub)" <gi...@apache.org>.
junrushao commented on PR #15764:
URL: https://github.com/apache/tvm/pull/15764#issuecomment-1730016305
CC @LeshengJin could you please review this PR?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org