Repository: Lightning-AI/pytorch-lightning
File: src/lightning/fabric/strategies/model_parallel.py (Python)

Code before:
        else:  # everything not a module or optimizer is considered metadata
            converted = obj.state_dict() if isinstance(obj, _Stateful) else obj
            target_dict = metadata
        _apply_filter(key, filter or {}, converted, target_dict)

    if full_state_dict:
        if _is_sharded_checkpoint(path):
            shutil.rmtree(path)
        converted_state.update(metadata)
        if rank == 0:
            torch.save(converted_state, path)
    else:
        if path.is_file():
            path.unlink()
        path.mkdir(parents=True, exist_ok=True)
        _distributed_checkpoint_save(converted_state, path)
        if rank == 0:
            torch.save(metadata, path / _METADATA_FILENAME)


def _load_checkpoint(
    path: Path,
    state: dict[str, Union[Module, Optimizer, Any]],
    strict: bool = True,
    optimizer_states_from_list: bool = False,
    weights_only: bool = False,
) -> dict[str, Any]:
    from torch.distributed.checkpoint.state_dict import (
        StateDictOptions,
        get_model_state_dict,
        get_optimizer_state_dict,
        set_optimizer_state_dict,
    )

    modules = {key: module for key, module in state.items() if _has_dtensor_modules(module)}
    if len(modules) == 0:
        raise ValueError(
            "Could not find a distributed model in the provided checkpoint state. Please provide the model as"
            " part of the state like so: `load_checkpoint(..., state={'model': model, ...})`. Make sure"
            " you set up the model (and optimizers if any) through the strategy before loading the checkpoint."
        )
    optimizers = {key: optim for key, optim in state.items() if isinstance(optim, Optimizer)}
    if len(modules) > 1:
        raise ValueError(
            "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."
        )
    module_key, module = list(modules.items())[0]

    if _is_sharded_checkpoint(path):

Review comment (suggestion):
I thought we would keep it unset as None until the very lat point we need to pas a bool value so we wont add any assumption to the user code

Code after:
            target_dict = converted_state
        else:  # everything not a module or optimizer is considered metadata
            converted = obj.state_dict() if isinstance(obj, _Stateful) else obj
            target_dict = metadata
        _apply_filter(key, filter or {}, converted, target_dict)

    if full_state_dict:
        if _is_sharded_checkpoint(path):
            shutil.rmtree(path)
        converted_state.update(metadata)
        if rank == 0:
            torch.save(converted_state, path)
    else:
        if path.is_file():
            path.unlink()
        path.mkdir(parents=True, exist_ok=True)
        _distributed_checkpoint_save(converted_state, path)
        if rank == 0:
            torch.save(metadata, path / _METADATA_FILENAME)


def _load_checkpoint(
    path: Path,
    state: dict[str, Union[Module, Optimizer, Any]],
    strict: bool = True,
    optimizer_states_from_list: bool = False,
    weights_only: Optional[bool] = None,
) -> dict[str, Any]:
    from torch.distributed.checkpoint.state_dict import (
        StateDictOptions,
        get_model_state_dict,
        get_optimizer_state_dict,
        set_optimizer_state_dict,
    )

    modules = {key: module for key, module in state.items() if _has_dtensor_modules(module)}
    if len(modules) == 0:
        raise ValueError(
            "Could not find a distributed model in the provided checkpoint state. Please provide the model as"
            " part of the state like so: `load_checkpoint(..., state={'model': model, ...})`. Make sure"
            " you set up the model (and optimizers if any) through the strategy before loading the checkpoint."
        )
    optimizers = {key: optim for key, optim in state.items() if isinstance(optim, Optimizer)}
    if len(modules) > 1:
        raise ValueError(
            "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."
        )
    module_key, module = list(modules.items())[0]