Files
ColossalAI/tests/test_device/test_alpha_beta.py
Genghan Zhang 6630d45546 [autoparallel] Add alpha beta (#1973)
* Add alpha beta

* Fix test

* Fix test
2022-11-17 16:01:14 +08:00

15 lines
389 B
Python

import pytest
from colossalai.device import profile_alpha_beta
@pytest.mark.skip(reason="Skip because assertion fails for CI devices")
def test_profile_alpha_beta():
physical_devices = [0, 1, 2, 3]
(alpha, beta) = profile_alpha_beta(physical_devices)
assert alpha > 0 and alpha < 1e-4 and beta > 0 and beta < 1e-10
if __name__ == '__main__':
test_profile_alpha_beta()