add test for early return with dp size 1

Signed-off-by: Hao Wu <skyw@nvidia.com>

Author: skyw

tests/unit_tests/distributed/test_param_and_grad_buffer.py

@@ -382,7 +382,8 @@ def test_start_param_sync_dp_size_1():
     """When dp_size == 1 (e.g., expt_dp_size == 1), start_param_sync should set
     param_gather_dispatched=True and return immediately without launching any
     all-gather collective."""
-    Utils.initialize_model_parallel(tensor_model_parallel_size=8)
+    world_size = torch.distributed.get_world_size()
+    Utils.initialize_model_parallel(tensor_model_parallel_size=world_size)
 
     ddp_config = DistributedDataParallelConfig(
         grad_reduce_in_fp32=True,