Lightning-AI/pytorch-lightning · error · ValueError

Unknown state_dict_type: {self._state_dict_type}

Error message

Unknown state_dict_type: {self._state_dict_type}

What it means

Defensive fallback at the end of XLAFSDPStrategy.load_checkpoint: the strategy's _state_dict_type attribute did not match 'sharded' or 'full'. Since the constructor validates allowed values, this normally only fires if the attribute was mutated after construction or an internal bug/new state_dict_type was introduced.

Source

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

                    " 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
            kwargs.setdefault("compute_dtype", precision._desired_dtype)
        kwargs = _auto_wrap_policy_kwargs(self._auto_wrap_policy, kwargs)
        return _activation_checkpointing_kwargs(self._activation_checkpointing_policy, kwargs)

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Only use supported state_dict_type values: 'sharded' or 'full'
2. Do not mutate _state_dict_type after construction; pass state_dict_type=... to XLAFSDPStrategy(...) at init
3. If subclassing, ensure you keep the attribute within supported values or override load_checkpoint coherently

Example fix

# before
strategy._state_dict_type = 'sharded2'

# after
strategy = XLAFSDPStrategy(state_dict_type='sharded')

assert strategy._state_dict_type in ('sharded', 'full')

When it happens

Trigger: Assigning strategy._state_dict_type = 'something' directly, subclassing XLAFSDPStrategy and overriding the attribute, or a version mismatch where the code and config expect different state_dict_type values.

Common situations: Monkey-patching or subclass customizations; carrying pickled strategy objects between versions with different supported state_dict_types.

Understand the failure class

Background: "Invalid configuration value" and "Unsupported/Unknown setting value" errors: why libraries reject your config strings, numbers, and types — this error's family across 30 libraries.

Related errors

• Could not find a XLAFSDP model in the provided checkpoint st
• Found multiple XLAFSDP modules in the given state. Saving ch
• Multihost setups do not have a shared filesystem, so the che
• Got `XLAFSDPStrategy.load_checkpoint(..., state={state!r})`
• Loading a single module or optimizer object from a checkpoin

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