github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-30 15:01:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

[NFC] polish colossalai/testing/comparison.py code style. (#1558)

Browse Source
This commit is contained in:
Super Daniel 2022-09-08 15:25:06 +08:00 committed by Frank Lee
parent 7cc052f6c0
commit 8328917348
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
colossalai/testing/comparison.py
View File

@ -7,15 +7,19 @@ from torch.distributed import ProcessGroup
def assert_equal(a: Tensor, b: Tensor):
assert torch.all(a == b), f'expected a and b to be equal but they are not, {a} vs {b}'
def assert_not_equal(a: Tensor, b: Tensor):
assert not torch.all(a == b), f'expected a and b to be not equal but they are, {a} vs {b}'
def assert_close(a: Tensor, b: Tensor, rtol: float = 1e-5, atol: float = 1e-8):
assert torch.allclose(a, b, rtol=rtol, atol=atol), f'expected a and b to be close but they are not, {a} vs {b}'
def assert_close_loose(a: Tensor, b: Tensor, rtol: float = 1e-3, atol: float = 1e-3):
assert_close(a, b, rtol, atol)
def assert_equal_in_group(tensor: Tensor, process_group: ProcessGroup = None):
# all gather tensors from different ranks
world_size = dist.get_world_size(process_group)
@ -25,5 +29,5 @@ def assert_equal_in_group(tensor: Tensor, process_group: ProcessGroup = None):
# check if they are equal one by one
for i in range(world_size - 1):
a = tensor_list[i]
b = tensor_list[i+1]
b = tensor_list[i + 1]
assert torch.all(a == b), f'expected tensors on rank {i} and {i+1} to be equal but they are not, {a} vs {b}'