wehub-resource-sync
51 lines
1.7 KiB
Python
51 lines
1.7 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: (3e-2, 2e-3)}[x.dtype]
|
|
return torch.allclose(x, y, rtol=rtol, atol=atol)
|
|
|
|
|
|
def run_softmax_reference(input):
|
|
return torch.nn.functional.softmax(input, dim=-1)
|
|
|
|
|
|
def run_softmax_ds(input, use_triton_ops=False):
|
|
if use_triton_ops:
|
|
from deepspeed.ops.transformer.inference.triton import softmax
|
|
# return torch.empty_like(input)
|
|
return softmax(input)
|
|
|
|
assert use_triton_ops, "Only triton softmax is supported for now"
|
|
|
|
|
|
@pytest.mark.inference_ops
|
|
@pytest.mark.parametrize("batch", [1, 2])
|
|
@pytest.mark.parametrize("sequence", [1, 128, 255, 1232])
|
|
@pytest.mark.parametrize("channels", [512, 4096])
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
|
|
@pytest.mark.parametrize("use_triton_ops", [True])
|
|
def test_softmax(batch, sequence, channels, dtype, use_triton_ops):
|
|
if not deepspeed.get_accelerator().is_triton_supported():
|
|
pytest.skip("triton is not supported on this system")
|
|
|
|
device = deepspeed.accelerator.get_accelerator().device_name()
|
|
input_ds = torch.randn((batch, sequence, channels), dtype=dtype, device=device)
|
|
input_ref = input_ds.clone().detach()
|
|
|
|
ds_out = run_softmax_ds(input_ds, use_triton_ops)
|
|
ref_out = run_softmax_reference(input_ref)
|
|
assert (allclose(ds_out, ref_out))
|