This is an automated email from the ASF dual-hosted git repository.
junrushao pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 35e8404f17 [Disco] Expose `DiscoWorker` and `ndarray_cache_support` in
header (#16153)
35e8404f17 is described below
commit 35e8404f17a2aae2e00110748ecce515395a5428
Author: Lesheng Jin <[email protected]>
AuthorDate: Sun Dec 10 18:25:40 2023 +0800
[Disco] Expose `DiscoWorker` and `ndarray_cache_support` in header (#16153)
---
{src => include/tvm}/runtime/disco/builtin.h | 38 ++++-
.../tvm/runtime/disco/disco_worker.h | 71 +--------
include/tvm/runtime/disco/session.h | 20 ++-
.../tvm}/runtime/relax_vm/ndarray_cache_support.h | 39 ++---
python/tvm/relax/frontend/nn/core.py | 2 +
python/tvm/relax/frontend/nn/op.py | 39 +++++
python/tvm/runtime/disco/process_pool.py | 16 +-
python/tvm/runtime/disco/session.py | 3 +-
src/runtime/disco/bcast_session.h | 3 +-
src/runtime/disco/builtin.cc | 5 +-
src/runtime/disco/{worker.cc => disco_worker.cc} | 8 +-
src/runtime/disco/disco_worker_thread.h | 83 +++++++++++
src/runtime/disco/loader.cc | 89 ++++++++++--
src/runtime/disco/nccl/nccl.cc | 5 +-
src/runtime/disco/process_session.cc | 9 +-
src/runtime/disco/session.cc | 3 +-
src/runtime/disco/threaded_session.cc | 3 +-
src/runtime/disco/utils.h | 32 +---
src/runtime/relax_vm/ndarray_cache_support.cc | 161 ++++++++++-----------
19 files changed, 379 insertions(+), 250 deletions(-)
diff --git a/src/runtime/disco/builtin.h b/include/tvm/runtime/disco/builtin.h
similarity index 81%
rename from src/runtime/disco/builtin.h
rename to include/tvm/runtime/disco/builtin.h
index cfbf2e2477..3847aef3f2 100644
--- a/src/runtime/disco/builtin.h
+++ b/include/tvm/runtime/disco/builtin.h
@@ -25,11 +25,37 @@
#include <string>
-#include "./utils.h"
-
namespace tvm {
namespace runtime {
+/*!
+ * \brief Possible kinds of reduction operations.
+ */
+enum class ReduceKind : int32_t {
+ kSum = 0,
+ kProd = 1,
+ kMin = 2,
+ kMax = 3,
+ kAvg = 4,
+};
+
+/*! \brief Converts `ReduceKind` to string */
+inline std::string ReduceKind2String(ReduceKind kind) {
+ switch (kind) {
+ case ReduceKind::kSum:
+ return "kSum";
+ case ReduceKind::kProd:
+ return "kProd";
+ case ReduceKind::kMin:
+ return "kMin";
+ case ReduceKind::kMax:
+ return "kMax";
+ case ReduceKind::kAvg:
+ return "kAvg";
+ }
+ LOG(FATAL) << "ValueError: Unknown ReduceKind: " << static_cast<int>(kind);
+}
+
/*!
* \brief Load a runtime Module, then create and initialize a RelaxVM
* \param path The path to the runtime Module (a DSO file) to be loaded
@@ -49,19 +75,19 @@ NDArray DiscoEmptyNDArray(ShapeTuple shape, DataType dtype,
Device device);
* \brief Perform an allreduce operation using the underlying communication
library
* \param send The array send to perform allreduce on
* \param reduce_kind The kind of reduction operation (e.g. sum, avg, min, max)
- * \return The outcome of allreduce
+ * \param recv The array receives 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
+ * \param recv The array receives the outcome of allgather
*/
void AllGather(NDArray send, NDArray recv);
/*!
* \brief Perform a broadcast operation from worker-0
- * \param buffer The buffer to be broadcasted
- * \return The result buffer
+ * \param send The buffer to be broadcasted
+ * \param recv The buffer receives the broadcasted array
*/
void BroadcastFromWorker0(NDArray send, NDArray recv);
/*!
diff --git a/src/runtime/disco/worker.h
b/include/tvm/runtime/disco/disco_worker.h
similarity index 60%
rename from src/runtime/disco/worker.h
rename to include/tvm/runtime/disco/disco_worker.h
index e948fa1668..0c666150d4 100644
--- a/src/runtime/disco/worker.h
+++ b/include/tvm/runtime/disco/disco_worker.h
@@ -17,44 +17,22 @@
* under the License.
*/
/*!
- * \file worker.h
+ * \file disco_worker.h
* \brief This file defines a worker in Disco. A worker can be launched in a
separate thread or
* process as long as the channel supports bi-directional communication
in-between the worker and
* the controler.
*/
-#ifndef TVM_RUNTIME_DISCO_WORKER_H_
-#define TVM_RUNTIME_DISCO_WORKER_H_
+#ifndef TVM_RUNTIME_DISCO_DISCO_WORKER_H_
+#define TVM_RUNTIME_DISCO_DISCO_WORKER_H_
#include <tvm/runtime/disco/session.h>
#include <tvm/runtime/packed_func.h>
-#include <atomic>
-#include <condition_variable>
-#include <memory>
-#include <mutex>
-#include <queue>
-#include <thread>
-#include <utility>
#include <vector>
namespace tvm {
namespace runtime {
-/*!
- * \brief A special communication channel between controler and worker-0,
- * assuming they are always collocated in the same process.
- */
-class WorkerZeroData {
- public:
- /*!
- * \brief The host-side arrays to passed to worker-0 for special uses, for
example,
- * copy-to-worker0 and copy-from-worker0
- */
- std::queue<NDArray> host_arrays;
- /*! \brief The mutex that guards `host_arrays` */
- std::mutex queue_mutex_;
-};
-
/*!
* \brief A worker in Disco. It takes a channel to communication with the
controler.
* The worker can be run in a separate thread or process as long as the
channel supports
@@ -65,6 +43,7 @@ class DiscoWorker {
/*!
* \brief Construct a worker.
* \param worker_id The id of the worker.
+ * \param num_workers The number of the workers.
* \param worker_zero_data The data shared between worker-0 and the
controler. It's a nullptr if
* the worker is not worker-0.
* \param channel The communication channel between the worker and the
controler.
@@ -111,46 +90,6 @@ class DiscoWorker {
friend struct DiscoWorker::Impl;
};
-/*!
- * \brief A worker thread in Disco, which upon creation, launches a new thread
to run the
- * DiscoWorker.
- * \sa DiscoWorker
- */
-class DiscoWorkerThread {
- public:
- /*!
- * \brief Construct a worker thread.
- * \param worker_id The id of the worker.
- * \param num_workers The total number of workers.
- * \param worker_zero_data_ The data shared between worker-0 and the
controler. It's a nullptr if
- * the worker is not worker-0.
- */
- explicit DiscoWorkerThread(int worker_id, int num_workers, WorkerZeroData*
worker_zero_data_);
-
- /*! \brief Move constructor. */
- explicit DiscoWorkerThread(DiscoWorkerThread&& other)
- : channel(std::move(other.channel)),
- worker(std::move(other.worker)),
- thread(std::move(other.thread)) {}
-
- /*! \brief Copy constructor is disabled */
- DiscoWorkerThread(const DiscoWorkerThread& other) = delete;
-
- /*! \brief Destructor that joins the thread before destruction */
- ~DiscoWorkerThread() {
- if (this->thread != nullptr) {
- this->thread->join();
- }
- }
-
- /*! \brief The communication channel between the controler and the worker */
- std::unique_ptr<DiscoChannel> channel;
- /*! \brief The worker whose internal state is visible to the controler */
- std::unique_ptr<DiscoWorker> worker;
- /*! \brief The thread that runs the worker's main loop. */
- std::unique_ptr<std::thread> thread;
-};
-
} // namespace runtime
} // namespace tvm
-#endif // TVM_RUNTIME_DISCO_WORKER_H_
+#endif // TVM_RUNTIME_DISCO_DISCO_WORKER_H_
diff --git a/include/tvm/runtime/disco/session.h
b/include/tvm/runtime/disco/session.h
index 3de519e9cc..5e745166b0 100644
--- a/include/tvm/runtime/disco/session.h
+++ b/include/tvm/runtime/disco/session.h
@@ -76,6 +76,7 @@
#include <tvm/runtime/object.h>
#include <tvm/runtime/packed_func.h>
+#include <queue>
#include <string>
#include <utility>
@@ -270,10 +271,12 @@ class Session : public ObjectRef {
* and returns a PackedFunc, which takes an integer `worker_id` as the input
and returns None.
* When `worker-id` is 0, it shuts down the process pool; Otherwise, it
retursn a tuple
* (read_fd, writefd) used to communicate with the corresponding worker.
+ * \param entrypoint The entrypoint of DiscoWorker main worker function.
* \note Worker-0 is always co-located with the controler as a separate
thread, and therefore
* worker-0 does not exist in the process pool.
*/
- TVM_DLL static Session ProcessSession(int num_workers, String
process_pool_creator);
+ TVM_DLL static Session ProcessSession(int num_workers, String
process_pool_creator,
+ String entrypoint);
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(Session, ObjectRef,
SessionObj);
};
@@ -294,6 +297,21 @@ class DiscoChannel {
virtual TVMArgs RecvReply() = 0;
};
+/*!
+ * \brief A special communication channel between controler and worker-0,
+ * assuming they are always collocated in the same process.
+ */
+class WorkerZeroData {
+ public:
+ /*!
+ * \brief The host-side arrays to passed to worker-0 for special uses, for
example,
+ * copy-to-worker0 and copy-from-worker0
+ */
+ std::queue<NDArray> host_arrays;
+ /*! \brief The mutex that guards `host_arrays` */
+ std::mutex queue_mutex_;
+};
+
// Implementation details
DRefObj::~DRefObj() {
diff --git a/src/runtime/relax_vm/ndarray_cache_support.h
b/include/tvm/runtime/relax_vm/ndarray_cache_support.h
similarity index 77%
rename from src/runtime/relax_vm/ndarray_cache_support.h
rename to include/tvm/runtime/relax_vm/ndarray_cache_support.h
index c1beb5a946..3d8b639ee4 100644
--- a/src/runtime/relax_vm/ndarray_cache_support.h
+++ b/include/tvm/runtime/relax_vm/ndarray_cache_support.h
@@ -42,10 +42,11 @@ struct NDArrayCacheMetadata {
* \brief Load the parameter from raw data.
* \param device The device to load the parameter onto.
* \param raw_data The raw data stream
- * \param f_load The function to load the parameter from raw data.
+ * \param staging_buffer The buffer to be used to avoid extra OpenCL
copies. Pass in a nullptr
+ * in other cases
*/
NDArray Load(Device device, const std::string* raw_data,
- std::function<void(NDArray, const void*, int64_t)> f_load)
const;
+ Optional<NDArray>* staging_buffer = nullptr) const;
/*! \brief Name of the parameter */
std::string name;
@@ -61,6 +62,12 @@ struct NDArrayCacheMetadata {
int64_t byte_offset;
};
+ /*! \brief Load a FileRecord into memory */
+ Array<NDArray> Load(Device device, //
+ const std::string& path_prefix, //
+ std::string* raw_data_buffer, //
+ Optional<NDArray>* staging_buffer = nullptr) const;
+
/*! \brief Relative path to the bin file */
std::string data_path;
/*! \brief Format of the file */
@@ -75,34 +82,12 @@ struct NDArrayCacheMetadata {
/*! \brief The path to the `ndarray-cache.json` file */
std::string path;
- /*! \brief Load the metadata from a specific path */
+ /*! \brief Load the metadata from a specific directory */
+ static NDArrayCacheMetadata Load(const std::string& path);
+ /*! \brief Load the metadata from a given JSON string */
static NDArrayCacheMetadata LoadFromStr(const std::string& json_str, const
std::string& path);
};
-/*!
- * \brief Information of sharding function,
- * including the shard function name and extra parameters.
- */
-struct ShardInfo {
- struct TensorInfo {
- ShapeTuple shape;
- DataType dtype;
- };
- struct ShardFunc {
- std::string name;
- TensorInfo output_info;
- std::vector<int64_t> params;
- };
- std::vector<ShardFunc> funcs;
-};
-
-/*!
- * \brief Load the shard information from dist
- * \param path Path to the file to be loaded
- * \return Mapping from parameter name to its shard dim
- */
-std::unordered_map<std::string, ShardInfo> LoadShardInfoFromStr(const
std::string& json_str);
-
} // namespace relax_vm
} // namespace runtime
} // namespace tvm
diff --git a/python/tvm/relax/frontend/nn/core.py
b/python/tvm/relax/frontend/nn/core.py
index ba48b449fc..c7d745c721 100644
--- a/python/tvm/relax/frontend/nn/core.py
+++ b/python/tvm/relax/frontend/nn/core.py
@@ -176,6 +176,7 @@ class Parameter(Tensor):
"""
_data: Optional[NDArray]
+ attrs: Dict[str, Any]
def __init__(
self,
@@ -196,6 +197,7 @@ class Parameter(Tensor):
dtype = get_default_dtype()
super().__init__(_expr=_tensor_placeholder("param", shape,
dtype=dtype)._expr)
self._data = None
+ self.attrs = OrderedDict()
@property
def data(self) -> Optional[NDArray]:
diff --git a/python/tvm/relax/frontend/nn/op.py
b/python/tvm/relax/frontend/nn/op.py
index b95ceac4ed..76af0044c3 100644
--- a/python/tvm/relax/frontend/nn/op.py
+++ b/python/tvm/relax/frontend/nn/op.py
@@ -1372,6 +1372,45 @@ def interpolate(
)
+def ccl_allreduce(x: Tensor, op_type: str = "sum", name="ccl_allreduce"):
+ """CCL Allreduce operator
+
+ Parameters
+ ----------
+ x : Tensor
+ The input tensor.
+ op_type: str
+ The type of reduction operation to be applied to the input data.
+ Now "sum", "prod", "min", "max" and "avg" are supported.
+ name : str
+ Name hint for this operation.
+
+ Returns
+ -------
+ result : Tensor
+ The result tensor of allreduce.
+ """
+ return _wrap_nested(_op.ccl.allreduce(x._expr, op_type), name)
+
+
+def ccl_broadcast_from_worker0(x: Tensor, name="broadcast_from_worker"):
+ """Broadcast data from worker-0 to all other workers.
+
+ Parameters
+ ----------
+ x : Tensor
+ The tensor to be broadcast.
+ name : str
+ Name hint for this operation.
+
+ Returns
+ -------
+ result : Tensor
+ The same tensor, which has been broadcast to all other workers.
+ """
+ return _wrap_nested(_op.ccl.broadcast_from_worker0(x._expr), name)
+
+
def tensor_expr_op(
tensor_expr_func: Callable,
name_hint: str,
diff --git a/python/tvm/runtime/disco/process_pool.py
b/python/tvm/runtime/disco/process_pool.py
index 836744dba6..e91d855953 100644
--- a/python/tvm/runtime/disco/process_pool.py
+++ b/python/tvm/runtime/disco/process_pool.py
@@ -47,9 +47,17 @@ class DiscoPopenWorker:
The standard error streams handler specified for the popen process.
"""
- def __init__(self, worker_id: int, num_workers: int, stdout=None,
stderr=None):
+ def __init__( # pylint: disable=too-many-arguments
+ self,
+ worker_id: int,
+ num_workers: int,
+ entrypoint: str = "tvm.exec.disco_worker",
+ stdout=None,
+ stderr=None,
+ ):
self.worker_id = worker_id
self.num_workers = num_workers
+ self.entrypoint = entrypoint
self._proc = None
self._stdout = stdout
self._stderr = stderr
@@ -109,7 +117,7 @@ class DiscoPopenWorker:
cmd = [
sys.executable,
"-m",
- "tvm.exec.disco_worker",
+ self.entrypoint,
str(self.worker_id),
str(self.num_workers),
]
@@ -164,9 +172,9 @@ def _kill_child_processes(pid):
@register_func("runtime.disco.create_process_pool")
-def _create_process_pool(num_workers: int):
+def _create_process_pool(num_workers: int, entrypoint: str):
"""Create a process pool where the workers' are [1, num_workers)."""
- pool = [DiscoPopenWorker(i, num_workers) for i in range(1, num_workers)]
+ pool = [DiscoPopenWorker(i, num_workers, entrypoint) for i in range(1,
num_workers)]
def result_func(worker_id: int):
nonlocal pool
diff --git a/python/tvm/runtime/disco/session.py
b/python/tvm/runtime/disco/session.py
index a7defd96cd..8d0ff57a32 100644
--- a/python/tvm/runtime/disco/session.py
+++ b/python/tvm/runtime/disco/session.py
@@ -360,10 +360,11 @@ class ThreadedSession(Session):
class ProcessSession(Session):
"""A Disco session backed by pipe-based multi-processing."""
- def __init__(self, num_workers: int) -> None:
+ def __init__(self, num_workers: int, entrypoint: str) -> None:
self.__init_handle_by_constructor__(
_ffi_api.SessionProcess, # type: ignore # pylint:
disable=no-member
num_workers,
+ entrypoint,
"runtime.disco.create_process_pool",
)
diff --git a/src/runtime/disco/bcast_session.h
b/src/runtime/disco/bcast_session.h
index 772e14c087..1a4df634b7 100644
--- a/src/runtime/disco/bcast_session.h
+++ b/src/runtime/disco/bcast_session.h
@@ -19,13 +19,12 @@
#ifndef TVM_RUNTIME_DISCO_BCAST_SESSION_H_
#define TVM_RUNTIME_DISCO_BCAST_SESSION_H_
+#include <tvm/runtime/disco/disco_worker.h>
#include <tvm/runtime/disco/session.h>
#include <string>
#include <vector>
-#include "./worker.h"
-
namespace tvm {
namespace runtime {
diff --git a/src/runtime/disco/builtin.cc b/src/runtime/disco/builtin.cc
index 514d633fa6..51fe4c13fc 100644
--- a/src/runtime/disco/builtin.cc
+++ b/src/runtime/disco/builtin.cc
@@ -16,10 +16,10 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include "./builtin.h"
-
#include <dlpack/dlpack.h>
#include <tvm/runtime/container/shape_tuple.h>
+#include <tvm/runtime/disco/builtin.h>
+#include <tvm/runtime/disco/disco_worker.h>
#include <tvm/runtime/disco/session.h>
#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/registry.h>
@@ -28,7 +28,6 @@
#include <sstream>
#include "./utils.h"
-#include "./worker.h"
namespace tvm {
namespace runtime {
diff --git a/src/runtime/disco/worker.cc b/src/runtime/disco/disco_worker.cc
similarity index 98%
rename from src/runtime/disco/worker.cc
rename to src/runtime/disco/disco_worker.cc
index 9192215dda..d3c6d6a383 100644
--- a/src/runtime/disco/worker.cc
+++ b/src/runtime/disco/disco_worker.cc
@@ -16,17 +16,13 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include "./worker.h"
-
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/disco/builtin.h>
+#include <tvm/runtime/disco/disco_worker.h>
#include <tvm/runtime/disco/session.h>
#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/registry.h>
-#include <thread>
-
#include "../../support/process_id.h"
-#include "./builtin.h"
#include "./protocol.h"
namespace tvm {
diff --git a/src/runtime/disco/disco_worker_thread.h
b/src/runtime/disco/disco_worker_thread.h
new file mode 100644
index 0000000000..67742cdd04
--- /dev/null
+++ b/src/runtime/disco/disco_worker_thread.h
@@ -0,0 +1,83 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file disco_worker_thread.h
+ * \brief This file defines a worker in Disco. A worker can be launched in a
separate thread or
+ * process as long as the channel supports bi-directional communication
in-between the worker and
+ * the controler.
+ */
+#ifndef TVM_RUNTIME_DISCO_DISCO_WORKER_THREAD_H_
+#define TVM_RUNTIME_DISCO_DISCO_WORKER_THREAD_H_
+
+#include <tvm/runtime/disco/disco_worker.h>
+#include <tvm/runtime/disco/session.h>
+#include <tvm/runtime/packed_func.h>
+
+#include <memory>
+#include <thread>
+#include <utility>
+
+namespace tvm {
+namespace runtime {
+
+/*!
+ * \brief A worker thread in Disco, which upon creation, launches a new thread
to run the
+ * DiscoWorker.
+ * \sa DiscoWorker
+ */
+class DiscoWorkerThread {
+ public:
+ /*!
+ * \brief Construct a worker thread.
+ * \param worker_id The id of the worker.
+ * \param num_workers The total number of workers.
+ * \param worker_zero_data_ The data shared between worker-0 and the
controler. It's a nullptr if
+ * the worker is not worker-0.
+ * \note This method is implemented in threaded worker, because it depends
on creation of a
+ * sub-class of DiscoChannel, DiscoThreadChannel, which is hidden from the
public interface.
+ */
+ explicit DiscoWorkerThread(int worker_id, int num_workers, WorkerZeroData*
worker_zero_data_);
+
+ /*! \brief Move constructor. */
+ explicit DiscoWorkerThread(DiscoWorkerThread&& other)
+ : channel(std::move(other.channel)),
+ worker(std::move(other.worker)),
+ thread(std::move(other.thread)) {}
+
+ /*! \brief Copy constructor is disabled */
+ DiscoWorkerThread(const DiscoWorkerThread& other) = delete;
+
+ /*! \brief Destructor that joins the thread before destruction */
+ ~DiscoWorkerThread() {
+ if (this->thread != nullptr) {
+ this->thread->join();
+ }
+ }
+
+ /*! \brief The communication channel between the controler and the worker */
+ std::unique_ptr<DiscoChannel> channel;
+ /*! \brief The worker whose internal state is visible to the controler */
+ std::unique_ptr<DiscoWorker> worker;
+ /*! \brief The thread that runs the worker's main loop. */
+ std::unique_ptr<std::thread> thread;
+};
+
+} // namespace runtime
+} // namespace tvm
+#endif // TVM_RUNTIME_DISCO_DISCO_WORKER_THREAD_H_
diff --git a/src/runtime/disco/loader.cc b/src/runtime/disco/loader.cc
index c931baa942..7a5d978946 100644
--- a/src/runtime/disco/loader.cc
+++ b/src/runtime/disco/loader.cc
@@ -16,9 +16,16 @@
* specific language governing permissions and limitations
* under the License.
*/
+#define PICOJSON_USE_INT64
+#ifndef __STDC_FORMAT_MACROS
+#define __STDC_FORMAT_MACROS
+#endif
+#include <picojson.h>
#include <tvm/runtime/data_type.h>
+#include <tvm/runtime/disco/builtin.h>
#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/registry.h>
+#include <tvm/runtime/relax_vm/ndarray_cache_support.h>
#include <functional>
#include <numeric>
@@ -27,8 +34,6 @@
#include <vector>
#include "../file_utils.h"
-#include "../relax_vm/ndarray_cache_support.h"
-#include "./builtin.h"
#include "./utils.h"
namespace tvm {
@@ -37,7 +42,75 @@ namespace runtime {
using relax_vm::NDArrayCacheMetadata;
using FileRecord = NDArrayCacheMetadata::FileRecord;
using ParamRecord = NDArrayCacheMetadata::FileRecord::ParamRecord;
-using relax_vm::ShardInfo;
+
+struct ShardInfo {
+ struct TensorInfo {
+ ShapeTuple shape;
+ DataType dtype;
+ };
+ struct ShardFunc {
+ std::string name;
+ TensorInfo output_info;
+ std::vector<int64_t> params;
+ };
+ std::vector<ShardFunc> funcs;
+};
+
+template <typename ExpectedType>
+inline ExpectedType AsType(const picojson::value& json) {
+ ICHECK(json.is<ExpectedType>());
+ return json.get<ExpectedType>();
+}
+
+template <typename ValueType>
+inline ValueType GetValue(const picojson::object& json, const std::string&
key) {
+ return AsType<ValueType>(json.at(key));
+}
+
+std::unordered_map<std::string, ShardInfo> LoadShardInfoFromStr(const
std::string& json_str);
+ShardInfo::TensorInfo LoadTensorInfoFromJSON(const picojson::array&
json_tensor_info) {
+ CHECK_EQ(json_tensor_info.size(), 2) << "ValueError: Invalid tensor info
JSON";
+ picojson::array shape_json = AsType<picojson::array>(json_tensor_info[0]);
+ int ndim = shape_json.size();
+ std::vector<int64_t> shape;
+ shape.reserve(ndim);
+ for (int i = 0; i < ndim; ++i) {
+ shape.push_back(AsType<int64_t>(shape_json[i]));
+ }
+ std::string dtype = AsType<std::string>(json_tensor_info[1]);
+ return ShardInfo::TensorInfo{ShapeTuple(std::move(shape)),
DataType(String2DLDataType(dtype))};
+}
+
+ShardInfo::ShardFunc LoadShardFuncFromJSON(const picojson::array&
json_shard_func) {
+ int n = json_shard_func.size();
+ ShardInfo::ShardFunc shard_info;
+ shard_info.name = AsType<std::string>(json_shard_func[0]);
+ shard_info.output_info =
LoadTensorInfoFromJSON(AsType<picojson::array>(json_shard_func[1]));
+ shard_info.params.reserve(n - 2);
+ for (int i = 2; i < n; ++i) {
+ shard_info.params.push_back(AsType<int64_t>(json_shard_func[i]));
+ }
+ return shard_info;
+}
+
+std::unordered_map<std::string, ShardInfo> LoadShardInfoFromStr(const
std::string& json_str) {
+ picojson::value json_info;
+ picojson::parse(json_info, json_str);
+ picojson::object json_obj = AsType<picojson::object>(json_info);
+ std::unordered_map<std::string, ShardInfo> result;
+ for (auto kv : json_obj) {
+ std::string name = kv.first;
+ picojson::array json_shard_funcs = AsType<picojson::array>(kv.second);
+ ShardInfo info;
+ std::vector<ShardInfo::ShardFunc>& shard_funcs = info.funcs;
+ shard_funcs.reserve(json_shard_funcs.size());
+ for (const picojson::value& json_shard_func : json_shard_funcs) {
+
shard_funcs.push_back(LoadShardFuncFromJSON(AsType<picojson::array>(json_shard_func)));
+ }
+ result[name] = info;
+ }
+ return result;
+}
/*! \brief An object that helps to load parameters in shards. */
class ShardLoaderObj : public Object {
@@ -114,7 +187,7 @@ ObjectRef ShardLoaderObj::Create(const std::string&
path_to_metadata, const std:
n->metadata_ = NDArrayCacheMetadata::LoadFromStr(metadata, path_to_metadata);
n->current_file_ = nullptr;
n->param_info_.clear();
- std::unordered_map<std::string, ShardInfo> shards =
relax_vm::LoadShardInfoFromStr(shard_info);
+ std::unordered_map<std::string, ShardInfo> shards =
LoadShardInfoFromStr(shard_info);
for (const FileRecord& file_record : n->metadata_.records) {
for (const ParamRecord& param_record : file_record.records) {
const std::string& name = param_record.name;
@@ -181,9 +254,7 @@ NDArray ShardLoaderObj::LoadParamOnWorker0(int
weight_index) const {
std::string file_name = GetSiblingPath(this->metadata_.path,
file->data_path);
LoadBinaryFromFile(file_name, &this->current_file_stream_);
}
- return param->Load(
- device, &this->current_file_stream_,
- [](NDArray param, const void* data, size_t nbytes) {
param.CopyFromBytes(data, nbytes); });
+ return param->Load(device, &this->current_file_stream_);
};
if (worker_id == 0) {
@@ -230,9 +301,7 @@ NDArray ShardLoaderObj::LoadDirect(int weight_index) const {
std::string file_name = GetSiblingPath(this->metadata_.path,
file->data_path);
LoadBinaryFromFile(file_name, &this->current_file_stream_);
}
- return param->Load(
- device, &this->current_file_stream_,
- [](NDArray param, const void* data, size_t nbytes) {
param.CopyFromBytes(data, nbytes); });
+ return param->Load(device, &this->current_file_stream_);
}
NDArray ShardLoaderObj::Load(int weight_index) const {
diff --git a/src/runtime/disco/nccl/nccl.cc b/src/runtime/disco/nccl/nccl.cc
index e61306377f..61c307c673 100644
--- a/src/runtime/disco/nccl/nccl.cc
+++ b/src/runtime/disco/nccl/nccl.cc
@@ -19,6 +19,7 @@
#include <dlpack/dlpack.h>
#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/disco/builtin.h>
#include <tvm/runtime/disco/session.h>
#include <tvm/runtime/registry.h>
@@ -336,10 +337,10 @@ TVM_REGISTER_GLOBAL("runtime.disco." TVM_DISCO_CCL_NAME
".init_ccl_per_worker")
TVM_REGISTER_GLOBAL("runtime.disco." TVM_DISCO_CCL_NAME ".allreduce")
.set_body_typed([](NDArray send, int kind, NDArray recv) {
CHECK(0 <= kind && kind <= 4) << "ValueError: Unknown ReduceKind: " <<
kind;
- AllReduce(send, static_cast<ReduceKind>(kind), recv);
+ nccl::AllReduce(send, static_cast<ReduceKind>(kind), recv);
});
TVM_REGISTER_GLOBAL("runtime.disco." TVM_DISCO_CCL_NAME ".allgather")
- .set_body_typed([](NDArray send, NDArray recv) { AllGather(send, recv); });
+ .set_body_typed([](NDArray send, NDArray recv) { nccl::AllGather(send,
recv); });
TVM_REGISTER_GLOBAL("runtime.disco." TVM_DISCO_CCL_NAME
".broadcast_from_worker0")
.set_body_typed(BroadcastFromWorker0);
TVM_REGISTER_GLOBAL("runtime.disco." TVM_DISCO_CCL_NAME
".scatter_from_worker0")
diff --git a/src/runtime/disco/process_session.cc
b/src/runtime/disco/process_session.cc
index ece93d07bf..467a635181 100644
--- a/src/runtime/disco/process_session.cc
+++ b/src/runtime/disco/process_session.cc
@@ -16,6 +16,8 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/disco/disco_worker.h>
#include <tvm/runtime/object.h>
#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/registry.h>
@@ -28,9 +30,8 @@
#include "../../support/pipe.h"
#include "../minrpc/rpc_reference.h"
#include "./bcast_session.h"
+#include "./disco_worker_thread.h"
#include "./protocol.h"
-#include "./worker.h"
-#include "tvm/runtime/c_runtime_api.h"
namespace tvm {
namespace runtime {
@@ -191,11 +192,11 @@ class ProcessSessionObj final : public BcastSessionObj {
TVM_REGISTER_OBJECT_TYPE(DiscoDebugObject);
TVM_REGISTER_OBJECT_TYPE(ProcessSessionObj);
-Session Session::ProcessSession(int num_workers, String process_pool_creator) {
+Session Session::ProcessSession(int num_workers, String process_pool_creator,
String entrypoint) {
const PackedFunc* pf = Registry::Get(process_pool_creator);
CHECK(pf) << "ValueError: Cannot find function " << process_pool_creator
<< " in the registry. Please check if it is registered.";
- PackedFunc process_pool = (*pf)(num_workers);
+ PackedFunc process_pool = (*pf)(num_workers, entrypoint);
auto n = make_object<ProcessSessionObj>(num_workers, process_pool);
return Session(n);
}
diff --git a/src/runtime/disco/session.cc b/src/runtime/disco/session.cc
index c588af82f1..12339c4fa5 100644
--- a/src/runtime/disco/session.cc
+++ b/src/runtime/disco/session.cc
@@ -16,12 +16,11 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include <tvm/runtime/disco/disco_worker.h>
#include <tvm/runtime/disco/session.h>
#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/registry.h>
-#include "./worker.h"
-
namespace tvm {
namespace runtime {
diff --git a/src/runtime/disco/threaded_session.cc
b/src/runtime/disco/threaded_session.cc
index 349601fd03..985601aeb6 100644
--- a/src/runtime/disco/threaded_session.cc
+++ b/src/runtime/disco/threaded_session.cc
@@ -18,6 +18,7 @@
*/
#include <dmlc/io.h>
#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/disco/disco_worker.h>
#include <tvm/runtime/object.h>
#include <condition_variable>
@@ -30,8 +31,8 @@
#include "../../support/ring_buffer.h"
#include "../minrpc/rpc_reference.h"
#include "./bcast_session.h"
+#include "./disco_worker_thread.h"
#include "./protocol.h"
-#include "./worker.h"
namespace tvm {
namespace runtime {
diff --git a/src/runtime/disco/utils.h b/src/runtime/disco/utils.h
index ace3873a38..0c177e36e9 100644
--- a/src/runtime/disco/utils.h
+++ b/src/runtime/disco/utils.h
@@ -20,12 +20,10 @@
#define TVM_RUNTIME_DISCO_UTILS_H_
#include <dlpack/dlpack.h>
-#include <tvm/runtime/disco/session.h>
+#include <tvm/runtime/disco/disco_worker.h>
#include <string>
-#include "./worker.h"
-
namespace tvm {
namespace runtime {
@@ -36,34 +34,6 @@ inline Device UseDefaultDeviceIfNone(Device device) {
return device;
}
-/*!
- * \brief Possible kinds of reduction operations.
- */
-enum class ReduceKind : int32_t {
- kSum = 0,
- kProd = 1,
- kMin = 2,
- kMax = 3,
- kAvg = 4,
-};
-
-/*! \brief Converts `ReduceKind` to string */
-inline std::string ReduceKind2String(ReduceKind kind) {
- switch (kind) {
- case ReduceKind::kSum:
- return "kSum";
- case ReduceKind::kProd:
- return "kProd";
- case ReduceKind::kMin:
- return "kMin";
- case ReduceKind::kMax:
- return "kMax";
- case ReduceKind::kAvg:
- return "kAvg";
- }
- LOG(FATAL) << "ValueError: Unknown ReduceKind: " << static_cast<int>(kind);
-}
-
/*!
* \brief Converts a 1-d shape tuple to an integer.
* \note At the time of scaffolding Disco, RelaxVM has not provided mature
support for standalone
diff --git a/src/runtime/relax_vm/ndarray_cache_support.cc
b/src/runtime/relax_vm/ndarray_cache_support.cc
index 25f1fd282e..613c70bb44 100644
--- a/src/runtime/relax_vm/ndarray_cache_support.cc
+++ b/src/runtime/relax_vm/ndarray_cache_support.cc
@@ -38,11 +38,10 @@
#ifndef __STDC_FORMAT_MACROS
#define __STDC_FORMAT_MACROS
#endif
-#include "./ndarray_cache_support.h"
-
#include <picojson.h>
#include <tvm/runtime/ndarray.h>
#include <tvm/runtime/registry.h>
+#include <tvm/runtime/relax_vm/ndarray_cache_support.h>
#include <string>
#include <vector>
@@ -55,13 +54,13 @@ namespace runtime {
namespace relax_vm {
template <typename ExpectedType>
-ExpectedType AsType(const picojson::value& json) {
+inline ExpectedType AsType(const picojson::value& json) {
ICHECK(json.is<ExpectedType>());
return json.get<ExpectedType>();
}
template <typename ValueType>
-ValueType GetValue(const picojson::object& json, const std::string& key) {
+inline ValueType GetValue(const picojson::object& json, const std::string&
key) {
return AsType<ValueType>(json.at(key));
}
@@ -111,59 +110,63 @@ NDArrayCacheMetadata JSONAsNDArrayCacheMetadata(const
picojson::object& json) {
NDArrayCacheMetadata NDArrayCacheMetadata::LoadFromStr(const std::string&
json_str,
const std::string&
path) {
picojson::value json_info;
- picojson::parse(json_info, json_str);
+ {
+ std::string err = picojson::parse(json_info, json_str);
+ if (!err.empty()) {
+ LOG(FATAL) << "Failed to parse JSON: err. The JSON string is:" <<
json_str;
+ }
+ CHECK(json_info.is<picojson::object>())
+ << "ValueError: The given string is not a JSON object: " << json_str;
+ }
NDArrayCacheMetadata result =
JSONAsNDArrayCacheMetadata(AsType<picojson::object>(json_info));
result.path = path;
return result;
}
-ShardInfo::TensorInfo LoadTensorInfoFromJSON(const picojson::array&
json_tensor_info) {
- CHECK_EQ(json_tensor_info.size(), 2) << "ValueError: Invalid tensor info
JSON";
- picojson::array shape_json = AsType<picojson::array>(json_tensor_info[0]);
- int ndim = shape_json.size();
- std::vector<int64_t> shape;
- shape.reserve(ndim);
- for (int i = 0; i < ndim; ++i) {
- shape.push_back(AsType<int64_t>(shape_json[i]));
+NDArrayCacheMetadata NDArrayCacheMetadata::Load(const std::string& path) {
+ picojson::value json_info;
+ {
+ std::string json_str;
+ LoadBinaryFromFile(path + "/ndarray-cache.json", &json_str);
+ std::string err = picojson::parse(json_info, json_str);
+ if (!err.empty()) {
+ LOG(FATAL) << "Failed to parse JSON: err. The JSON string is:" <<
json_str;
+ }
+ CHECK(json_info.is<picojson::object>())
+ << "ValueError: The given string is not a JSON object: " << json_str;
}
- std::string dtype = AsType<std::string>(json_tensor_info[1]);
- return ShardInfo::TensorInfo{ShapeTuple(std::move(shape)),
DataType(String2DLDataType(dtype))};
+ NDArrayCacheMetadata result =
JSONAsNDArrayCacheMetadata(AsType<picojson::object>(json_info));
+ result.path = path;
+ return result;
}
-ShardInfo::ShardFunc LoadShardFuncFromJSON(const picojson::array&
json_shard_func) {
- int n = json_shard_func.size();
- ShardInfo::ShardFunc shard_info;
- shard_info.name = AsType<std::string>(json_shard_func[0]);
- shard_info.output_info =
LoadTensorInfoFromJSON(AsType<picojson::array>(json_shard_func[1]));
- shard_info.params.reserve(n - 2);
- for (int i = 2; i < n; ++i) {
- shard_info.params.push_back(AsType<int64_t>(json_shard_func[i]));
+void CopyNDArrayFromBytes(NDArray param, const void* data, size_t nbytes,
+ Optional<NDArray>* staging_buffer) {
+ Device device = param->device;
+ if (device.device_type != kDLOpenCL || staging_buffer == nullptr) {
+ param.CopyFromBytes(data, nbytes);
+ return;
}
- return shard_info;
-}
-
-std::unordered_map<std::string, ShardInfo> LoadShardInfoFromStr(const
std::string& json_str) {
- picojson::value json_info;
- picojson::parse(json_info, json_str);
- picojson::object json_obj = AsType<picojson::object>(json_info);
- std::unordered_map<std::string, ShardInfo> result;
- for (auto kv : json_obj) {
- std::string name = kv.first;
- picojson::array json_shard_funcs = AsType<picojson::array>(kv.second);
- ShardInfo info;
- std::vector<ShardInfo::ShardFunc>& shard_funcs = info.funcs;
- shard_funcs.reserve(json_shard_funcs.size());
- for (const picojson::value& json_shard_func : json_shard_funcs) {
-
shard_funcs.push_back(LoadShardFuncFromJSON(AsType<picojson::array>(json_shard_func)));
+ // Special handle for OpenCL runtime.
+ // It creates a host side memory mirror, for every cl_mem that tries to copy
data from host
+ // which can cause memory issue. Her we use a large staging buffer to
postpone deallocation
+ if (staging_buffer->defined()) {
+ size_t curr_size =
runtime::GetDataSize(*(staging_buffer->value().operator->()));
+ if (curr_size < nbytes) {
+ *staging_buffer = NullOpt;
}
- result[name] = info;
}
- return result;
+ if (!staging_buffer->defined()) {
+ *staging_buffer = NDArray::Empty(param.Shape(), param->dtype,
param->device);
+ }
+ NDArray staging_view = staging_buffer->value().CreateView(param.Shape(),
param->dtype);
+ staging_view.CopyFromBytes(data, nbytes);
+ param.CopyFrom(staging_view);
+ TVMSynchronize(device.device_type, device.device_id, nullptr);
}
NDArray NDArrayCacheMetadata::FileRecord::ParamRecord::Load(
- Device device, const std::string* raw_data,
- std::function<void(NDArray, const void*, int64_t)> f_load) const {
+ Device device, const std::string* raw_data, Optional<NDArray>*
staging_buffer) const {
NDArray arr = NDArray::Empty(shape, dtype, device);
if (dtype == DataType::Float(32) && format == "f32-to-bf16") {
// decode bf16 to f32
@@ -173,13 +176,30 @@ NDArray
NDArrayCacheMetadata::FileRecord::ParamRecord::Load(
for (size_t i = 0; i < buffer.size(); ++i) {
decoded[i] = static_cast<uint32_t>(buffer[i]) << 16;
}
- f_load(arr, decoded.data(), decoded.size() * sizeof(uint32_t));
+ CopyNDArrayFromBytes(arr, decoded.data(), decoded.size() *
sizeof(uint32_t), staging_buffer);
} else {
- f_load(arr, raw_data->data() + byte_offset, nbytes);
+ CopyNDArrayFromBytes(arr, raw_data->data() + byte_offset, nbytes,
staging_buffer);
}
return arr;
}
+Array<NDArray> NDArrayCacheMetadata::FileRecord::Load(Device device,
+ const std::string&
path_prefix, //
+ std::string*
raw_data_buffer, //
+ Optional<NDArray>*
staging_buffer) const {
+ LoadBinaryFromFile(path_prefix + "/" + this->data_path, raw_data_buffer);
+ CHECK_EQ(this->format, "raw-shard") << "ValueError: Only `raw-shard` format
is supported";
+ CHECK_EQ(this->nbytes, raw_data_buffer->length())
+ << "ValueError: Encountered an corrupted parameter shard. It means it is
not downloaded "
+ "completely or downloading is interrupted. Please try to download
again.";
+ Array<NDArray> result;
+ result.reserve(this->records.size());
+ for (const ParamRecord& nd_rec : this->records) {
+ result.push_back(nd_rec.Load(device, raw_data_buffer, staging_buffer));
+ }
+ return result;
+}
+
/*!
* A NDArray cache to store pre-loaded arrays in the system.
*/
@@ -217,53 +237,26 @@ class NDArrayCache {
/*!
* \brief Load parameters from path and append them.
- *
* \param cache_path The cache to path.
* \param device_type The type of device to be loaded.
* \param device_id The device id.
*/
static void Load(const std::string& cache_path, int device_type, int
device_id) {
DLDevice device{static_cast<DLDeviceType>(device_type), device_id};
- std::string json_str;
- LoadBinaryFromFile(cache_path + "/ndarray-cache.json", &json_str);
- NDArrayCacheMetadata metadata =
NDArrayCacheMetadata::LoadFromStr(json_str, cache_path);
+ NDArrayCacheMetadata metadata = NDArrayCacheMetadata::Load(cache_path);
Optional<NDArray> staging_buffer;
- auto fcopy_param_from_bytes = [&](NDArray param, const void* data, size_t
nbytes) {
- if (device_type != kDLOpenCL) {
- param.CopyFromBytes(data, nbytes);
- return;
- }
- // special handle OpenCL
- // OpenCL runtime can create a host side memory mirror
- // for every cl_mem that tries to copy data from host
- // which can cause memory issue.
- // We use a single staging buffer here
- // that get de-allocated later
- if (staging_buffer.defined()) {
- size_t curr_size =
runtime::GetDataSize(*(staging_buffer.value().operator->()));
- if (curr_size < nbytes) {
- staging_buffer = NullOpt;
- }
- }
- if (!staging_buffer.defined()) {
- staging_buffer = NDArray::Empty(param.Shape(), param->dtype,
param->device);
- }
- NDArray staging_view = staging_buffer.value().CreateView(param.Shape(),
param->dtype);
- staging_view.CopyFromBytes(data, nbytes);
- param.CopyFrom(staging_view);
- TVMSynchronize(device_type, device_id, nullptr);
- };
-
- Map<String, NDArray> result;
std::string raw_data;
- for (const auto& shard_rec : metadata.records) {
- LoadBinaryFromFile(cache_path + "/" + shard_rec.data_path, &raw_data);
- CHECK_EQ(shard_rec.format, "raw-shard") << "ValueError: Only `raw-shard`
format is supported";
- CHECK_EQ(shard_rec.nbytes, raw_data.length())
- << "ValueError: Parameters are not loaded properly. Please check
your parameter shards "
- "and git lfs installation";
- for (const auto& nd_rec : shard_rec.records) {
- Update(nd_rec.name, nd_rec.Load(device, &raw_data,
fcopy_param_from_bytes), true);
+ Array<NDArray> params;
+ for (const NDArrayCacheMetadata::FileRecord& shard_rec : metadata.records)
{
+ try {
+ params = shard_rec.Load(device, cache_path, &raw_data,
&staging_buffer);
+ } catch (const dmlc::Error& e) {
+ LOG(FATAL) << "ValueError: Error when loading parameters from " <<
shard_rec.data_path
+ << ": " << e.what();
+ }
+ int num_params = params.size();
+ for (int i = 0; i < num_params; ++i) {
+ Update(shard_rec.records[i].name, params[i], true);
}
}
}
