{"record":{"id":"bbaf9a4e84c175ef","repo":"Lightning-AI/pytorch-lightning","slug":"skipping-the-training-step-by-returning-none-in","errorCode":null,"errorMessage":"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.","messagePattern":"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\\.","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"critical","filePath":"src/lightning/pytorch/loops/optimization/automatic.py","lineNumber":323,"sourceCode":"        self.optim_progress.optimizer.zero_grad.increment_completed()\n\n    def _training_step(self, kwargs: OrderedDict) -> ClosureResult:\n        \"\"\"Performs the actual train step with the tied hooks.\n\n        Args:\n            kwargs: the kwargs passed down to the hooks.\n\n        Returns:\n            A ``ClosureResult`` containing the training step output.\n\n        \"\"\"\n        trainer = self.trainer\n\n        training_step_output = call._call_strategy_hook(trainer, \"training_step\", *kwargs.values())\n        self.trainer.strategy.post_training_step()  # unused hook - call anyway for backward compatibility\n\n        if training_step_output is None and trainer.world_size > 1:\n            raise RuntimeError(\n                \"Skipping the `training_step` by returning None in distributed training is not supported.\"\n                \" It is recommended that you rewrite your training logic to avoid having to skip the step in the first\"\n                \" place.\"\n            )\n\n        return self.output_result_cls.from_training_step_output(training_step_output, trainer.accumulate_grad_batches)\n","sourceCodeStart":305,"sourceCodeEnd":330,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/pytorch/loops/optimization/automatic.py#L305-L330","documentation":"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.","triggerScenarios":"`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.","commonSituations":"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.","solutions":["Rewrite training_step to always return a finite loss tensor (e.g. zero loss for skipped batches)","Ensure every rank sees the same number of batches (balanced DistributedSampler, no rank-local filtering)","Move conditional logic to dataloader-level filtering applied identically on all ranks"],"exampleFix":"# before\ndef training_step(self, batch, batch_idx):\n    if batch[0].shape[0] < 2:\n        return None  # hangs DDP\n    return self.loss_fn(self(batch[0]), batch[1])\n\n# after\ndef training_step(self, batch, batch_idx):\n    loss = self.loss_fn(self(batch[0]), batch[1])\n    if batch[0].shape[0] < 2:\n        loss = loss * 0.0  # keep graph, stay in sync\n    return loss","handlingStrategy":"validation","validationCode":"# before multi-GPU runs, ensure training_step never returns None:\ndef training_step(self, batch, batch_idx):\n    loss = self.compute_loss(batch)\n    assert loss is not None\n    return loss","typeGuard":"def always_returns_loss(fn, batch) -> bool:\n    return fn(batch) is not None","tryCatchPattern":null,"preventionTips":["Never conditionally return None from training_step in distributed training","Test with Trainer(devices=2, strategy='ddp') locally before cluster runs","Filter data at the dataset/sampler level uniformly across ranks"],"tags":["pytorch-lightning","distributed","ddp","training-step","sync"],"backgroundTag":"distributed-step-desync","analyzedSha":"9fed5c27d2a62ff0efd6c3573599921d6ff67c14","analyzedAt":"2026-08-28T11:52:41.083Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}