diff --git a/specforge/cli.py b/specforge/cli.py index 92fed155a..2456bba88 100644 --- a/specforge/cli.py +++ b/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: diff --git a/specforge/distributed.py b/specforge/distributed.py index d8e7d00ca..7f4643165 100644 --- a/specforge/distributed.py +++ b/specforge/distributed.py @@ -210,7 +210,7 @@ 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 @@ -218,8 +218,12 @@ def destroy_distributed(): # 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 diff --git a/tests/test_runtime/test_npu_portability.py b/tests/test_runtime/test_npu_portability.py index 906005785..2777f4e67 100644 --- a/tests/test_runtime/test_npu_portability.py +++ b/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):