Lightning-AI/pytorch-lightning · error · MisconfigurationException

Skipping backward by returning `None` from your `training_st

Error message

Skipping backward by returning `None` from your `training_step` is not supported by `DeepSpeed`

What it means

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.

Source

Thrown at src/lightning/pytorch/plugins/precision/deepspeed.py:134

        deepspeed_engine: deepspeed.DeepSpeedEngine = model.trainer.model
        deepspeed_engine.backward(tensor, *args, **kwargs)

    @override
    def optimizer_step(  # type: ignore[override]
        self,
        optimizer: Steppable,
        model: "pl.LightningModule",
        closure: Callable[[], Any],
        **kwargs: Any,
    ) -> Any:
        if isinstance(optimizer, LBFGS):
            raise MisconfigurationException("DeepSpeed and the LBFGS optimizer are not compatible.")
        closure_result = closure()
        self._after_closure(model, optimizer)
        skipped_backward = closure_result is None
        # in manual optimization, the closure does not return a value
        if model.automatic_optimization and skipped_backward:
            raise MisconfigurationException(
                "Skipping backward by returning `None` from your `training_step` is not supported by `DeepSpeed`"
            )
        # DeepSpeed handles the optimizer step internally
        deepspeed_engine: deepspeed.DeepSpeedEngine = model.trainer.model
        return deepspeed_engine.step(**kwargs)

    @override
    def clip_gradients(
        self,
        optimizer: Optimizer,
        clip_val: Union[int, float] = 0.0,
        gradient_clip_algorithm: GradClipAlgorithmType = GradClipAlgorithmType.NORM,
    ) -> None:
        """DeepSpeed handles gradient clipping internally."""

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Always return a loss tensor from training_step; to skip a batch, return the loss multiplied by 0 or accumulate masks instead
2. Return loss = loss * valid_mask.float().mean() (or similar) so backward still runs
3. If you truly need skipped steps, switch to manual optimization and manage DeepSpeed's engine yourself

Example fix

# before
def training_step(self, batch, batch_idx):
    x, y = batch
    if x is None:
        return None  # unsupported under DeepSpeed
    loss = self.model(x, y).loss
    return loss

# after
def training_step(self, batch, batch_idx):
    x, y = batch
    loss = self.model(x, y).loss
    return loss

# static check: training_step must not return None under deepspeed
import inspect

def training_step_returns_loss(cls) -> bool:
    src = inspect.getsource(cls.training_step)
    return 'return None' not in src

if strategy == 'deepspeed':
    assert training_step_returns_loss(type(model))

def loss_is_trainable(loss) -> bool:
    import torch
    return isinstance(loss, torch.Tensor) and loss.requires_grad

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• In automatic_optimization, when `training_step` returns a di
• In automatic optimization, `training_step` must return a Ten
• Skipping backward by returning `None` from your `training_st
• No models were set up for backward. Did you forget to call `
• When using multiple models + deepspeed, please provide the m

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