Lightning-AI/pytorch-lightning · error · NotImplementedError

XLAFSDP only supports a single model instance with 'model' a

Error message

XLAFSDP only supports a single model instance with 'model' as the key.

What it means

For state_dict_type='full' loading, XLAFSDP support is limited to exactly one model stored under the key 'model' in the state dict; the loader does a targeted state['model'].load_state_dict(...). Any other key name, a missing 'model' key, or a non-nn.Module value is not implemented. Optimizers and other state entries are ignored in this mode (a warning covers that).

Source

Thrown at src/lightning/fabric/strategies/xla_fsdp.py:611

        if self._state_dict_type == "full":
            if not path.is_file():
                raise ValueError(
                    f"The path {str(path)!r} does not point to a valid full checkpoint. Make sure the path points to a"
                    " directory with a full XLAFSDP checkpoint."
                )
            if len(optimizers) > 0 or len(state.keys() - modules.keys() - optimizers.keys()) > 0:
                rank_zero_warn(
                    "Loading a full checkpoint will only load the full model."
                    " The optimizer and any additional metadata are not included."
                )
            if len(modules) > 0:
                raise ValueError(
                    "Found a XLAFSDP model in the provided checkpoint state."
                    " Please provide the model without any XLAFSDP wrapper."
                )
            if "model" not in state or not isinstance(model := state["model"], torch.nn.Module):
                raise NotImplementedError("XLAFSDP only supports a single model instance with 'model' as the key.")
            full_ckpt = torch.load(path, weights_only=weights_only)
            model.load_state_dict(full_ckpt.pop("model"), strict=strict)
            return full_ckpt

        raise ValueError(f"Unknown state_dict_type: {self._state_dict_type}")

    @classmethod
    @override
    def register_strategies(cls, strategy_registry: _StrategyRegistry) -> None:
        strategy_registry.register("xla_fsdp", cls, description=cls.__name__)

    def _parse_fsdp_kwargs(self) -> dict:
        # this needs to be delayed because `self.precision` isn't available at init
        kwargs = self._fsdp_kwargs.copy()
        precision = self.precision
        if isinstance(precision, XLAPrecision):
            # the `compute_dtype` will be passed to the `auto_wrapper_callable` automatically, so we don't need to pass
            # it when creating it

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Use exactly state={'model': model} with an unwrapped torch.nn.Module instance
2. Store/restore optimizer state separately (it is not restored by full-checkpoint loading) — e.g. save optimizer state as part of the checkpoint's extra dict and apply manually

Example fix

# before
fabric.load_checkpoint(path, state={'net': model})

# after
fabric.load_checkpoint(path, state={'model': model})

assert 'model' in state and hasattr(state['model'], 'load_state_dict'), "full load requires state={'model': nn.Module}"

import torch.nn as nn

def full_state_is_valid(state: dict) -> bool:
    return isinstance(state, dict) and isinstance(state.get('model'), nn.Module)

When it happens

Trigger: fabric.load_checkpoint(path, state={'net': model}) (wrong key), state={'model': state_dict_tensor} (not a Module), or state={'model': m, 'optimizer': opt} where value under 'model' is not an nn.Module.

Common situations: Using generic key names from other Lightning workflows; passing model.state_dict() instead of the module; expecting multi-object restore from full checkpoints.

Related errors

• Loading a single module or optimizer object from a checkpoin
• The path {str(path)!r} does not point to a valid full checkp
• Found a XLAFSDP model in the provided checkpoint state. Plea
• Got `XLAFSDPStrategy.load_checkpoint(..., state={state!r})`
• The path {str(file)!r} does not point to valid sharded check

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