using pytest parametrize

tests/test_zero_data_parallel/test_zero_param_mgr.py

@@ -1,41 +1,39 @@
 #!/usr/bin/env python
 # -*- encoding: utf-8 -*-
 
-import os
 from functools import partial
-from pathlib import Path
-
-import colossalai
 import pytest
+
 import torch
 import torch.multiprocessing as mp
+
+import colossalai
 from colossalai.zero.sharded_model.param_manager import Zero3ParameterManager
 from colossalai.core import global_context as gpc
 from colossalai.context.parallel_mode import ParallelMode
 from colossalai.utils import free_port
 from common import CONFIG
 
+
 def run_shard_shape_check(rank, world_size, port):
-    colossalai.launch(config=CONFIG,
-                      rank=rank,
-                      world_size=world_size,
-                      host='localhost',
-                      port=port,
-                      backend='nccl')
-    
+    colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+
     model = torch.nn.Linear(2, 4 * world_size)
     gpc.init_parallel_groups()
-    Zero3ParameterManager(module=model, process_group=gpc.get_group(ParallelMode.DATA), offload_config=CONFIG.get('offload_param_config'))
+    Zero3ParameterManager(module=model,
+                          process_group=gpc.get_group(ParallelMode.DATA),
+                          offload_config=CONFIG.get('offload_param_config'))
 
-    assert(model.weight.numel() == 4 * 2)
-    assert(model.bias.numel() == 4)
+    assert (model.weight.numel() == 4 * 2)
+    assert (model.bias.numel() == 4)
 
 
 @pytest.mark.dist
-def test_run_shard_shape():
-    world_size = 2
+@pytest.mark.parametrize("world_size", [1, 2, 4])
+def test_run_shard_shape(world_size):
     run_func = partial(run_shard_shape_check, world_size=world_size, port=free_port())
     mp.spawn(run_func, nprocs=world_size)
 
+
 if __name__ == '__main__':
-    test_run_shard_shape()
+    test_run_shard_shape(2)