You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by ju...@apache.org on 2023/09/03 06:10:49 UTC
[tvm] branch unity updated: [Disco] Introduce ShardLoader (#15655)
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 f1246d0078 [Disco] Introduce ShardLoader (#15655)
f1246d0078 is described below
commit f1246d0078ba46f8e9195acd4a35f84b7dea0c80
Author: Junru Shao <ju...@apache.org>
AuthorDate: Sat Sep 2 23:10:43 2023 -0700
[Disco] Introduce ShardLoader (#15655)
This PR introduces `ShardLoader`, an object that allows convenient
sharding of each parameter, assuming there is a single shard dimension
and the sharding strategy is even. The sharding can be performed
efficiently on device (e.g. CUDA) and scattered to each worker with
NCCL support.
The shard loading process could be further compiled to Relax IRModule
to be executed by each worker with the function signature below:
```python
@R.function
def main(loader: R.Object) -> R.Tuple(
R.Tensor((64, 64), "float32"),
R.Tensor((16, 128), "float32"),
):
R.func_attr({"global_symbol": "main"})
with R.dataflow():
lv0 = R.call_pure_packed(
"runtime.disco.ShardLoaderLoad",
loader,
R.shape([0]),
sinfo_args=R.Tensor((64, 64), "float32"),
)
lv1 = R.call_pure_packed(
"runtime.disco.ShardLoaderLoad",
loader,
R.shape([1]),
sinfo_args=R.Tensor((16, 128), "float32"),
)
lv2 = R.tuple(lv0, lv1)
R.output(lv2)
return lv2
```
---
src/runtime/disco/loader.cc | 191 ++++++++++++++++++++++++++++++++++++++
tests/python/disco/test_loader.py | 178 +++++++++++++++++++++++++++++++++++
2 files changed, 369 insertions(+)
diff --git a/src/runtime/disco/loader.cc b/src/runtime/disco/loader.cc
new file mode 100644
index 0000000000..6f9054296a
--- /dev/null
+++ b/src/runtime/disco/loader.cc
@@ -0,0 +1,191 @@
+/*
+ * 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.
+ */
+#include <tvm/runtime/packed_func.h>
+#include <tvm/runtime/registry.h>
+
+#include <functional>
+#include <numeric>
+#include <string>
+#include <vector>
+
+#include "../file_utils.h"
+#include "../relax_vm/ndarray_cache_support.h"
+#include "./builtin.h"
+#include "./utils.h"
+
+namespace tvm {
+namespace runtime {
+
+using relax_vm::NDArrayCacheMetadata;
+using FileRecord = NDArrayCacheMetadata::FileRecord;
+using ParamRecord = NDArrayCacheMetadata::FileRecord::ParamRecord;
+using relax_vm::LoadShardInfoFromStr;
+
+/*! \brief An object that helps to load parameters in shards. */
+class ShardLoaderObj : public Object {
+ public:
+ /*! \brief Create a shard loader. */
+ static ObjectRef Create(const std::string& path_to_metadata, const std::string& metadata,
+ const std::string& shard_info,
+ TypedPackedFunc<void(DLTensor*, int, DLTensor*)> f_shard);
+ /*! \brief Load the i-th parameter */
+ NDArray Load(int weight_index) const;
+ /*! \brief Slice the given tensor at a specific dimension */
+ NDArray Shard(NDArray source, int dim, int num_slices) const;
+
+ static constexpr const char* _type_key = "runtime.disco.ShardLoader";
+ TVM_DECLARE_FINAL_OBJECT_INFO(ShardLoaderObj, Object);
+
+ public:
+ /*! \brief Information of how each weight is stored and sharded */
+ struct ShardInfo {
+ const FileRecord* file;
+ const ParamRecord* param;
+ int shard_dim;
+ };
+ /*! \brief The metadata loaded from `ndarray-cache.json` */
+ NDArrayCacheMetadata metadata_;
+ /*! \brief Sharding information for each weight */
+ std::vector<ShardInfo> shard_info_;
+ /*! \brief A method to slice a 3-D tensor */
+ TypedPackedFunc<void(DLTensor*, int, DLTensor*)> f_shard_;
+ /*! \brief The current file opened to load weights in it */
+ mutable const FileRecord* current_file_;
+ /*! \brief The context of the current file to be loaded from */
+ mutable std::string current_file_stream_;
+};
+
+TVM_REGISTER_OBJECT_TYPE(ShardLoaderObj);
+
+/*!
+ * \brief Get the shape of a result tensor if it is scattered along a given axis.
+ * \param shape The shape of the input tensor.
+ * \param dim The axis along which the tensor is scattered.
+ * \param num_shards The number of shards.
+ * \return The shape of the result tensor.
+ */
+inline std::vector<ShapeTuple::index_type> ShardShape(const ShapeTuple& shape, int dim,
+ int num_shards) {
+ CHECK(0 <= dim && dim < static_cast<int>(shape.size()))
+ << "ValueError: Cannot scatter at dim " << dim << ", because "
+ << "shape is " << shape << ".";
+ CHECK_EQ(shape[dim] % num_shards, 0)
+ << "ValueError: The shape " << shape << " cannot be scattered at dim " << dim << " into "
+ << num_shards << " shards.";
+ std::vector<ShapeTupleObj::index_type> result{shape.begin(), shape.end()};
+ result[dim] /= num_shards;
+ return result;
+}
+
+ObjectRef ShardLoaderObj::Create(const std::string& path_to_metadata, const std::string& metadata,
+ const std::string& shard_info,
+ TypedPackedFunc<void(DLTensor*, int, DLTensor*)> f_shard) {
+ ObjectPtr<ShardLoaderObj> n = make_object<ShardLoaderObj>();
+ n->f_shard_ = f_shard;
+ n->metadata_ = NDArrayCacheMetadata::LoadFromStr(metadata, path_to_metadata);
+ n->current_file_ = nullptr;
+ n->shard_info_.clear();
+ std::unordered_map<std::string, int> 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;
+ int shard_id = shards.count(name) ? shards[name] : -1;
+ n->shard_info_.push_back(ShardInfo{&file_record, ¶m_record, shard_id});
+ }
+ }
+ return ObjectRef(std::move(n));
+}
+
+std::string GetSiblingPath(const std::string& path, const std::string& filename) {
+ size_t found = path.find_last_of("/\\");
+ if (found != std::string::npos) {
+ return path.substr(0, found + 1) + filename;
+ }
+ LOG(FATAL) << "ValueError: Cannot find the parent directory: " << path;
+}
+
+NDArray ShardLoaderObj::Load(int weight_index) const {
+ DiscoWorker* worker = DiscoWorker::ThreadLocal();
+ int shard_idx = worker->worker_id;
+ Device device = worker->default_device;
+ const auto& shard_info = shard_info_[weight_index];
+ const ParamRecord* param = shard_info.param;
+ const FileRecord* file = shard_info.file;
+ int shard_dim = shard_info.shard_dim;
+ int num_shards = worker->num_workers;
+ Optional<NDArray> send = NullOpt;
+ if (shard_idx == 0) {
+ if (file != current_file_) {
+ current_file_ = file;
+ std::string file_name = GetSiblingPath(this->metadata_.path, file->data_path);
+ LoadBinaryFromFile(file_name, &this->current_file_stream_);
+ }
+ auto f_load = [](NDArray param, const void* data, size_t nbytes) {
+ param.CopyFromBytes(data, nbytes);
+ };
+ send = this->Shard(param->Load(device, &this->current_file_stream_, f_load), shard_dim,
+ num_shards);
+ }
+ NDArray recv =
+ NDArray::Empty(ShardShape(param->shape, shard_dim, num_shards), param->dtype, device);
+ ScatterFromWorker0(send, recv);
+ return recv;
+}
+
+NDArray ShardLoaderObj::Shard(NDArray source, int dim, int num_slices) const {
+ ICHECK(dim >= 0 && dim < source->ndim);
+ // Assemble a flattened 3d src tensor
+ int64_t src_flat[3] = {1, 1, 1};
+ {
+ const int64_t* s = source.Shape().data();
+ int ndim = source->ndim;
+ src_flat[0] = std::accumulate(&s[0], &s[dim], 1, std::multiplies<int64_t>());
+ src_flat[1] = s[dim];
+ src_flat[2] = std::accumulate(&s[dim + 1], &s[ndim], 1, std::multiplies<int64_t>());
+ }
+ DLTensor src_tensor = *source.operator->();
+ src_tensor.ndim = 3;
+ src_tensor.shape = src_flat;
+ // Assmeble a flattened 4d dst tensor
+ int64_t dst_flat[4] = {num_slices, src_flat[0], src_flat[1] / num_slices, src_flat[2]};
+ NDArray destination{nullptr};
+ {
+ std::vector<ShapeTuple::index_type> dst_shape = ShardShape(source.Shape(), dim, num_slices);
+ dst_shape.insert(dst_shape.begin(), static_cast<ShapeTuple::index_type>(num_slices));
+ destination = NDArray::Empty(dst_shape, source->dtype, source->device);
+ }
+ DLTensor dst_tensor = *destination.operator->();
+ dst_tensor.ndim = 4;
+ dst_tensor.shape = dst_flat;
+ // Copy slices using the API
+ this->f_shard_(&src_tensor, num_slices, &dst_tensor);
+ return destination;
+}
+
+TVM_REGISTER_GLOBAL("runtime.disco.ShardLoader").set_body_typed(ShardLoaderObj::Create);
+TVM_REGISTER_GLOBAL("runtime.disco.ShardLoaderLoad")
+ .set_body_typed([](ObjectRef loader_obj, ShapeTuple weight_index) {
+ const auto* loader = loader_obj.as<ShardLoaderObj>();
+ CHECK(loader != nullptr) << "TypeError: Expected ShardLoaderObj, but gets: "
+ << loader_obj->GetTypeKey();
+ return loader->Load(IntegerFromShapeTuple(weight_index));
+ });
+
+} // namespace runtime
+} // namespace tvm
diff --git a/tests/python/disco/test_loader.py b/tests/python/disco/test_loader.py
new file mode 100644
index 0000000000..baa6745b00
--- /dev/null
+++ b/tests/python/disco/test_loader.py
@@ -0,0 +1,178 @@
+# 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.
+"""Test sharded loader"""
+# pylint: disable=missing-docstring
+import json
+import tempfile
+
+import numpy as np
+
+from tvm import dlight as dl
+from tvm import relax as rx
+from tvm._ffi import register_func
+from tvm.contrib import tvmjs
+from tvm.runtime import ShapeTuple
+from tvm.runtime import disco as di
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.target import Target
+
+
+@register_func("tests.disco.shard_with_numpy", override=True)
+def _shard_with_numpy(src, num_shards, tgt):
+ s_0, s_1, s_2 = src.shape
+ tgt.copyfrom(src.numpy().reshape(s_0, num_shards, s_1 // num_shards, s_2).transpose(1, 0, 2, 3))
+
+
+def _create_loader(sess, path, param_dict, shard_info):
+ path_ndarray_cache = path + "/ndarray-cache.json"
+ tvmjs.dump_ndarray_cache(param_dict, path, encode_format="raw")
+ with open(path_ndarray_cache, "r", encoding="utf-8") as i_f:
+ ndarray_cache = i_f.read()
+ shard_with_numpy = sess.get_global_func("tests.disco.shard_with_numpy")
+ loader_create = sess.get_global_func("runtime.disco.ShardLoader")
+ loader = loader_create(path_ndarray_cache, ndarray_cache, shard_info, shard_with_numpy)
+ return loader
+
+
+def test_load_shard():
+ devices = [1, 2]
+ param_dict = {
+ "x_0": np.random.uniform(size=[64, 128]).astype("float16"),
+ "x_1": np.random.uniform(size=[32, 128]).astype("float32"),
+ }
+ shard_info = json.dumps(
+ {
+ "x_0": 1,
+ "x_1": 0,
+ }
+ )
+ with tempfile.TemporaryDirectory() as path:
+ sess = di.ThreadedSession(num_workers=len(devices))
+ sess.init_ccl("nccl", *devices)
+ loader = _create_loader(sess, path, param_dict, shard_info)
+ loader_load = sess.get_global_func("runtime.disco.ShardLoaderLoad")
+ d_0 = loader_load(loader, ShapeTuple([0]))
+ d_1 = loader_load(loader, ShapeTuple([1]))
+ np.testing.assert_equal(
+ param_dict["x_0"][:, 0:64],
+ d_0.debug_get_from_remote(0).numpy(),
+ )
+ np.testing.assert_equal(
+ param_dict["x_0"][:, 64:128],
+ d_0.debug_get_from_remote(1).numpy(),
+ )
+ np.testing.assert_equal(
+ param_dict["x_1"][0:16, :],
+ d_1.debug_get_from_remote(0).numpy(),
+ )
+ np.testing.assert_equal(
+ param_dict["x_1"][16:32, :],
+ d_1.debug_get_from_remote(1).numpy(),
+ )
+
+
+def test_load_shard_in_relax():
+ devices = [1, 2]
+ param_dict = {
+ "x_0": np.random.uniform(size=[64, 128]).astype("float16"),
+ "x_1": np.random.uniform(size=[32, 128]).astype("float32"),
+ }
+ shard_info = json.dumps(
+ {
+ "x_0": 1,
+ "x_1": 0,
+ }
+ )
+
+ # pylint: disable=invalid-name
+ @I.ir_module
+ class Module: # pylint: disable=too-few-public-methods
+ @R.function
+ def main(
+ loader: R.Object,
+ ) -> R.Tuple(R.Tensor((64, 64), "float32"), R.Tensor((16, 128), "float32"),):
+ R.func_attr({"global_symbol": "main"})
+ with R.dataflow():
+ lv0: R.Tensor((64, 64), "float32") = R.call_pure_packed(
+ "runtime.disco.ShardLoaderLoad",
+ loader,
+ R.shape([0]),
+ sinfo_args=R.Tensor((64, 64), "float32"),
+ )
+ lv1: R.Tensor((16, 128), "float32") = R.call_pure_packed(
+ "runtime.disco.ShardLoaderLoad",
+ loader,
+ R.shape([1]),
+ sinfo_args=R.Tensor((16, 128), "float32"),
+ )
+ lv2 = R.tuple(lv0, lv1)
+ R.output(lv2)
+ return lv2
+
+ # pylint: enable=invalid-name
+ def relax_build(mod, target):
+ with target:
+ mod = rx.get_pipeline("zero")(mod) # pylint: disable=no-value-for-parameter
+ mod = dl.ApplyDefaultSchedule( # pylint: disable=not-callable
+ dl.gpu.Matmul(),
+ dl.gpu.GEMV(),
+ dl.gpu.Reduction(),
+ dl.gpu.GeneralReduction(),
+ dl.gpu.Fallback(),
+ )(mod)
+ return rx.build(mod, target="cuda")
+
+ target = Target(
+ {
+ "kind": "cuda",
+ "max_shared_memory_per_block": 49152,
+ "max_threads_per_block": 1024,
+ "thread_warp_size": 32,
+ "registers_per_block": 65536,
+ "arch": "sm_80",
+ }
+ )
+ with tempfile.TemporaryDirectory() as tmpdir:
+ dso_path = tmpdir + "/test.so"
+ relax_build(Module, target).export_library(dso_path)
+ sess = di.ThreadedSession(num_workers=len(devices))
+ sess.init_ccl("nccl", *devices)
+ mod = sess.load_vm_module(dso_path)
+ loader = _create_loader(sess, tmpdir, param_dict, shard_info)
+ result = mod["main"](loader)
+ np.testing.assert_equal(
+ param_dict["x_0"][:, 0:64],
+ result.debug_get_from_remote(0)[0].numpy(),
+ )
+ np.testing.assert_equal(
+ param_dict["x_0"][:, 64:128],
+ result.debug_get_from_remote(1)[0].numpy(),
+ )
+ np.testing.assert_equal(
+ param_dict["x_1"][0:16, :],
+ result.debug_get_from_remote(0)[1].numpy(),
+ )
+ np.testing.assert_equal(
+ param_dict["x_1"][16:32, :],
+ result.debug_get_from_remote(1)[1].numpy(),
+ )
+
+
+if __name__ == "__main__":
+ test_load_shard()
+ test_load_shard_in_relax()