chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,175 @@
|
||||
# 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 pathlib
|
||||
import tempfile
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import relax as R
|
||||
from tvm.script import tirx as T
|
||||
from tvm.testing import env
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(tvm.runtime.disco is None, reason="disco runtime is not available")
|
||||
@pytest.mark.skipif(not env.has_nccl(), reason="need nccl")
|
||||
@pytest.mark.skipif(not env.has_multi_gpu(), reason="need multiple gpus")
|
||||
def test_callback():
|
||||
"""Simulate lazy loading of parameters in a callback
|
||||
|
||||
The output of a lazy parameter loading, which would accept a
|
||||
callback to load the parameters.
|
||||
"""
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class Module:
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def slice_A(
|
||||
A: T.Buffer((4, 4), "int32"),
|
||||
rank: T.int64,
|
||||
A_sharded: T.Buffer((2, 4), "int32"),
|
||||
):
|
||||
for i, j in T.grid(2, 4):
|
||||
with T.sblock("slice_A"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
A_sharded[vi, vj] = A[rank * 2 + vi, vj]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def slice_B(
|
||||
B: T.Buffer((2, 2), "float32"),
|
||||
rank: T.int64,
|
||||
B_sharded: T.Buffer((2, 1), "float32"),
|
||||
):
|
||||
for i in range(2):
|
||||
with T.sblock("slice_B"):
|
||||
vi = T.axis.spatial(2, i)
|
||||
B_sharded[vi, 0] = B[vi, rank]
|
||||
|
||||
@R.function
|
||||
def transform_params(
|
||||
rank_arg: R.Prim("int64"),
|
||||
fget_item: R.Callable([R.Any, R.Prim("int64")], R.Any),
|
||||
):
|
||||
cls = Module
|
||||
|
||||
A = fget_item(R.str("A"), R.prim_value(0))
|
||||
A = R.match_cast(A, R.Tensor([4, 4], "int32"))
|
||||
A = R.call_tir(
|
||||
cls.slice_A,
|
||||
(A, rank_arg),
|
||||
out_ty=R.Tensor([2, 4], "int32"),
|
||||
)
|
||||
|
||||
B = fget_item(R.str("B"), R.prim_value(1))
|
||||
B = R.match_cast(B, R.Tensor([2, 2], "float32"))
|
||||
B = R.call_tir(
|
||||
cls.slice_B,
|
||||
(B, rank_arg),
|
||||
out_ty=R.Tensor([2, 1], "float32"),
|
||||
)
|
||||
|
||||
return (A, B)
|
||||
|
||||
pipeline = tvm.ir.transform.Sequential(
|
||||
[
|
||||
tvm.relax.transform.LegalizeOps(),
|
||||
tvm.s_tir.dlight.ApplyDefaultSchedule(tvm.s_tir.dlight.gpu.Fallback()),
|
||||
],
|
||||
name="pipeline",
|
||||
)
|
||||
|
||||
with tvm.target.Target("cuda"):
|
||||
mod = Module
|
||||
mod = pipeline(mod)
|
||||
built = tvm.compile(mod, "cuda")
|
||||
|
||||
num_shards = 2
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_dir = pathlib.Path(temp_dir)
|
||||
|
||||
# TODO(Lunderberg): Update `disco.Session.load_vm_module` to
|
||||
# allow a `tvm.runtime.Module` argument. This would avoid the
|
||||
# need for a temporary file.
|
||||
shlib_path = temp_dir.joinpath("libtemp.so")
|
||||
built.export_library(shlib_path)
|
||||
|
||||
def run_and_check():
|
||||
session = tvm.runtime.disco.ProcessSession(num_workers=num_shards)
|
||||
try:
|
||||
session.import_python_module("tvm.exec.disco_worker")
|
||||
session.init_ccl("nccl", *range(num_shards))
|
||||
|
||||
worker_device = session.get_global_func("runtime.disco.device")()
|
||||
worker_id = session.get_global_func("runtime.disco.worker_rank")()
|
||||
callback_maker = session.get_global_func("tests.disco.test_callback")
|
||||
fget_item = callback_maker(worker_device)
|
||||
vm = session.load_vm_module(shlib_path.as_posix())
|
||||
transform_params = vm["transform_params"]
|
||||
|
||||
params = transform_params(worker_id, fget_item)
|
||||
|
||||
# Worker 0 is the same PID as the controlling scope, so
|
||||
# `debug_get_from_remote(0)` returns the Tensor containing
|
||||
# the output.
|
||||
params_gpu0 = params.debug_get_from_remote(0)
|
||||
assert params_gpu0[0].device == tvm.cuda(0)
|
||||
assert params_gpu0[1].device == tvm.cuda(0)
|
||||
np.testing.assert_array_equal(
|
||||
params_gpu0[0].numpy(),
|
||||
[
|
||||
[0, 1, 2, 3],
|
||||
[4, 5, 6, 7],
|
||||
],
|
||||
)
|
||||
np.testing.assert_array_equal(
|
||||
params_gpu0[1].numpy(),
|
||||
[[0], [2]],
|
||||
)
|
||||
|
||||
# Worker 1 is a different PID altogether, so
|
||||
# `debug_get_from_remote(1)` returns a new Tensor within the
|
||||
# calling scope's PID.
|
||||
params_gpu1 = params.debug_get_from_remote(1)
|
||||
assert params_gpu1[0].device == tvm.cpu()
|
||||
assert params_gpu1[1].device == tvm.cpu()
|
||||
np.testing.assert_array_equal(
|
||||
params_gpu1[0].numpy(),
|
||||
[
|
||||
[8, 9, 10, 11],
|
||||
[12, 13, 14, 15],
|
||||
],
|
||||
)
|
||||
np.testing.assert_array_equal(
|
||||
params_gpu1[1].numpy(),
|
||||
[[1], [3]],
|
||||
)
|
||||
finally:
|
||||
session.shutdown()
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
Reference in New Issue
Block a user