Function load_diff_checkpoint

accessory/util/tensor_parallel.py:387–422

r"""This method loads model from a *_diff format checkpoint. The behavior of loading a diff checkpoint is different from loading a regular checkpoint: In case a key is in the given ``existing_keys``, the new value of the tensor is the value in the state_dict plus the old value. Note

(
    model: nn.Module, state_dict: Dict[str, torch.Tensor],
    existing_keys: Set[str],
)

def load_diff_checkpoint(
    model: nn.Module, state_dict: Dict[str, torch.Tensor],
    existing_keys: Set[str],
) -> Tuple[List[str], List[str]]:
    r"""This method loads model from a *_diff format checkpoint. The behavior
    of loading a diff checkpoint is different from loading a regular
    checkpoint: In case a key is in the given ``existing_keys``, the new value
    of the tensor is the value in the state_dict plus the old value.

    Note:
        The input ``state_dict`` will be changed in-place to save memory.

    Args:
        model (nn.Module): The model to load the state dict into.
        state_dict (Dict[str, torch.Tensor]): The state dict to be loaded into
            the model.
        existing_keys (Set[str]): A set of keys that have appeared in the
            previous checkpoints. If a key is in this set, the corresponding
            value from the state dict will be added to the value in the model;
            otherwise the value in the model is considered uninitialized and
            is directly set to the value in the state dict.

    Returns:
        Tuple[List[str], List[str]]: A pair of lists including missing keys and
            unexpected keys, following the regular
            ``torch.nn.Module.load_stat_dict``.
    """
    model_state_dict = model.state_dict()
    for key in list(state_dict.keys()):
        if key in existing_keys and key in model_state_dict:
            orig_value = model_state_dict[key]
            diff_value = state_dict[key]
            orig_value = orig_value.to(diff_value.device)
            diff_value = diff_value.to(orig_value.dtype)
            state_dict[key] = orig_value + diff_value
    return model.load_state_dict(state_dict, strict=False)


def load_tensor_parallel_model_list(