fix: abort process groups on rank-local training failure

specforge/cli.py

@@ -135,15 +135,22 @@ def _train(resolved) -> int:
         sp_ulysses_size=cfg.training.sp_ulysses_size,
         sp_ring_size=cfg.training.sp_ring_size,
     )
+    failed = True
     try:
         import torch.distributed as dist
 
         _validate_world_size(cfg, dist.get_world_size())
         from specforge.application import build_application_run
 
-        return build_application_run(resolved).run()
+        result = build_application_run(resolved).run()
+        failed = False
+        return result
     finally:
-        destroy_distributed()
+        # A rank-local data/capture error can leave peers inside a CUDA or FSDP
+        # collective. Collective NCCL destruction then hides the real exception
+        # and prevents torchrun from terminating the peers. Abort communicators
+        # on the exceptional path so the originating traceback reaches elastic.
+        destroy_distributed(abort=failed)
 
 
 def _config_for_role(cfg: Config, role: str) -> Config:

specforge/distributed.py

@@ -210,16 +210,20 @@ def init_distributed(
     _DP_DEVICE_MESH = dist.DeviceMesh.from_group(dp_group, device_type=device_type)
 
 
-def destroy_distributed():
+def destroy_distributed(*, abort: bool = False):
     global _DEVICE_MESH, _TP_DEVICE_MESH, _TP_GROUP
     global _DP_DEVICE_MESH, _DP_GROUP, _DRAFT_DP_GROUP, _DRAFT_SP_GROUP
     global _SP_ULYSSES_GROUP, _SP_RING_GROUP
     # Teardown must never crash the process. Several handles can alias the same
     # underlying group (e.g. DP and draft-DP when there is no sequence
     # parallelism), and degenerate single-rank SP groups (created when
     # sp_ulysses_size == 1 or sp_ring_size == 1) are not registered in torch's
-    # process-group map and would raise on destroy. Destroy each distinct, valid
-    # sub-group at most once, then tear down the default group.
+    # process-group map and would raise on destroy. Clean up each distinct,
+    # valid group at most once. On rank-local failure, ProcessGroup.abort() is
+    # intentionally non-collective: peers may be blocked inside an NCCL
+    # collective, and a collective destroy would hang there and hide the
+    # originating exception from torchrun/elastic.
+    default_group = dist.group.WORLD if dist.is_initialized() else None
     seen = set()
     for group in (
         _TP_GROUP,
@@ -228,23 +232,23 @@ def destroy_distributed():
         _SP_RING_GROUP,
         _DRAFT_DP_GROUP,
         _DRAFT_SP_GROUP,
+        # The all-ranks DP group may alias the default group; the seen-set
+        # keeps the default group from being torn down twice.
+        default_group,
     ):
         if group is None or id(group) in seen:
             continue
         seen.add(id(group))
         try:
-            dist.destroy_process_group(group)
+            abort_group = getattr(group, "abort", None)
+            if abort and callable(abort_group):
+                abort_group()
+            else:
+                dist.destroy_process_group(group)
         except Exception:
             # Group not registered (e.g. degenerate single-rank SP group) or
             # already destroyed.
             pass
-    # The all-ranks DP group may alias the default group, in which case
-    # destroying it above already tore the default group down.
-    if dist.is_initialized():
-        try:
-            dist.destroy_process_group()
-        except Exception:
-            pass
 
     # Process-group and DeviceMesh objects are invalid after teardown. Keeping
     # them reachable makes a later single-process load look initialized while

tests/test_runtime/test_npu_portability.py

@@ -145,6 +145,34 @@ def test_destroy_clears_every_cached_group_and_mesh(self):
             for name, value in saved.items():
                 setattr(sf_dist, name, value)
 
+    def test_failure_teardown_aborts_each_distinct_process_group(self):
+        names = NPUDistributedTest._GLOBAL_NAMES
+        saved = {name: getattr(sf_dist, name) for name in names}
+        subgroup = mock.Mock()
+        world = mock.Mock()
+        try:
+            for name in names:
+                setattr(sf_dist, name, None)
+            sf_dist._TP_GROUP = subgroup
+            sf_dist._DP_GROUP = subgroup
+            with (
+                mock.patch.object(sf_dist.dist, "is_initialized", return_value=True),
+                mock.patch.object(
+                    sf_dist.dist,
+                    "group",
+                    SimpleNamespace(WORLD=world),
+                ),
+                mock.patch.object(sf_dist.dist, "destroy_process_group") as destroy,
+            ):
+                sf_dist.destroy_distributed(abort=True)
+
+            subgroup.abort.assert_called_once_with()
+            world.abort.assert_called_once_with()
+            destroy.assert_not_called()
+        finally:
+            for name, value in saved.items():
+                setattr(sf_dist, name, value)
+
 
 class NPURNGTest(unittest.TestCase):
     def test_checkpoint_round_trip_uses_bound_npu_rng(self):