[compatibility] support torch 2.2 (#5875)

* Support Pytorch 2.2.2 * keep build_on_pr file and update .compatibility
hpcaitech · Jul 15, 2024 · d46c7a6 · d46c7a6
1 parent 1c961b2
commit d46c7a6
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/.compatibility b/.compatibility
@@ -1 +1,2 @@
 2.1.0-12.1.0
+2.2.2-12.1.0
diff --git a/colossalai/tensor/d_tensor/layout_converter.py b/colossalai/tensor/d_tensor/layout_converter.py
@@ -473,7 +473,7 @@ def _group_alive_check(cached_comm_action_sequence):
                 for process_group in used_process_groups:
                     try:
                         dist.get_rank(process_group)
-                    except RuntimeError as e:
+                    except (ValueError, RuntimeError) as e:
                         # If the group is not registered, it means it has been deleted
                         if str(e) == (
                             f"Group {process_group} is not registered, please create group with torch.distributed.new_group API"