github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2026-05-05 12:24:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

[refactor] move process group from _DistSpec to ColoTensor. (#1203)

Browse Source
This commit is contained in:
Jiarui Fang
2022-07-06 16:15:16 +08:00
committed by GitHub
parent 5da87ce35d
commit ae7d3f4927
34 changed files with 452 additions and 367 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
tests/test_tensor/test_parameter.py
View File

@@ -1,10 +1,16 @@
from colossalai.tensor import ColoParameter, ColoTensor
from colossalai.tensor import ColoParameter, ColoTensor, ColoTensorSpec, ProcessGroup
import torch
from numpy import allclose
import pytest
from _utils import tensor_equal
import colossalai
from colossalai.utils import free_port
@pytest.mark.skip
def test_multiinheritance():
colo_param = ColoParameter()
colossalai.launch(config={}, rank=0, world_size=1, host='localhost', port=free_port(), backend='nccl')
colo_param = ColoParameter(None, requires_grad=True)
assert colo_param.dist_spec.placement.value == 'r'
assert isinstance(colo_param, ColoTensor)
assert isinstance(colo_param, torch.nn.Parameter)
@@ -22,5 +28,6 @@ def test_multiinheritance():
clone_param = torch.clone(colo_param)
assert isinstance(clone_param, ColoTensor)
if __name__ == '__main__':
test_multiinheritance()
test_multiinheritance()