chore: import upstream snapshot with attribution

tests/python/disco/test_loader.py

@@ -0,0 +1,496 @@
+# 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.
+# ruff: noqa: F401, F841
+"""Test sharded loader"""
+
+# pylint: disable=missing-docstring
+import json
+import tempfile
+
+import numpy as np
+import pytest
+from tvm_ffi import Shape, register_global_func
+
+import tvm
+import tvm.testing
+from tvm import relax as rx
+from tvm.contrib import tvmjs
+from tvm.runtime import disco as di
+from tvm.s_tir import dlight as dl
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.target import Target
+from tvm.testing import env
+
+# `runtime.disco.compiled_ccl` is registered together with the CCL runtime
+# functions, so its absence means the disco CCL runtime is not in this build.
+_compiled_ccl = tvm.get_global_func("runtime.disco.compiled_ccl", allow_missing=True)
+if _compiled_ccl is None or _compiled_ccl() != "nccl":
+    pytest.skip("Disco NCCL support is not available", allow_module_level=True)
+
+# All tests in this file shard across two GPUs.
+pytestmark = [
+    pytest.mark.skipif(not env.has_multi_gpu(), reason="need multiple gpus"),
+]
+
+
+def _run_with_nccl_session(devices, func):
+    def run_and_check():
+        sess = di.ThreadedSession(num_workers=len(devices))
+        try:
+            sess.init_ccl("nccl", *devices)
+            return func(sess)
+        finally:
+            sess.shutdown()
+
+    return tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@register_global_func("tests.disco.shard_dim_0", override=True)
+def _shard_dim_0(src, num_shards, tgt):
+    s_0, s_1 = src.shape
+    tgt.copyfrom(src.numpy().reshape(num_shards, s_0 // num_shards, s_1))
+
+
+@register_global_func("tests.disco.shard_dim_1", override=True)
+def _shard_dim_1(src, num_shards, tgt):
+    s_0, s_1 = src.shape
+    tgt.copyfrom(src.numpy().reshape(s_0, num_shards, s_1 // num_shards).transpose(1, 0, 2))
+
+
+@register_global_func("tests.disco.shard_qkv_0", override=True)
+def _shard_qkv_0(src, num_shards, q_heads, kv_heads, tgt):
+    total_dim, hidden_size = src.shape
+    head_dim = total_dim // (q_heads + kv_heads + kv_heads)
+    q_dim = q_heads * head_dim
+    kv_dim = kv_heads * head_dim
+    w_q = src.numpy()[:q_dim, :].reshape(
+        num_shards,
+        q_heads // num_shards,
+        head_dim,
+        hidden_size,
+    )
+    w_k = src.numpy()[q_dim : q_dim + kv_dim, :].reshape(
+        num_shards,
+        kv_heads // num_shards,
+        head_dim,
+        hidden_size,
+    )
+    w_v = src.numpy()[q_dim + kv_dim :, :].reshape(
+        num_shards,
+        kv_heads // num_shards,
+        head_dim,
+        hidden_size,
+    )
+    w_qkv = np.concatenate([w_q, w_k, w_v], axis=1)
+    tgt.copyfrom(w_qkv)
+
+
+@register_global_func("tests.disco.shard_qkv_1", override=True)
+def _shard_qkv_1(src, tgt):
+    s, _, _, h = src.shape  # pylint: disable=invalid-name
+    tgt.copyfrom(src.numpy().reshape(s, -1, h))
+
+
+def _create_loader(sess, path, param_dict, shard_info):
+    path_tensor_cache = path + "/tensor-cache.json"
+    tvmjs.dump_tensor_cache(param_dict, path, encode_format="raw")
+    with open(path_tensor_cache, encoding="utf-8") as i_f:
+        tensor_cache = i_f.read()
+    loader_create = sess.get_global_func("runtime.disco.ShardLoader")
+    loader = loader_create(path_tensor_cache, tensor_cache, json.dumps(shard_info), None)
+    return loader
+
+
+def _simulate_presharded_weights(base_path, param_dict, num_shards, shard_info):
+    """Create fake weights to simulate those produced MLC-LLM's pre-sharding"""
+
+    sharded_params = {}
+
+    for key, ndarray in param_dict.items():
+        assert key in shard_info, f"ShardInfo lacks shard info about param: {key}"
+        shard_dim = shard_info[key]
+        sharded_params[key] = [
+            tvm.runtime.tensor(np_shard)
+            for np_shard in np.split(ndarray, num_shards, axis=shard_dim)
+        ]
+
+    # Re-order so that the parameter order is sorted first by shard,
+    # then by parameter.  This matches the ordering used by MLC-LLM,
+    # and avoids having *.bin files that must be accessed by more than
+    # one worker.
+    sharded_params = {
+        f"{key}_shard-{i + 1}-of-{num_shards}": shards[i]
+        for i in range(num_shards)
+        for key, shards in sharded_params.items()
+    }
+
+    tvmjs.dump_tensor_cache(
+        sharded_params,
+        base_path,
+        encode_format="raw",
+    )
+
+
+def test_load_shard():
+    devices = [0, 1]
+    num_shards = len(devices)
+    param_dict = {
+        "x_0": np.random.uniform(size=[64, 128]).astype("float16"),
+        "x_1": np.random.uniform(size=[32, 128]).astype("float32"),
+    }
+    shard_info = {
+        "x_0": [
+            [
+                "tests.disco.shard_dim_1",
+                [(num_shards, 64, 64), "float16"],
+                num_shards,
+            ],
+        ],
+        "x_1": [
+            [
+                "tests.disco.shard_dim_0",
+                [(num_shards, 16, 128), "float32"],
+                num_shards,
+            ]
+        ],
+    }
+    with tempfile.TemporaryDirectory() as path:
+
+        def run_test(sess):
+            loader = _create_loader(sess, path, param_dict, shard_info)
+            loader_load = sess.get_global_func("runtime.disco.ShardLoaderLoad")
+            d_0 = loader_load(loader, Shape([0]))
+            d_1 = loader_load(loader, Shape([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(),
+            )
+
+        _run_with_nccl_session(devices, run_test)
+
+
+def _create_presharded_loader(sess, path):
+    path_tensor_cache = path + "/tensor-cache.json"
+    with open(path_tensor_cache, encoding="utf-8") as i_f:
+        tensor_cache = i_f.read()
+    loader_create = sess.get_global_func("runtime.disco.ShardLoader")
+    loader = loader_create(path_tensor_cache, tensor_cache, json.dumps({}), None)
+    return loader
+
+
+def test_load_presharded():
+    devices = [0, 1]
+    param_dict = {
+        "x_0": np.random.uniform(size=[64, 128]).astype("float16"),
+        "x_1": np.random.uniform(size=[32, 128]).astype("float32"),
+    }
+    shard_info = {
+        "x_0": 1,
+        "x_1": 0,
+    }
+
+    with tempfile.TemporaryDirectory() as path:
+        _simulate_presharded_weights(path, param_dict, len(devices), shard_info)
+
+        def run_test(sess):
+            loader = _create_presharded_loader(sess, path)
+            loader_load = sess.get_global_func("runtime.disco.ShardLoaderLoadPresharded")
+
+            d_0 = loader_load(loader, Shape([0]))
+            d_1 = loader_load(loader, Shape([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(),
+            )
+
+        _run_with_nccl_session(devices, run_test)
+
+
+def test_load_shard_in_relax():
+    devices = [0, 1]
+    num_shards = len(devices)
+    param_dict = {
+        "x_0": np.random.uniform(size=[64, 128]).astype("float16"),
+        "x_1": np.random.uniform(size=[32, 128]).astype("float32"),
+    }
+    shard_info = {
+        "x_0": [
+            [
+                "tests.disco.shard_dim_1",
+                [(num_shards, 64, 64), "float16"],
+                num_shards,
+            ],
+        ],
+        "x_1": [
+            [
+                "tests.disco.shard_dim_0",
+                [(num_shards, 16, 128), "float32"],
+                num_shards,
+            ]
+        ],
+    }
+
+    # pylint: disable=invalid-name
+    @I.ir_module
+    class Module:  # pylint: disable=too-few-public-methods
+        @R.function
+        def main(
+            loader: R.Any,
+        ) -> 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]),
+                    ty_args=R.Tensor((64, 64), "float32"),
+                )
+                lv1: R.Tensor((16, 128), "float32") = R.call_pure_packed(
+                    "runtime.disco.ShardLoaderLoad",
+                    loader,
+                    R.shape([1]),
+                    ty_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
+            return tvm.compile(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)
+
+        def run_test(sess):
+            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(),
+            )
+
+        _run_with_nccl_session(devices, run_test)
+
+
+def test_load_shard_all():
+    devices = [0, 1]
+    num_shards = len(devices)
+    param_dict = {
+        "param_0": np.random.uniform(size=[64, 128]).astype("float16"),
+        "param_1": np.random.uniform(size=[32, 128]).astype("float32"),
+    }
+    shard_info = {
+        "param_0": [
+            [
+                "tests.disco.shard_dim_1",
+                [(num_shards, 64, 64), "float16"],
+                num_shards,
+            ],
+        ],
+        "param_1": [
+            [
+                "tests.disco.shard_dim_0",
+                [(2, 16, 128), "float32"],
+                num_shards,
+            ]
+        ],
+    }
+    with tempfile.TemporaryDirectory() as path:
+
+        def run_test(sess):
+            loader = _create_loader(sess, path, param_dict, shard_info)
+            loader_load = sess.get_global_func("runtime.disco.ShardLoaderLoadAll")
+            params = loader_load(loader)
+            p_0 = params.debug_get_from_remote(0)
+            p_1 = params.debug_get_from_remote(1)
+            np.testing.assert_equal(param_dict["param_0"][:, 0:64], p_0[0].numpy())
+            np.testing.assert_equal(param_dict["param_0"][:, 64:128], p_1[0].numpy())
+            np.testing.assert_equal(param_dict["param_1"][0:16, :], p_0[1].numpy())
+            np.testing.assert_equal(param_dict["param_1"][16:32, :], p_1[1].numpy())
+
+        _run_with_nccl_session(devices, run_test)
+
+
+def test_load_all_presharded():
+    devices = [0, 1]
+    num_shards = len(devices)
+    param_dict = {
+        "param_0": np.random.uniform(size=[64, 128]).astype("float16"),
+        "param_1": np.random.uniform(size=[32, 128]).astype("float32"),
+    }
+    shard_info = {
+        "param_0": 0,
+        "param_1": 1,
+    }
+    with tempfile.TemporaryDirectory() as path:
+        _simulate_presharded_weights(path, param_dict, len(devices), shard_info)
+
+        def run_test(sess):
+            loader = _create_presharded_loader(sess, path)
+            loader_load = sess.get_global_func("runtime.disco.ShardLoaderLoadAllPresharded")
+            params = loader_load(loader)
+
+            p_0 = params.debug_get_from_remote(0)
+            p_1 = params.debug_get_from_remote(1)
+
+            np.testing.assert_equal(param_dict["param_0"][0:32, :], p_0[0].numpy())
+            np.testing.assert_equal(param_dict["param_0"][32:64, :], p_1[0].numpy())
+            np.testing.assert_equal(param_dict["param_1"][:, 0:64], p_0[1].numpy())
+            np.testing.assert_equal(param_dict["param_1"][:, 64:128], p_1[1].numpy())
+
+        _run_with_nccl_session(devices, run_test)
+
+
+def test_load_shard_broadcast():
+    devices = [0, 1]
+    param_dict = {
+        "param_0": np.random.uniform(size=[64, 128]).astype("float16"),
+        "param_1": np.random.uniform(size=[32, 128]).astype("float32"),
+    }
+    shard_info = {}
+    with tempfile.TemporaryDirectory() as path:
+
+        def run_test(sess):
+            loader = _create_loader(sess, path, param_dict, shard_info)
+            loader_load = sess.get_global_func("runtime.disco.ShardLoaderLoadAll")
+            params = loader_load(loader)
+            p_0 = params.debug_get_from_remote(0)
+            p_1 = params.debug_get_from_remote(1)
+            np.testing.assert_equal(param_dict["param_0"], p_0[0].numpy())
+            np.testing.assert_equal(param_dict["param_0"], p_1[0].numpy())
+            np.testing.assert_equal(param_dict["param_1"], p_0[1].numpy())
+            np.testing.assert_equal(param_dict["param_1"], p_1[1].numpy())
+
+        _run_with_nccl_session(devices, run_test)
+
+
+def test_load_qkv_proj_shard():  # pylint: disable=too-many-locals
+    devices = [0, 1]
+    num_shards = len(devices)
+    q_heads = 8
+    kv_heads = 10
+    head_dim = 10
+    hidden_size = 20
+    w_q = np.random.uniform(size=[q_heads * head_dim, hidden_size]).astype("float16")
+    w_k = np.random.uniform(size=[kv_heads * head_dim, hidden_size]).astype("float16")
+    w_v = np.random.uniform(size=[kv_heads * head_dim, hidden_size]).astype("float16")
+    w_qkv = np.concatenate([w_q, w_k, w_v], axis=0)
+    param_dict = {"w_qkv": w_qkv}
+    np_qkv = np.concatenate(
+        [
+            w_q.reshape((num_shards, q_heads // num_shards, head_dim, hidden_size)),
+            w_k.reshape((num_shards, kv_heads // num_shards, head_dim, hidden_size)),
+            w_v.reshape((num_shards, kv_heads // num_shards, head_dim, hidden_size)),
+        ],
+        axis=1,
+    ).reshape((num_shards, -1, hidden_size))
+
+    shard_info = {
+        "w_qkv": [
+            [
+                "tests.disco.shard_qkv_0",
+                [
+                    (num_shards, (q_heads + kv_heads * 2) // num_shards, head_dim, hidden_size),
+                    "float16",
+                ],
+                num_shards,
+                q_heads,
+                kv_heads,
+            ],
+            [
+                "tests.disco.shard_qkv_1",
+                [
+                    (num_shards, (q_heads + kv_heads * 2) // num_shards * head_dim, hidden_size),
+                    "float16",
+                ],
+            ],
+        ],
+    }
+
+    with tempfile.TemporaryDirectory() as path:
+
+        def run_test(sess):
+            loader = _create_loader(sess, path, param_dict, shard_info)
+            loader_load = sess.get_global_func("runtime.disco.ShardLoaderLoad")
+            d_0 = loader_load(loader, Shape([0]))
+            np.testing.assert_equal(
+                np_qkv[0],
+                d_0.debug_get_from_remote(0).numpy(),
+            )
+            np.testing.assert_equal(
+                np_qkv[1],
+                d_0.debug_get_from_remote(1).numpy(),
+            )
+
+        _run_with_nccl_session(devices, run_test)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()