[checkpointio]support distributed checkpoint io for model saving.

colossalai/booster/plugin/hybrid_parallel_plugin.py

@@ -78,6 +78,9 @@ def __init__(
         self.require_grad_sync = True
         self.overlap_allgather = overlap_allgather
         self.use_fp8 = use_fp8
+        self.param_origin_shape = {}
+        for name, param in module.named_parameters():
+            self.param_origin_shape[name] = param.shape
 
         shardformer = ShardFormer(shard_config)
         if custom_policy is not None:

colossalai/checkpoint_io/distributed_checkpoint_utils.py

@@ -0,0 +1,238 @@
+import json
+import os
+from typing import Dict
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from torch.distributed.distributed_c10d import _get_default_group
+
+from colossalai.interface import ModelWrapper
+from colossalai.shardformer.layer.parallel_module import ParallelModule
+from contextlib import contextmanager
+
+from .utils import (
+    load_state_dict,
+    search_tp_partition_dim,
+)
+
+MODEL_META_PREFIX = "pytorch_model-meta-dist-"
+MODEL_WEIGHT_PREFIX = "pytorch_model-dist-"
+SHARD_META_SUFFIX = ".index.json"
+UNSHARD_META_SUFFIX = ".json"
+
+
+@contextmanager
+def RestoreDefaultStateDictBehavior(model):
+    original_methods = {}
+    for name, module in model.named_modules():
+        if isinstance(module, ParallelModule):
+            original_methods[module] = (module._save_to_state_dict, module._load_from_state_dict)
+            module._save_to_state_dict = nn.Module._save_to_state_dict.__get__(module, nn.Module)
+            module._load_from_state_dict = nn.Module._load_from_state_dict.__get__(module, nn.Module)
+    try:
+        yield model
+    finally:
+        for module, original_method in original_methods.items():
+            module._save_to_state_dict, module._load_from_state_dict  = original_method
+
+
+def save_metadata(model_metadata, metadata_file, checkpoint_file=None, total_size=None):
+    metadata_dicts = {
+        "checkpoint_version": "1.0",
+        "total_size": total_size,
+        "metadata": {},
+    }
+    for name, data in model_metadata.items():
+        metadata_dicts["metadata"][name] = {}
+        for k, v in data.items():
+            if isinstance(v, torch.Tensor):
+                v = v.tolist()
+            metadata_dicts["metadata"][name][k] = v
+        if checkpoint_file is not None:
+            metadata_dicts["metadata"][name]["file"] = checkpoint_file
+        metadata_dicts["metadata"][name]["rank"] = dist.get_rank(_get_default_group())
+    with open(metadata_file, "w") as json_file:
+        json.dump(metadata_dicts, json_file, indent=4)
+
+
+def load_metadata(checkpoint: str):
+    metadata_dict = {}
+    for filename in os.listdir(checkpoint):
+        if filename.startswith(MODEL_META_PREFIX) and filename.endswith(".json"):
+            file_path = os.path.join(checkpoint, filename)
+            try:
+                with open(file_path, "r") as f:
+                    metadata_json = json.load(f)
+                    for name, item in metadata_json["metadata"].items():
+                        if name not in metadata_dict:
+                            metadata_dict[name] = {}
+                            metadata_dict[name]["global_shape"] = item["global_shape"]
+                            metadata_dict[name]["shards"] = {}
+                        else:
+                            assert metadata_dict[name]["global_shape"] == item["global_shape"]
+                        shard = {item["rank"]: {}}
+                        for k, v in item.items():
+                            if k == "rank":
+                                continue
+                            shard[item["rank"]][k] = v
+                        metadata_dict[name]["shards"].update(shard)
+            except (json.JSONDecodeError, IOError) as e:
+                print(f"Unable to load file {file_path}: {e}")
+    return metadata_dict
+
+
+def find_covering_shards(shards, target_offsets, target_lengths):
+    """
+    Parameters:
+
+    shards: A list containing information about all shards.
+    target_offsets: A one-dimensional array representing the starting position of the target tensor in each dimension.
+    target_lengths: A one-dimensional array representing the lengths of the target tensor in each dimension.
+    Returns:
+
+    A list of all shards that cover the target range.
+    """
+    target_start = target_offsets
+    target_end = [start + length for start, length in zip(target_offsets, target_lengths)]
+
+    covering_shards = {}
+
+    global_shape = None
+    total_lengths = None
+    for rank, shard in shards.items():
+        shard_start = shard["offsets"]
+        shard_lengths = shard["lengths"]
+        if global_shape == None:
+            global_shape = shard["global_shape"]
+            total_lengths = [0] * len(global_shape)
+        shard_end = [start + length for start, length in zip(shard_start, shard_lengths)]
+
+        overlap = any(
+            not (target_end[dim] <= shard_start[dim] or target_start[dim] >= shard_end[dim])
+            for dim in range(len(target_start))
+        )
+        if overlap:
+            covering_shards.update({rank: shard})
+        for dim in range(len(shard_start)):
+            total_lengths[dim] = max(total_lengths[dim], shard_start[dim] + shard_lengths[dim])
+
+    assert total_lengths == global_shape
+    return covering_shards
+
+
+def extract_weight_from_shard_partial(shard, target_offsets, target_lengths):
+    """
+    Extract the target range of weights from shard data, supporting partial overlap.
+
+    param shard: A dictionary containing shard data, including 'offsets', 'lengths', and 'weight'.
+    param target_offsets: A 1D array indicating the starting position of the target tensor in each dimension.
+    param target_lengths: A 1D array indicating the length of the target tensor in each dimension.
+    return: The extracted sub-tensor of the target weights and its position within the target range.
+    """
+    shard_offsets = shard["offsets"]
+    shard_lengths = shard["lengths"]
+    weight = shard["weight"]
+
+    slices = []
+    target_slices = []
+
+    for dim, (t_offset, t_length, s_offset, s_length) in enumerate(
+        zip(target_offsets, target_lengths, shard_offsets, shard_lengths)
+    ):
+        intersection_start = max(t_offset, s_offset)
+        intersection_end = min(t_offset + t_length, s_offset + s_length)
+
+        if intersection_start >= intersection_end:
+            return None, None
+
+        shard_slice_start = intersection_start - s_offset
+        shard_slice_end = intersection_end - s_offset
+        slices.append(slice(shard_slice_start, shard_slice_end))
+
+        target_slice_start = intersection_start - t_offset
+        target_slice_end = intersection_end - t_offset
+        target_slices.append(slice(target_slice_start, target_slice_end))
+
+    target_weight = weight[tuple(slices)]
+    return target_weight, target_slices
+
+
+def assemble_tensor_from_shards_partial(shards, target_offsets, target_lengths, dtype):
+    target_tensor = torch.zeros(target_lengths, dtype=dtype)
+
+    for rank, shard in shards.items():
+        target_weight, target_slices = extract_weight_from_shard_partial(shard, target_offsets, target_lengths)
+
+        if target_weight is not None and target_slices is not None:
+            target_tensor[tuple(target_slices)] = target_weight
+
+    return target_tensor
+
+
+def is_pytorch_model_meta_dist_file(checkpoint_index_file):
+    if MODEL_META_PREFIX in str(checkpoint_index_file):
+        return True
+    return False
+
+
+def load_dist_model(
+    model_metadata: Dict,
+    checkpoint: str,
+):
+    """
+    Load model from a single file with the given path of checkpoint.
+
+    Args:
+        model (nn.Module): The model to be loaded.
+        checkpoint_index_file (str): Path to the checkpoint file.
+        strict (bool, optional): For name matching during loading state_dict. Defaults to False.
+                                    This argument should be manually set to False since not all params in checkpoint are needed for each device when pipeline is enabled.
+    """
+    metadata_loaded = load_metadata(checkpoint)
+
+    load_files = {}
+    covered_shards = {}
+    for key, item in model_metadata.items():
+        offsets = item["offsets"]
+        lengths = item["lengths"]
+        assert (
+            item["global_shape"] == metadata_loaded[key]["global_shape"]
+        ), f"{item['global_shape']}, {metadata_loaded[key]['global_shape']}"
+        shards = metadata_loaded[key]["shards"]
+        covering_shards = find_covering_shards(shards=shards, target_offsets=offsets, target_lengths=lengths)
+        covered_shards[key] = covering_shards
+        for rank, shard in covering_shards.items():
+            if rank not in load_files:
+                load_files[rank] = set()
+            load_files[rank].add(shard["file"])
+
+    dtype = None
+    for rank, files in load_files.items():
+        for file in files:
+            file_path = os.path.join(checkpoint, file)
+            state_dict_shard = load_state_dict(file_path)
+            for key, weight in state_dict_shard.items():
+                if key not in covered_shards or rank not in covered_shards[key]:
+                    continue
+                if dtype == None:
+                    dtype = weight.dtype
+                covered_shards[key][rank]["weight"] = weight
+    state_dict = {}
+    for key, shards in covered_shards.items():
+        state = assemble_tensor_from_shards_partial(
+            shards, model_metadata[key]["offsets"], model_metadata[key]["lengths"], dtype=dtype
+        )
+        state_dict[key] = state
+
+    return state_dict
+
+def get_dist_files_name(weights_name, dist_id):
+    weights_name = weights_name.replace(".bin", f"-dist-{dist_id:05d}-shard.bin")
+    weights_name = weights_name.replace(".safetensors", f"-dist-{dist_id:05d}-shard.safetensors")
+    return weights_name
+
+def get_dist_meta_file_name(checkpoint, dist_id, use_safetensors):
+    if use_safetensors:
+        return os.path.join(checkpoint, f"{MODEL_META_PREFIX}{dist_id:05d}{SHARD_META_SUFFIX}")
+    return os.path.join(checkpoint, f"{MODEL_META_PREFIX}{dist_id:05d}{UNSHARD_META_SUFFIX}")