Dev/rebase 0.14.0 (#847)

Signed-off-by: JaredforReal <w13431838023@gmail.com>
Signed-off-by: root <root@hk01dgx039.cm.cluster>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
Signed-off-by: TangPeng <85704592@qq.com>
Signed-off-by: XU Mingshi <91017482+mxuax@users.noreply.github.com>
Signed-off-by: mxuax <mxuax@connect.ust.hk>
Signed-off-by: Sihyeon Jang <sihyeon.jang@navercorp.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Signed-off-by: iwzbi <wzbi@zju.edu.cn>
Signed-off-by: ZeldaHuang <hzm414167@alibaba-inc.com>
Signed-off-by: Yuhan Liu <30294295+liuyuhanalex@users.noreply.github.com>
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
Signed-off-by: samithuang <285365963@qq.com>
Signed-off-by: linyueqian <linyueqian@outlook.com>
Signed-off-by: David Chen <530634352@qq.com>
Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Signed-off-by: wangyu31577 <wangyu31577@hundsun.com>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: John Liu BUAA <liukecheng97@gmail.com>
Signed-off-by: Dinesh G <G.Dinesh@ibm.com>
Signed-off-by: gDINESH13 <dinesh13g@gmail.com>
Signed-off-by: root <root@hk01dgx028.cm.cluster>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
Co-authored-by: JaredforReal <w13431838023@gmail.com>
Co-authored-by: root <root@hk01dgx039.cm.cluster>
Co-authored-by: JustQJ <37905360+JustQJ@users.noreply.github.com>
Co-authored-by: XU Mingshi <91017482+mxuax@users.noreply.github.com>
Co-authored-by: Sihyeon Jang <uneedsihyeon@gmail.com>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
Co-authored-by: catcat <108673086+iwzbi@users.noreply.github.com>
Co-authored-by: Ziming Huang <hzm414167@alibaba-inc.com>
Co-authored-by: Yuhan Liu <30294295+liuyuhanalex@users.noreply.github.com>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>
Co-authored-by: Samit <285365963@qq.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
Co-authored-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: wangyu <53896905+yenuo26@users.noreply.github.com>
Co-authored-by: wangyu31577 <wangyu31577@hundsun.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: John Liu BUAA <liukecheng97@gmail.com>
Co-authored-by: D!NE$H <67671800+gDINESH13@users.noreply.github.com>
Co-authored-by: root <root@hk01dgx028.cm.cluster>

tests/diffusion/distributed/test_comm.py

@@ -48,6 +48,11 @@ def test_4d_identity(
     use_sync: bool,
 ):
     """Test that two consecutive all-to-all operations return the original input."""
+    # Skip if not enough GPUs available
+    available_gpus = torch_device.device_count()
+    if available_gpus < world_size:
+        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
+
     # Ensure num_heads is divisible by world_size
     if num_heads % world_size != 0:
         pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")
@@ -177,6 +182,11 @@ def test_5d_identity(
     use_sync: bool,
 ):
     """Test that two consecutive all-to-all operations return the original input."""
+    # Skip if not enough GPUs available
+    available_gpus = torch_device.device_count()
+    if available_gpus < world_size:
+        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
+
     # Ensure num_heads is divisible by world_size
     if num_heads % world_size != 0:
         pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")
@@ -305,6 +315,11 @@ def test_ring_p2p(
     head_size: int,
 ):
     """Test Ring P2P communication (send_recv)."""
+    # Skip if not enough GPUs available
+    available_gpus = torch_device.device_count()
+    if available_gpus < world_size:
+        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
+
     torch.multiprocessing.spawn(
         _test_ring_p2p_worker,
         args=(world_size, dtype, batch_size, num_heads, head_size),