Lightning-AI/pytorch-lightning · critical · RuntimeError

Skipping the `training_step` by returning None in distribute

Error message

Skipping the `training_step` by returning None in distributed training is not supported. It is recommended that you rewrite your training logic to avoid having to skip the step in the first place.

What it means

Raised in _AutomaticOptimization._training_step when training_step returns None while world_size > 1. Skipping a step in only some distributed ranks desynchronizes collective operations (all-reduce of gradients), hanging or corrupting DDP training, so Lightning forbids it.

Source

Thrown at src/lightning/pytorch/loops/optimization/automatic.py:323

        self.optim_progress.optimizer.zero_grad.increment_completed()

    def _training_step(self, kwargs: OrderedDict) -> ClosureResult:
        """Performs the actual train step with the tied hooks.

        Args:
            kwargs: the kwargs passed down to the hooks.

        Returns:
            A ``ClosureResult`` containing the training step output.

        """
        trainer = self.trainer

        training_step_output = call._call_strategy_hook(trainer, "training_step", *kwargs.values())
        self.trainer.strategy.post_training_step()  # unused hook - call anyway for backward compatibility

        if training_step_output is None and trainer.world_size > 1:
            raise RuntimeError(
                "Skipping the `training_step` by returning None in distributed training is not supported."
                " It is recommended that you rewrite your training logic to avoid having to skip the step in the first"
                " place."
            )

        return self.output_result_cls.from_training_step_output(training_step_output, trainer.accumulate_grad_batches)

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Rewrite training_step to always return a finite loss tensor (e.g. zero loss for skipped batches)
2. Ensure every rank sees the same number of batches (balanced DistributedSampler, no rank-local filtering)
3. Move conditional logic to dataloader-level filtering applied identically on all ranks

Example fix

# before
def training_step(self, batch, batch_idx):
    if batch[0].shape[0] < 2:
        return None  # hangs DDP
    return self.loss_fn(self(batch[0]), batch[1])

# after
def training_step(self, batch, batch_idx):
    loss = self.loss_fn(self(batch[0]), batch[1])
    if batch[0].shape[0] < 2:
        loss = loss * 0.0  # keep graph, stay in sync
    return loss

# before multi-GPU runs, ensure training_step never returns None:
def training_step(self, batch, batch_idx):
    loss = self.compute_loss(batch)
    assert loss is not None
    return loss

def always_returns_loss(fn, batch) -> bool:
    return fn(batch) is not None

When it happens

Trigger: `if some_condition: return None` inside training_step while running DDP/DeepSpeed/ddp_spawn with multiple devices; conditional data filtering that only triggers on some ranks; NaN guards that skip batches per-rank.

Common situations: Adding batch-skipping logic that worked single-GPU, then scaling to trainer = Trainer(devices=2, strategy='ddp'); imbalanced or rank-dependent data where one rank exhausts or filters batches earlier.

Related errors

• Trying to inject a modified sampler into the batch sampler;
• It is recommended to use `self.log({result_metric.meta.name!
• To use Fabric with more than one device, you must call `.lau
• `{type(self).__name__}` does not own a group. HINT: try `col
• `{type(self).__name__}` already owns a group.

AI-assisted analysis of Lightning-AI/pytorch-lightning@9fed5c27d2 (2026-08-28). Data as JSON: /api/errors/bbaf9a4e84c175ef. Report an issue: GitHub.