Add broadcast_sigterm_every_n_steps Trainer parameter

Allow users to control how often SIGTERM status is broadcast across
ranks, reducing CPU-GPU sync overhead for fast training loops while
preserving the default every-step behavior.

Closes #21487

Author: c-pozzi

src/lightning/pytorch/loops/training_epoch_loop.py

@@ -94,6 +94,7 @@ def __init__(self, trainer: "pl.Trainer", min_steps: Optional[int] = None, max_s
         self._batches_that_stepped: int = 0
         self._restart_stage = RestartStage.NONE
         self._skip_next_val = False
+        self._sigterm_broadcast_step: int = 0
 
     @property
     def total_batch_idx(self) -> int:
@@ -297,7 +298,10 @@ def advance(self, data_fetcher: _DataFetcher) -> None:
         # =====================================================================
 
         if torch.distributed.is_available() and torch.distributed.is_initialized() and self.trainer.world_size > 1:
-            self._broadcast_sigterm_tensor()
+            self._sigterm_broadcast_step += 1
+            if self._sigterm_broadcast_step >= self.trainer.broadcast_sigterm_every_n_steps:
+                self._sigterm_broadcast_step = 0
+                self._broadcast_sigterm_tensor()
 
         # =====================================================================
 

src/lightning/pytorch/trainer/trainer.py

@@ -129,6 +129,7 @@ def __init__(
         plugins: Optional[Union[_PLUGIN_INPUT, list[_PLUGIN_INPUT]]] = None,
         sync_batchnorm: bool = False,
         reload_dataloaders_every_n_epochs: int = 0,
+        broadcast_sigterm_every_n_steps: int = 1,
         default_root_dir: Optional[_PATH] = None,
         enable_autolog_hparams: bool = True,
         model_registry: Optional[str] = None,
@@ -300,6 +301,13 @@ def __init__(
             reload_dataloaders_every_n_epochs: Set to a positive integer to reload dataloaders every n epochs.
                 Default: ``0``.
 
+            broadcast_sigterm_every_n_steps: How often (in training steps) to broadcast SIGTERM status across
+                ranks in distributed training. The default ``1`` broadcasts every step. Higher values reduce
+                the overhead of the NCCL broadcast at the cost of increased SIGTERM detection latency
+                (worst case: ``(N-1) * step_time``). This is useful for fast training loops where the
+                per-step broadcast cost is significant relative to the step time.
+                Default: ``1``.
+
             default_root_dir: Default path for logs and weights when no logger/ckpt_callback passed.
                 Default: ``os.getcwd()``.
                 Can be remote file paths such as `s3://mybucket/path` or 'hdfs://path/'
@@ -447,6 +455,11 @@ def __init__(
         self.predict_loop = _PredictionLoop(self, inference_mode=inference_mode)
 
         self.accumulate_grad_batches = accumulate_grad_batches
+        if broadcast_sigterm_every_n_steps < 1:
+            raise ValueError(
+                f"`broadcast_sigterm_every_n_steps` must be >= 1, got {broadcast_sigterm_every_n_steps}."
+            )
+        self.broadcast_sigterm_every_n_steps = broadcast_sigterm_every_n_steps
 
         # init callbacks
         # Declare attributes to be set in _callback_connector on_trainer_init

tests/tests_pytorch/loops/test_training_epoch_loop.py

@@ -228,3 +228,36 @@ def train_and_resume(dataloader, resume_step, expected_warning):
 
     # Resume mid-epoch, stateful dataloader -> no warning
     train_and_resume(dataloader=StatefulIterable(), resume_step=1, expected_warning=False)
+
+
+def test_broadcast_sigterm_every_n_steps_validation():
+    """Test that invalid values for broadcast_sigterm_every_n_steps are rejected."""
+    with pytest.raises(ValueError, match="broadcast_sigterm_every_n_steps` must be >= 1"):
+        Trainer(broadcast_sigterm_every_n_steps=0)
+    with pytest.raises(ValueError, match="broadcast_sigterm_every_n_steps` must be >= 1"):
+        Trainer(broadcast_sigterm_every_n_steps=-1)
+
+
+def test_broadcast_sigterm_every_n_steps_default():
+    """Test that the default value broadcasts every step."""
+    trainer = Trainer()
+    assert trainer.broadcast_sigterm_every_n_steps == 1
+
+
+@pytest.mark.parametrize("n_steps", [1, 5, 10])
+def test_broadcast_sigterm_interval(n_steps):
+    """Test that _broadcast_sigterm_tensor is called at the correct interval."""
+    trainer = Trainer(broadcast_sigterm_every_n_steps=n_steps)
+    epoch_loop = trainer.fit_loop.epoch_loop
+
+    total_steps = 20
+    broadcast_call_count = 0
+
+    for _ in range(total_steps):
+        epoch_loop._sigterm_broadcast_step += 1
+        if epoch_loop._sigterm_broadcast_step >= trainer.broadcast_sigterm_every_n_steps:
+            epoch_loop._sigterm_broadcast_step = 0
+            broadcast_call_count += 1
+
+    assert broadcast_call_count == total_steps // n_steps
+    assert epoch_loop._sigterm_broadcast_step == total_steps % n_steps