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 implemented with XLA. Please, open an issue in `https://github.com/Lightning-AI/pytorch-lightning/issues` requesting this feature.

What it means

In automatic optimization, Lightning detects that training_step returned None (i.e. you skipped backward). The XLA precision plugin's optimizer_step explicitly rejects this because the XLA graph execution path (xm.mark_step) does not support skipping backward. The maintainers ask users to open a feature request if they need it.

Source

Thrown at src/lightning/pytorch/plugins/precision/xla.py:80

    @override
    def optimizer_step(  # type: ignore[override]
        self,
        optimizer: Optimizable,
        model: "pl.LightningModule",
        closure: Callable[[], Any],
        **kwargs: Any,
    ) -> Any:
        import torch_xla.core.xla_model as xm

        closure = partial(self._xla_wrap_closure, optimizer, closure)
        closure = partial(self._wrap_closure, model, optimizer, closure)
        closure_result = optimizer.step(closure=closure, **kwargs)
        xm.mark_step()
        skipped_backward = closure_result is None
        # in manual optimization, the closure does not return a value
        if model.automatic_optimization and skipped_backward:
            # we lack coverage here so disable this - something to explore if there's demand
            raise MisconfigurationException(
                "Skipping backward by returning `None` from your `training_step` is not implemented with XLA."
                " Please, open an issue in `https://github.com/Lightning-AI/pytorch-lightning/issues`"
                " requesting this feature."
            )
        return closure_result

    @override
    def teardown(self) -> None:
        os.environ.pop("XLA_USE_BF16", None)
        os.environ.pop("XLA_USE_F16", None)

    def _xla_wrap_closure(self, optimizer: Optimizable, closure: Callable[[], Any]) -> Any:
        import torch_xla.core.xla_model as xm

        closure_result = closure()
        xm.reduce_gradients(optimizer)
        return closure_result

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Always return a loss tensor from training_step (never None) when using XLA with automatic optimization
2. If you must skip batches, gate the batch in train_dataloader/on_train_batch_start instead of returning None
3. If you genuinely need skipped backward on XLA, open the referenced GitHub issue requesting the feature

Example fix

# before
def training_step(self, batch, batch_idx):
    if batch is None:
        return None  # triggers MisconfigurationException on XLA
    loss = self(batch).loss
    self.log("loss", loss)

# after
def training_step(self, batch, batch_idx):
    loss = self(batch).loss  # always compute and return loss
    self.log("loss", loss)
    return loss

out = model.training_step(batch, batch_idx)
assert out is not None, "training_step must return a loss on XLA"

When it happens

Trigger: Your LightningModule.training_step returns None (or implicitly returns None) while automatic_optimization is True and you use the XLA precision plugin / XLA strategy; optimizer_step then raises MisconfigurationException.

Common situations: Porting code that conditionally skips batches (e.g. return None on empty batch) to TPU; refactoring training_step and accidentally dropping the loss return.

Related errors

• Skipping backward by returning `None` from your `training_st
• Device should be CPU, got {device} instead.
• str(_XLA_AVAILABLE)
• `Trainer.save_checkpoint(..., storage_options=...)` with `st
• str(_XLA_AVAILABLE)

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