Repository: hpcaitech/ColossalAI
File: colossalai/checkpoint_io/utils.py (Python)

Code before:
            
        if ret_block != None:
            yield ret_block, ret_block_size

    yield current_block, current_block_size
    

def build_index(state_dict_shard: Iterator[Tuple[OrderedDict, int]], shards_total_num: int, use_safetensors: bool, variant: str):
    # If we only have one shard, we return it
    weights_name = SAFE_WEIGHTS_NAME if use_safetensors else WEIGHTS_NAME
    weights_name = add_variant(weights_name, variant)

    save_index_file = SAFE_WEIGHTS_INDEX_NAME if use_safetensors else WEIGHTS_INDEX_NAME
    save_index_file = add_variant(save_index_file, variant)

    if shards_total_num == 1:
        return {weights_name: next(state_dict_shard)[0]}, None
    
    weight_map = {}
    shards = {}
    total_size = 0
    # shard_pair like (shard, shard_size)
    for idx, shard_pair in enumerate(state_dict_shard):
        shard_file = weights_name.replace(".bin", f"-{idx+1:05d}-of-{shards_total_num:05d}.bin")
        shard_file = shard_file.replace(
            ".safetensors", f"-{idx + 1:05d}-of-{shards_total_num:05d}.safetensors"
        )
        shards[shard_file] = shard_pair[0]
        for key in shard_pair[0].keys():
            weight_map[key] = shard_file
        total_size = total_size + shard_pair[1]

    # Add the metadata
    metadata = {"total_size": total_size}
    index = {"metadata": metadata, "weight_map": weight_map}
    shards_index = {save_index_file: index}
    return shards, shards_index


def load_shard_state_dict(checkpoint_file: Path, use_safetensors: bool =False):
    """
    load shard state dict into model
    """
    if use_safetensors and not checkpoint_file.suffix == ".safetensors":
        raise Exception("load the model using `safetensors`, but no file endwith .safetensors")
    if use_safetensors:
        from safetensors.torch import safe_open
        from safetensors.torch import load_file as safe_load_file
        with safe_open(checkpoint_file, framework="pt") as f:
            metadata = f.metadata()
        if metadata["format"] != "pt":

Review comment (suggestion):
@FrankLeeeee Can the file be like `pytorch_model-00001.bin`? It wastes much time and memory to collect checkpoint twice for gemini.

Code after:

    yield current_block, current_block_size


def load_shard_state_dict(checkpoint_file: Path, use_safetensors: bool =False):
    """
    load shard state dict into model
    """
    if use_safetensors and not checkpoint_file.suffix == ".safetensors":
        raise Exception("load the model using `safetensors`, but no file endwith .safetensors")
    if use_safetensors:
        from safetensors.torch import safe_open
        from safetensors.torch import load_file as safe_load_file
        with safe_open(checkpoint_file, framework="pt") as f:
            metadata = f.metadata()
        if metadata["format"] != "pt":
            raise NotImplementedError(
                f"Conversion from a {metadata['format']} safetensors archive to PyTorch is not implemented yet."
            )
        return safe_load_file(checkpoint_file)
    else:
        return torch.load(checkpoint_file)
    
def load_state_dict_into_model(model: nn.Module, state_dict: torch.Tensor, missing_keys: List, strict: bool = False, load_sub_module: bool = True):
    r"""Copies parameters and buffers from :attr:`state_dict` into
    this module and its descendants. 

    Args:
        state_dict (dict): a dict containing parameters and
            persistent buffers.
    """
    if not isinstance(state_dict, Mapping):
        raise TypeError("Expected state_dict to be dict-like, got {}.".format(type(state_dict)))

    unexpected_keys: List[str] = []
    sub_missing_keys: List[str] = []
    error_msgs: List[str] = []

    # copy state_dict so _load_from_state_dict can modify it
    metadata = getattr(state_dict, '_metadata', None)
    state_dict = OrderedDict(state_dict)
    if metadata is not None:
        state_dict._metadata = metadata

    def load(module: nn.Module, state_dict, prefix="", load_sub_module: bool = True):
        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
        args = (state_dict, prefix, local_metadata, True, sub_missing_keys, [], error_msgs)
        # Parameters of module and children will start with prefix. We can exit early if there are none in this
        # state_dict
        if len([key for key in state_dict if key.startswith(prefix)]) > 0:
            module._load_from_state_dict(*args)