[AMD CI] Add 2-GPU sgl-kernel Tests (#17555)

Co-authored-by: YC Tseng <yctseng@amd.com>
This commit is contained in:
Bingxu Chen
2026-01-22 21:48:52 -08:00
committed by GitHub
co-authored by YC Tseng
parent c0b5a180fe
commit 50a2e4345a
3 changed files with 64 additions and 2 deletions
@@ -12,12 +12,13 @@ This test compares:
2. Deterministic kernel (different batch size)
Usage:
python test_amd_deterministic_custom_allreduce.py
pytest test_amd_deterministic_custom_allreduce.py
"""
import multiprocessing as mp
import socket
import pytest
import torch
import torch.distributed as dist
@@ -266,5 +267,14 @@ def main():
p.join()
@pytest.mark.skipif(
not torch.cuda.is_available() or torch.cuda.device_count() < 2,
reason="Requires at least 2 CUDA GPUs",
)
def test_deterministic_custom_allreduce():
"""Test that deterministic custom all-reduce produces consistent results."""
main()
if __name__ == "__main__":
main()
@@ -10,12 +10,13 @@ This test compares:
2. Default all-reduce (different batch size) - typically NON-DETERMINISTIC for bfloat16
Usage:
python test_ar.py
pytest test_amd_nccl_allreduce_determinism.py
"""
import multiprocessing as mp
import socket
import pytest
import torch
import torch.distributed as dist
@@ -194,5 +195,14 @@ def main():
p.join()
@pytest.mark.skipif(
not torch.cuda.is_available() or torch.cuda.device_count() < 2,
reason="Requires at least 2 CUDA GPUs",
)
def test_nccl_allreduce_determinism():
"""Test NCCL all-reduce determinism behavior with varying batch sizes."""
main()
if __name__ == "__main__":
main()