github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[test] fix test: test_zero1_2

moe_sp
hxwang 2024-07-22 05:36:20 +00:00
parent c67e553fd3
commit e31d2ebcf7
No known key found for this signature in database
GPG Key ID: 0EC383D418F0B9F8
3 changed files with 5 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/zero/low_level/low_level_optim.py
View File

@ -880,7 +880,7 @@ class LowLevelZeroOptimizer(OptimizerWrapper):
return None return None
grad_flat = torch.empty((grad_store.world_size, *grad.shape), dtype=grad.dtype, device=grad.device) grad_flat = torch.empty((grad_store.world_size, *grad.shape), dtype=grad.dtype, device=grad.device)
dist.all_gather_into_tensor(grad_flat, grad, group=grad_store.torch_pg) dist.all_gather_into_tensor(grad_flat, grad, group=grad_store.torch_pg)
return grad_flat[: working_param.numel()].reshape_as(working_param) return grad_flat.view(-1)[: working_param.numel()].view_as(working_param)
def get_working_grads_by_group_id(self, group_id: int) -> List[Tensor]: def get_working_grads_by_group_id(self, group_id: int) -> List[Tensor]:
working_grads = [] working_grads = []

4
tests/test_shardformer/test_model/test_shard_deepseek.py
View File

@ -179,9 +179,9 @@ def run_dist(rank, world_size, port):
@pytest.mark.dist @pytest.mark.dist
@pytest.mark.parametrize("world_size", [4]) @pytest.mark.parametrize("world_size", [4])
@rerun_if_address_is_in_use() @rerun_if_address_is_in_use()
def test_mistral(world_size): def test_deepseek(world_size):
spawn(run_dist, world_size) spawn(run_dist, world_size)
if __name__ == "__main__": if __name__ == "__main__":
test_mistral(world_size=8) test_deepseek(world_size=4)

4
tests/test_shardformer/test_model/test_shard_mixtral.py
View File

@ -180,9 +180,9 @@ def run_dist(rank, world_size, port):
@pytest.mark.dist @pytest.mark.dist
@pytest.mark.parametrize("world_size", [4]) @pytest.mark.parametrize("world_size", [4])
@rerun_if_address_is_in_use() @rerun_if_address_is_in_use()
def test_mistral(world_size): def test_mixtral(world_size):
spawn(run_dist, world_size) spawn(run_dist, world_size)
if __name__ == "__main__": if __name__ == "__main__":
test_mistral(world_size=8) test_mixtral(world_size=4)