You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by ho...@apache.org on 2023/09/14 17:52:08 UTC
[tvm] branch unity updated (bd51b5b592 -> 0626ab7b16)
This is an automated email from the ASF dual-hosted git repository.
hongyij pushed a change to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
from bd51b5b592 [Unity] [Bugfix] Fix KeyError: 'padding' in _avg_pool2d implementation (#15734)
new 1ab89c07e4 change stream
new 0626ab7b16 update test
The 2 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
src/runtime/disco/nccl/nccl.cc | 18 ++++++++++++++++--
tests/python/disco/test_nccl.py | 11 +++--------
2 files changed, 19 insertions(+), 10 deletions(-)
[tvm] 02/02: update test
Posted by ho...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
hongyij pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 0626ab7b1630eca84ff6ffb852fb2c0040f6bbc8
Author: Hongyi Jin <ji...@gmail.com>
AuthorDate: Thu Sep 14 13:51:49 2023 -0400
update test
---
tests/python/disco/test_nccl.py | 11 +++--------
1 file changed, 3 insertions(+), 8 deletions(-)
diff --git a/tests/python/disco/test_nccl.py b/tests/python/disco/test_nccl.py
index f0f949ab80..e86c973fc2 100644
--- a/tests/python/disco/test_nccl.py
+++ b/tests/python/disco/test_nccl.py
@@ -22,6 +22,7 @@ import numpy as np
import pytest
import tvm
+import tvm.testing
from tvm import dlight as dl
from tvm import relax as rx
from tvm.runtime import disco as di
@@ -103,7 +104,7 @@ def test_scatter(session_kind):
@pytest.mark.parametrize("session_kind", _all_session_kinds)
def test_gather(session_kind):
- devices = [1, 2]
+ devices = [0, 1]
sess = session_kind(num_workers=len(devices))
sess.init_ccl("nccl", *devices)
@@ -376,10 +377,4 @@ def test_attention(session_kind): # pylint: disable=too-many-locals,too-many-st
if __name__ == "__main__":
- test_init(di.ProcessSession)
- test_allreduce(di.ProcessSession)
- test_broadcast_from_worker0(di.ProcessSession)
- test_scatter(di.ProcessSession)
- test_gather(di.ProcessSession)
- test_mlp(di.ProcessSession)
- test_attention(di.ProcessSession)
+ tvm.testing.main()
[tvm] 01/02: change stream
Posted by ho...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
hongyij pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 1ab89c07e423b20e8807d82f9435adf0c6bdaa6b
Author: Hongyi Jin <ji...@gmail.com>
AuthorDate: Thu Sep 14 13:39:18 2023 -0400
change stream
---
src/runtime/disco/nccl/nccl.cc | 18 ++++++++++++++++--
1 file changed, 16 insertions(+), 2 deletions(-)
diff --git a/src/runtime/disco/nccl/nccl.cc b/src/runtime/disco/nccl/nccl.cc
index e404e3c2bb..7ac8341eca 100644
--- a/src/runtime/disco/nccl/nccl.cc
+++ b/src/runtime/disco/nccl/nccl.cc
@@ -76,7 +76,6 @@ void InitCCLPerWorker(ShapeTuple device_ids, std::string unique_id_bytes) {
CUDA_CALL(cudaSetDevice(device_id));
CUDA_CALL(cudaStreamCreate(&ctx->stream));
Device device{DLDeviceType::kDLCUDA, device_id};
- DeviceAPI::Get(device)->SetStream(device, ctx->stream);
worker->default_device = device;
worker->ccl = "nccl";
ctx->worker = worker;
@@ -91,9 +90,12 @@ void AllReduce(NDArray send, ReduceKind reduce_kind, NDArray recv) {
NCCLThreadLocalContext* ctx = NCCLThreadLocalContext::Get();
ShapeTuple shape = send.Shape();
int64_t numel = shape->Product();
+ Device device = ctx->worker->default_device;
+ DeviceAPI::Get(device)->SyncStreamFromTo(device, nullptr, ctx->stream);
NCCL_CALL(ncclAllReduce(send->data, recv->data, numel,
/*datatype=*/AsNCCLDataType(DataType(send->dtype)),
/*op=*/AsNCCLRedOp(reduce_kind), ctx->comm, ctx->stream));
+ DeviceAPI::Get(device)->SyncStreamFromTo(device, ctx->stream, nullptr);
}
void BroadcastFromWorker0(NDArray send, NDArray recv) {
@@ -101,9 +103,12 @@ void BroadcastFromWorker0(NDArray send, NDArray recv) {
ICHECK(send.Shape()->Product() == recv.Shape()->Product());
ShapeTuple shape = send.Shape();
int64_t numel = shape->Product();
+ Device device = ctx->worker->default_device;
+ DeviceAPI::Get(device)->SyncStreamFromTo(device, nullptr, ctx->stream);
NCCL_CALL(ncclBroadcast(send->data, recv->data, numel,
/*datatype=*/AsNCCLDataType(DataType(send->dtype)),
/*root=*/0, ctx->comm, ctx->stream));
+ DeviceAPI::Get(device)->SyncStreamFromTo(device, ctx->stream, nullptr);
}
void ScatterFromWorker0(Optional<NDArray> send, NDArray recv) {
@@ -111,6 +116,8 @@ void ScatterFromWorker0(Optional<NDArray> send, NDArray recv) {
NCCLThreadLocalContext* ctx = NCCLThreadLocalContext::Get();
int worker_id = ctx->worker->worker_id;
int num_workers = ctx->worker->num_workers;
+ Device device = ctx->worker->default_device;
+ DeviceAPI::Get(device)->SyncStreamFromTo(device, nullptr, ctx->stream);
if (worker_id == 0) {
CHECK(send.defined()) << "ValueError: buffer `send` must be provided when worker_id == 0.";
NDArray buffer = send.value();
@@ -144,6 +151,7 @@ void ScatterFromWorker0(Optional<NDArray> send, NDArray recv) {
DataType dtype(recv->dtype);
NCCL_CALL(ncclRecv(recv->data, numel, AsNCCLDataType(dtype), 0, ctx->comm, ctx->stream));
NCCL_CALL(ncclGroupEnd());
+ DeviceAPI::Get(device)->SyncStreamFromTo(device, ctx->stream, nullptr);
}
void GatherToWorker0(NDArray send, Optional<NDArray> recv) {
@@ -151,6 +159,8 @@ void GatherToWorker0(NDArray send, Optional<NDArray> recv) {
NCCLThreadLocalContext* ctx = NCCLThreadLocalContext::Get();
int worker_id = ctx->worker->worker_id;
int num_workers = ctx->worker->num_workers;
+ Device device = ctx->worker->default_device;
+ DeviceAPI::Get(device)->SyncStreamFromTo(device, nullptr, ctx->stream);
if (worker_id == 0) {
CHECK(recv.defined()) << "ValueError: buffer `recv` must be provided when worker_id == 0.";
NDArray buffer = recv.value();
@@ -184,22 +194,26 @@ void GatherToWorker0(NDArray send, Optional<NDArray> recv) {
DataType dtype(send->dtype);
NCCL_CALL(ncclSend(send->data, numel, AsNCCLDataType(dtype), 0, ctx->comm, ctx->stream));
NCCL_CALL(ncclGroupEnd());
+ DeviceAPI::Get(device)->SyncStreamFromTo(device, ctx->stream, nullptr);
}
void RecvFromWorker0(NDArray buffer) {
NCCLThreadLocalContext* ctx = NCCLThreadLocalContext::Get();
CHECK_NE(ctx->worker->worker_id, 0)
<< "ValueError: Worker 0 is not allowed to call RecvFromWorker0.";
+ Device device = ctx->worker->default_device;
+ DeviceAPI::Get(device)->SyncStreamFromTo(device, nullptr, ctx->stream);
NCCL_CALL(ncclGroupStart());
NCCL_CALL(ncclRecv(buffer->data, buffer.Shape()->Product(), AsNCCLDataType(buffer.DataType()), 0,
ctx->comm, ctx->stream));
NCCL_CALL(ncclGroupEnd());
+ DeviceAPI::Get(device)->SyncStreamFromTo(device, ctx->stream, nullptr);
}
void SyncWorker() {
NCCLThreadLocalContext* ctx = NCCLThreadLocalContext::Get();
ICHECK(ctx->worker != nullptr);
- CUDA_CALL(cudaStreamSynchronize(ctx->stream));
+ CUDA_CALL(cudaStreamSynchronize(nullptr));
}
TVM_REGISTER_GLOBAL("runtime.disco.nccl.init_ccl").set_body_typed(InitCCL);