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

FSDPStrategy can produce model state dicts in 'sharded' or 'full' format depending on self._state_dict_type. lightning_module_state_dict selects a context manager per type and raises ValueError for any other value, guarding against misconfigured strategy arguments.

Source

Thrown at src/lightning/pytorch/strategies/fsdp.py:528

        cls._registered_strategies.append("fsdp")

        strategy_registry.register(
            "fsdp_cpu_offload",
            cls,
            description="Fully Sharded Data Parallel (FSDP) training with Full Sharding and CPU Offloading",
            cpu_offload=True,
        )
        cls._registered_strategies.append("fsdp_cpu_offload")

    @override
    def lightning_module_state_dict(self) -> dict[str, Any]:
        assert self.model is not None
        if self._state_dict_type == "sharded":
            state_dict_ctx = _get_sharded_state_dict_context(self.model)
        elif self._state_dict_type == "full":
            state_dict_ctx = _get_full_state_dict_context(self.model, world_size=self.world_size)
        else:
            raise ValueError(f"Unknown state_dict_type: {self._state_dict_type}")
        with state_dict_ctx:
            return self.model.state_dict()

    @override
    def load_model_state_dict(self, checkpoint: Mapping[str, Any], strict: bool = True) -> None:
        # Override to do nothing, FSDP already loaded the states in `load_checkpoint()`
        pass

    @override
    def optimizer_state(self, optimizer: Optimizer) -> dict[str, Tensor]:
        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
        from torch.distributed.fsdp import OptimStateKeyType

        if isinstance(optimizer, LightningOptimizer):
            optimizer = optimizer._optimizer

        assert self.model is not None
        if self._state_dict_type == "sharded":

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Use one of the two supported literals: `FSDPStrategy(state_dict_type="sharded")` or `"full"`
  2. Check the strategy attribute before checkpointing if it comes from config: assert it in {"sharded", "full"}

Example fix

# before
strategy = FSDPStrategy(state_dict_type="full_state_dict")

# after
strategy = FSDPStrategy(state_dict_type="full")
Defensive patterns

Strategy: validation

Validate before calling

VALID = {"sharded", "full"}
assert state_dict_type in VALID, f"state_dict_type must be one of {VALID}"
strategy = FSDPStrategy(state_dict_type=state_dict_type)

Type guard

def is_valid_state_dict_type(v: str) -> bool:
    return v in ("sharded", "full")

Prevention

When it happens

Trigger: Constructing `FSDPStrategy(state_dict_type=<anything other than "sharded"/"full">)` and then saving a checkpoint (or otherwise asking for the module state dict), which invokes lightning_module_state_dict.

Common situations: Typos like `state_dict_type="full_state_dict"` or `"sharded_state_dict"` (names from raw torch.distributed.fsdp APIs); passing an enum value where the string "sharded"/"full" is expected.

Related errors


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