{"record":{"id":"bda0d1491bb067a5","repo":"Lightning-AI/pytorch-lightning","slug":"found-multiple-distributed-models-in-the-given-sta-bda0d1","errorCode":null,"errorMessage":"Found multiple distributed models in the given state. Loading distributed checkpoints 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 distributed models in the given state\\. Loading distributed checkpoints 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/model_parallel.py","lineNumber":438,"sourceCode":"    weights_only: Optional[bool] = None,\n) -> dict[str, Any]:\n    from torch.distributed.checkpoint.state_dict import (\n        StateDictOptions,\n        get_model_state_dict,\n        get_optimizer_state_dict,\n        set_optimizer_state_dict,\n    )\n\n    modules = {key: module for key, module in state.items() if _has_dtensor_modules(module)}\n    if len(modules) == 0:\n        raise ValueError(\n            \"Could not find a distributed 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 distributed models in the given state. Loading distributed checkpoints 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_options = StateDictOptions(cpu_offload=True)\n\n        module_state = {module_key: get_model_state_dict(module)}\n        _distributed_checkpoint_load(module_state, path)\n        module.load_state_dict(module_state[module_key], strict=strict)\n\n        # the optimizer states must be loaded separately\n        for optim_key, optim in optimizers.items():\n            optim_state = {optim_key: get_optimizer_state_dict(module, optim)}\n            _distributed_checkpoint_load(optim_state, path)\n            set_optimizer_state_dict(module, optim, optim_state_dict=optim_state[optim_key], options=state_dict_options)","sourceCodeStart":420,"sourceCodeEnd":456,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/fabric/strategies/model_parallel.py#L420-L456","documentation":"ModelParallelStrategy's distributed checkpoint loading supports at most one distributed model per checkpoint. When more than one entry in the `state` dict contains modules with DTensor parameters, this ValueError is raised because the sharded checkpoint format ties shards to a single model's state.","triggerScenarios":"load_checkpoint(path, state={'model1': m1, 'model2': m2}) where both m1 and m2 have DTensor parameters (both set up through the parallel strategy).","commonSituations":"Ensembling / multi-model pipelines under ModelParallel or TensorParallel plugins; refactoring code that previously loaded several models in one call under a different strategy.","solutions":["Split into separate calls, one per model, each with its own checkpoint path: load_checkpoint(path1, state={'model1': m1}); load_checkpoint(path2, state={'model2': m2})","Only include one distributed model in the state dict and load other (non-distributed) objects separately"],"exampleFix":"# before\nstrategy.load_checkpoint(path, state={'model1': m1, 'model2': m2})\n# after\nstrategy.load_checkpoint(path1, state={'model1': m1})\nstrategy.load_checkpoint(path2, state={'model2': m2})","handlingStrategy":"validation","validationCode":"distributed = [k for k, v in state.items() if is_distributed_module(v)]\nassert len(distributed) <= 1, f'multiple distributed models: {distributed}; load one per checkpoint'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep one distributed model per checkpoint path","Load additional models with separate load_checkpoint calls"],"tags":["lightning","fabric","model-parallel","checkpoint","multiple-models"],"backgroundTag":"checkpoint-state-mismatch","analyzedSha":"9fed5c27d2a62ff0efd6c3573599921d6ff67c14","analyzedAt":"2026-08-28T11:52:41.083Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}