Function test_callback

tests/python/disco/test_callback.py:35–130

Simulate lazy loading of parameters in a callback The output of a lazy parameter loading, which would accept a callback to load the parameters.

()

@pytest.mark.gpu
@pytest.mark.skipif(not env.has_nccl(), reason="need nccl")
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.
    """

    @R.function
    def transform_params(
        rank_arg: R.Prim(value="rank"),
        fget_item: R.Callable([R.Object, R.Prim("int64")], R.Object),
    ):
        rank = T.int64()

        A = fget_item(R.str("A"), R.prim_value(0))
        A = R.match_cast(A, R.Tensor([4, 4], "int32"))
        A = R.strided_slice(A, axes=[0], begin=[rank * 2], end=[(rank + 1) * 2])

        B = fget_item(R.str("B"), R.prim_value(1))
        B = R.match_cast(B, R.Tensor([2, 2], "float32"))
        B = R.strided_slice(B, axes=[1], begin=[rank * 1], end=[(rank + 1) * 1])

        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 = tvm.IRModule.from_expr(transform_params)
        mod = pipeline(mod)
        built = tvm.compile(mod, "cuda")

    num_shards = 2

    session = tvm.runtime.disco.ProcessSession(num_workers=num_shards)
    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)

    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)
        vm = session.load_vm_module(shlib_path.as_posix())
        transform_params = vm["transform_params"]