[AMD CI] Add 2-GPU sgl-kernel Tests (#17555)
Co-authored-by: YC Tseng <yctseng@amd.com>
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user