wehub-resource-sync
59 lines
1.9 KiB
Python
59 lines
1.9 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import pytest
|
|
import torch
|
|
import deepspeed
|
|
from deepspeed.ops.op_builder import InferenceBuilder
|
|
|
|
if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]:
|
|
pytest.skip("Inference ops are not available on this system", allow_module_level=True)
|
|
|
|
|
|
def allclose(x, y):
|
|
assert x.dtype == y.dtype
|
|
rtol, atol = {torch.float32: (5e-4, 5e-5), torch.float16: (5e-2, 2e-3)}[x.dtype]
|
|
return torch.allclose(x, y, rtol=rtol, atol=atol)
|
|
|
|
|
|
def run_matmul_ref(a, b):
|
|
return torch.matmul(a, b)
|
|
|
|
|
|
def run_matmul_ds(a, b, use_triton_ops=False):
|
|
if use_triton_ops:
|
|
from deepspeed.ops.transformer.inference.triton import matmul_4d as matmul
|
|
return matmul(a, b)
|
|
|
|
assert use_triton_ops, "Only triton softmax is supported for now"
|
|
|
|
|
|
@pytest.mark.inference_ops
|
|
@pytest.mark.parametrize("B", [1, 2])
|
|
@pytest.mark.parametrize("H", [1, 2, 16])
|
|
@pytest.mark.parametrize("M", [1, 7, 8, 128])
|
|
@pytest.mark.parametrize("K", [2, 5, 16, 128])
|
|
@pytest.mark.parametrize("N", [1, 2, 8, 512])
|
|
@pytest.mark.parametrize("dtype", [torch.float16])
|
|
@pytest.mark.parametrize("use_triton_ops", [True])
|
|
def test_matmul_4d(B, H, M, K, N, dtype, use_triton_ops):
|
|
if not deepspeed.get_accelerator().is_triton_supported():
|
|
pytest.skip("triton is not supported on this system")
|
|
if not deepspeed.HAS_TRITON:
|
|
pytest.skip("triton is not installed")
|
|
|
|
# skip autotune in testing
|
|
from deepspeed.ops.transformer.inference.triton.matmul_ext import fp16_matmul
|
|
fp16_matmul.skip_autotune()
|
|
|
|
a_ds = torch.randn((B, H, M, K), dtype=dtype, device='cuda')
|
|
b_ds = torch.randn((B, H, K, N), dtype=dtype, device='cuda')
|
|
a_ref = a_ds.clone().detach()
|
|
b_ref = b_ds.clone().detach()
|
|
|
|
ds_out = run_matmul_ds(a_ds, b_ds, use_triton_ops)
|
|
ref_out = run_matmul_ref(a_ref, b_ref)
|
|
assert (allclose(ds_out, ref_out))
|