{"record":{"id":"c8580764987a0414","repo":"Lightning-AI/pytorch-lightning","slug":"found-multiple-fsdp-models-in-the-given-state-loa","errorCode":null,"errorMessage":"Found multiple FSDP models in the given state. Loading checkpoints with FSDP is currently limited to a single model per checkpoint. To load multiple models, call the load method for each model separately with a different path.","messagePattern":"Found multiple FSDP models in the given state\\. Loading checkpoints with FSDP is currently limited to a single model per checkpoint\\. To load multiple models, call the load method for each model separately with a different path\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"src/lightning/fabric/strategies/fsdp.py","lineNumber":563,"sourceCode":"\n        if isinstance(state, Optimizer):\n            raise NotImplementedError(\n                \"Loading a single optimizer object from a checkpoint is not supported yet with the FSDP strategy.\"\n            )\n\n        from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict\n        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP\n\n        modules = {key: module for key, module in state.items() if _has_fsdp_modules(module)}\n        if len(modules) == 0:\n            raise ValueError(\n                \"Could not find a FSDP model in the provided checkpoint state. Please provide the model as\"\n                \" part of the state like so: `load_checkpoint(..., state={'model': model, ...})`. Make sure\"\n                \" you set up the model (and optimizers if any) through the strategy before loading the checkpoint.\"\n            )\n        optimizers = {key: optim for key, optim in state.items() if isinstance(optim, Optimizer)}\n        if len(modules) > 1:\n            raise ValueError(\n                \"Found multiple FSDP models in the given state. Loading checkpoints with FSDP is\"\n                \" currently limited to a single model per checkpoint. To load multiple models, call the\"\n                \" load method for each model separately with a different path.\"\n            )\n        module_key, module = list(modules.items())[0]\n\n        if _is_sharded_checkpoint(path):\n            state_dict_ctx = _get_sharded_state_dict_context(module)\n\n            with state_dict_ctx:\n                module_state = {module_key: module.state_dict()}\n                _distributed_checkpoint_load(module_state, path)\n                module.load_state_dict(module_state[module_key], strict=strict)\n\n                if optimizers:\n                    # TODO: replace with newer APIs\n                    # https://github.com/pytorch/pytorch/issues/119800#issuecomment-1942156271\n                    reader = _get_distributed_checkpoint_reader(path)","sourceCodeStart":545,"sourceCodeEnd":581,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/fabric/strategies/fsdp.py#L545-L581","documentation":"Symmetric to the save-side restriction: FSDP checkpoint loading supports exactly one FSDP model per load call because the sharded checkpoint layout identifies one module. A state containing multiple FSDP modules is ambiguous and rejected.","triggerScenarios":"Calling load_checkpoint(path, {'model': model_a, 'vae': model_b}) where both entries contain FSDP-wrapped modules.","commonSituations":"Restoring a multi-component generative model (UNet + VAE + text encoder) from what used to be a single combined checkpoint; resuming knowledge-distillation setups with student and teacher both sharded.","solutions":["Load each model separately: load_checkpoint(path1, {'model': model_a}) then load_checkpoint(path2, {'model': model_b})","Restructure training so only one model is FSDP-wrapped and the others use a different strategy","If checkpoints must share a path, use distinct subdirectories per model"],"exampleFix":"# before\nstrategy.load_checkpoint('ckpt', state={'unet': unet, 'vae': vae})\n# after\nstrategy.load_checkpoint('ckpt/unet', state={'unet': unet})\nstrategy.load_checkpoint('ckpt/vae', state={'vae': vae})","handlingStrategy":"validation","validationCode":"from torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nmods = [k for k, v in state.items() if hasattr(v, 'modules') and any(isinstance(m, FSDP) for m in v.modules())]\nassert len(mods) <= 1, f'multiple FSDP models: {mods}'","typeGuard":"def single_fsdp_keys(state: dict) -> list[str]:\n    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP\n    return [k for k, v in state.items() if hasattr(v, 'modules') and any(isinstance(m, FSDP) for m in v.modules())]","tryCatchPattern":null,"preventionTips":["One load call per FSDP model","Store each model's checkpoint under its own path/subdirectory"],"tags":["fsdp","checkpoint","load","distributed","pytorch-lightning"],"backgroundTag":"fsdp-multi-model-checkpoint","analyzedSha":"9fed5c27d2a62ff0efd6c3573599921d6ff67c14","analyzedAt":"2026-08-28T11:52:41.083Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}