Function check_4gpu

tests/test_fp8/test_fp8_all_to_all_single.py:18–24

(shape, dtype, fp8_format)

@parameterize("dtype", [torch.bfloat16, torch.float16])
@parameterize("fp8_format", ["e4m3", "e5m2"])
def check_4gpu(shape, dtype, fp8_format):
    x = torch.rand(shape, dtype=dtype, device=get_accelerator().get_current_device())
    output = torch.empty_like(x)
    output_fp8 = torch.empty_like(x)
    all_to_all_single_fp8(output_fp8, x, group=_get_default_group(), fp8_format=fp8_format)
    dist.all_to_all_single(output, x, group=_get_default_group())
    assert_close(output, output_fp8, rtol=0.1, atol=0.1)


def run_dist(rank, world_size, port):