{"record":{"id":"5c1ced5d3a76a71c","repo":"Lightning-AI/pytorch-lightning","slug":"skipping-backward-by-returning-none-from-your-t","errorCode":null,"errorMessage":"Skipping backward by returning `None` from your `training_step` is not supported by `DeepSpeed`","messagePattern":"Skipping backward by returning `None` from your `training_step` is not supported by `DeepSpeed`","errorType":"exception","errorClass":"MisconfigurationException","httpStatus":null,"severity":"error","filePath":"src/lightning/pytorch/plugins/precision/deepspeed.py","lineNumber":134,"sourceCode":"        deepspeed_engine: deepspeed.DeepSpeedEngine = model.trainer.model\n        deepspeed_engine.backward(tensor, *args, **kwargs)\n\n    @override\n    def optimizer_step(  # type: ignore[override]\n        self,\n        optimizer: Steppable,\n        model: \"pl.LightningModule\",\n        closure: Callable[[], Any],\n        **kwargs: Any,\n    ) -> Any:\n        if isinstance(optimizer, LBFGS):\n            raise MisconfigurationException(\"DeepSpeed and the LBFGS optimizer are not compatible.\")\n        closure_result = closure()\n        self._after_closure(model, optimizer)\n        skipped_backward = closure_result is None\n        # in manual optimization, the closure does not return a value\n        if model.automatic_optimization and skipped_backward:\n            raise MisconfigurationException(\n                \"Skipping backward by returning `None` from your `training_step` is not supported by `DeepSpeed`\"\n            )\n        # DeepSpeed handles the optimizer step internally\n        deepspeed_engine: deepspeed.DeepSpeedEngine = model.trainer.model\n        return deepspeed_engine.step(**kwargs)\n\n    @override\n    def clip_gradients(\n        self,\n        optimizer: Optimizer,\n        clip_val: Union[int, float] = 0.0,\n        gradient_clip_algorithm: GradClipAlgorithmType = GradClipAlgorithmType.NORM,\n    ) -> None:\n        \"\"\"DeepSpeed handles gradient clipping internally.\"\"\"\n","sourceCodeStart":116,"sourceCodeEnd":149,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/pytorch/plugins/precision/deepspeed.py#L116-L149","documentation":"In automatic optimization, the training_step closure returned None, meaning backward was skipped, but DeepSpeed requires every step to have a computed loss for its internal engine bookkeeping. DeepSpeed's engine.step() assumes backward ran through its managed graph, so a skipped backward desynchronizes the engine.","triggerScenarios":"training_step returns None (or a loss that evaluates to None) under Trainer(strategy='deepspeed') with automatic_optimization=True; deepspeed's optimizer_step sees closure_result is None and raises.","commonSituations":"Conditionally returning None from training_step to skip batches (a pattern valid with vanilla DDP); RBG/curriculum logic that skips steps; refactoring code from single-device training to DeepSpeed where None-returns were tolerated.","solutions":["Always return a loss tensor from training_step; to skip a batch, return the loss multiplied by 0 or accumulate masks instead","Return loss = loss * valid_mask.float().mean() (or similar) so backward still runs","If you truly need skipped steps, switch to manual optimization and manage DeepSpeed's engine yourself"],"exampleFix":"# before\ndef training_step(self, batch, batch_idx):\n    x, y = batch\n    if x is None:\n        return None  # unsupported under DeepSpeed\n    loss = self.model(x, y).loss\n    return loss\n\n# after\ndef training_step(self, batch, batch_idx):\n    x, y = batch\n    loss = self.model(x, y).loss\n    return loss","handlingStrategy":"validation","validationCode":"# static check: training_step must not return None under deepspeed\nimport inspect\n\ndef training_step_returns_loss(cls) -> bool:\n    src = inspect.getsource(cls.training_step)\n    return 'return None' not in src\n\nif strategy == 'deepspeed':\n    assert training_step_returns_loss(type(model))","typeGuard":"def loss_is_trainable(loss) -> bool:\n    import torch\n    return isinstance(loss, torch.Tensor) and loss.requires_grad","tryCatchPattern":null,"preventionTips":["Never return None from training_step; skip batches by zero-weighting the loss","Add an assertion `assert loss is not None and loss.requires_grad` before the return"],"tags":["pytorch-lightning","deepspeed","training-step","backward","automatic-optimization"],"backgroundTag":"skipped-backward-unsupported","analyzedSha":"9fed5c27d2a62ff0efd6c3573599921d6ff67c14","analyzedAt":"2026-08-28T11:52:41.083Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}