diff --git a/.github/workflows/lint_check.yaml b/.github/workflows/lint_check.yaml index ab1a532e2..fe86bd05a 100644 --- a/.github/workflows/lint_check.yaml +++ b/.github/workflows/lint_check.yaml @@ -10,7 +10,7 @@ on: jobs: # lint check can be auto-executed by the workflow lint-check: - runs-on: ubuntu-latest + runs-on: ubuntu-20.04 steps: - uses: actions/checkout@v3 diff --git a/configs/7B_internlm2.py b/configs/7B_internlm2.py index 97758bba4..cedb881aa 100644 --- a/configs/7B_internlm2.py +++ b/configs/7B_internlm2.py @@ -38,6 +38,12 @@ async_upload=True, # async ckpt upload. (only work for boto3 ckpt) async_upload_tmp_folder="/dev/shm/internlm_tmp_ckpt/", # path for temporarily files during asynchronous upload. oss_snapshot_freq=int(CHECKPOINT_EVERY / 2), # snapshot ckpt save frequency. + # INFO: Universal ckpt is not compatible with the original ckpt. + # Default is to use async_save and not use broadcast_load + # as broadcast_load may cause loading performance degradation. + # NOTE: If using aysnc_save, there is a risk of losing the latest ckpt + # when there is a sudden training interruption. + universal_ckpt=dict(enable=False, aysnc_save=True, broadcast_load=False), ) TRAIN_FOLDER = None diff --git a/internlm/checkpoint/checkpoint_manager.py b/internlm/checkpoint/checkpoint_manager.py index 2f7f5d4ed..a935d05ec 100644 --- a/internlm/checkpoint/checkpoint_manager.py +++ b/internlm/checkpoint/checkpoint_manager.py @@ -8,6 +8,7 @@ import torch from internlm.accelerator import get_accelerator +from internlm.checkpoint.universal_checkpoint.api import universal_load, universal_save from internlm.core.context import ParallelMode from internlm.core.context import global_context as gpc from internlm.core.trainer import TrainState @@ -60,7 +61,7 @@ class CheckpointLoadContent: SCHEDULAER = "scheduler" -def try_load_internevo_ckpt(ckpt_mm, load_info, train_state: TrainState = None): +def try_load_internevo_ckpt(ckpt_mm, load_info, train_state: TrainState = None, universal_ckpt=False): """Tries to load a checkpoint from the given folder. Args: @@ -83,7 +84,19 @@ def try_load_internevo_ckpt(ckpt_mm, load_info, train_state: TrainState = None): """ load_content_str, load_ckpt_folder, load_content = process_load_info(load_info) - if load_content.need_load(CheckpointLoadContent.MODEL): + if universal_ckpt: + checkpoint_state = {} + if load_content.need_load(CheckpointLoadContent.MODEL): + checkpoint_state["model"] = ckpt_mm.model + if load_content.need_load(CheckpointLoadContent.OPIMIZER): + checkpoint_state["optimizer"] = ckpt_mm.optimizer + universal_load( + load_ckpt_folder, checkpoint_state, broadcast_checkpoint=gpc.config.ckpt.universal_ckpt.broadcast_load + ) + if gpc.is_rank_for_log(): + logger.warning("Finsh loading universal model checkpoint and optimizer checkpoint.") + + if not universal_ckpt and load_content.need_load(CheckpointLoadContent.MODEL): load_model_checkpoint(folder=load_ckpt_folder, model=ckpt_mm.model) load_content_str += f"{CheckpointLoadContent.MODEL}, " @@ -93,12 +106,12 @@ def try_load_internevo_ckpt(ckpt_mm, load_info, train_state: TrainState = None): load_context(load_ckpt_folder, train_state) # load optimizer states. - if load_content.need_load(CheckpointLoadContent.OPIMIZER): + if not universal_ckpt and load_content.need_load(CheckpointLoadContent.OPIMIZER): load_optimizer_checkpoint(load_ckpt_folder, ckpt_mm.optimizer) load_content_str += f"{CheckpointLoadContent.OPIMIZER}, " - else: - if gpc.is_rank_for_log(): - logger.warning("CheckpointManager has no 'optimizer', skip reload optim checkpoint!") + + if not load_content.need_load(CheckpointLoadContent.OPIMIZER) and gpc.is_rank_for_log(): + logger.warning("CheckpointManager has no 'optimizer', skip reload optim checkpoint!") # load lr scheduler states. if load_content.need_load(CheckpointLoadContent.SCHEDULAER): @@ -109,7 +122,7 @@ def try_load_internevo_ckpt(ckpt_mm, load_info, train_state: TrainState = None): if gpc.is_rank_for_log(): logger.warning("CheckpointManager has no 'lr_scheduler', skip reload lr_scheduler checkpoint!") - if not load_content.need_load(CheckpointLoadContent.OPIMIZER): + if not universal_ckpt and not load_content.need_load(CheckpointLoadContent.OPIMIZER): if ckpt_mm.lr_scheduler and train_state: gpc.config.only_load_lr = True load_optimizer_checkpoint(load_ckpt_folder, ckpt_mm.optimizer) @@ -440,6 +453,7 @@ def try_save_checkpoint(self, train_state, force=False): train_state=train_state, model_config=self.model_config, model_config_file=self.model_config_file, + universal_ckpt=gpc.config.ckpt.universal_ckpt.enable, ) if ( @@ -579,9 +593,15 @@ def try_resume_training(self, train_state: TrainState, current_time=""): load_path = self.load_ckpt_info["path"] load_content = self.load_ckpt_info["content"] load_type = self.load_ckpt_info["ckpt_type"] + universal_ckpt = gpc.config.ckpt.universal_ckpt.enable + kwargs = {} + + if universal_ckpt: + assert load_type == "internevo", "Only internevo ckpt support universal ckpt." + kwargs = {"universal_ckpt": universal_ckpt} load_func = CheckpointLoadMethod.get_ckpt_load_type_func(load_type) - load_content_str = load_func(self, self.load_ckpt_info, train_state) + load_content_str = load_func(self, self.load_ckpt_info, train_state, **kwargs) # If we only load model weight, we need rewrite zero optim's fp32 buffer. if ( @@ -609,6 +629,7 @@ def save_checkpoint( train_state: TrainState, model_config: Dict = None, model_config_file: str = None, + universal_ckpt=False, ): """ Save checkpoint to the given folder path. @@ -621,13 +642,20 @@ def save_checkpoint( if gpc.is_rank_for_log(): logger.info(f"Saving checkpoint to `{folder}` at batch count:{train_state.step_count}...") - timer("save-model").start() - save_model_checkpoint(folder=folder, model=model) - timer("save-model").stop() + if not universal_ckpt: + timer("save-model").start() + save_model_checkpoint(folder=folder, model=model) + timer("save-model").stop() - timer("save-optimizer").start() - save_optimizer_checkpoint(optim=optimizer, state_path=folder) - timer("save-optimizer").stop() + timer("save-optimizer").start() + save_optimizer_checkpoint(optim=optimizer, state_path=folder) + timer("save-optimizer").stop() + else: + universal_save( + path=folder, + checkpoint_state={"model": model, "optimizer": optimizer}, + async_checkpoint=gpc.config.ckpt.universal_ckpt.aysnc_save, + ) if ( hasattr(train_state, "data_state_dict") diff --git a/internlm/checkpoint/universal_checkpoint/__init__.py b/internlm/checkpoint/universal_checkpoint/__init__.py new file mode 100644 index 000000000..6d144c821 --- /dev/null +++ b/internlm/checkpoint/universal_checkpoint/__init__.py @@ -0,0 +1,6 @@ +from .api import universal_load, universal_save + +__all__ = [ + "universal_save", + "universal_load", +] diff --git a/internlm/checkpoint/universal_checkpoint/api.py b/internlm/checkpoint/universal_checkpoint/api.py new file mode 100644 index 000000000..bffa75dd4 --- /dev/null +++ b/internlm/checkpoint/universal_checkpoint/api.py @@ -0,0 +1,48 @@ +#!/usr/bin/env python +# -*- encoding: utf-8 -*- + +# adopted from https://github.com/volcengine/veScale/blob/main/vescale/checkpoint + +from .checkpointer import UniversalCheckpointer +from .common import CheckpointState + + +def universal_save(path: str, checkpoint_state: CheckpointState, async_checkpoint=False): + """ + Save a checkpoint to a given path + Args: + path: Defines the storage path for checkpoint. + checkpoint_state: A dictionary contains key-value pairs for model and optimizer. + - Model: Identified by 'model' key, value should be a model instance. + - Optimizer: Identified by 'optimizer' key, value should be an optimizer instance. + async_checkpoint: A boolean value indicating if saving checkpoint asynchronously, + i.e. after dumping tensors from GPU memory to Host memory, + the training program can continue training immediately. + Then universal_checkpoint will serialize tensors and dumping to + the persistent storage asynchronously. + Example: + >>> checkpoint_state = { "model": nn.Module, "optimizer": HybridZeroOptimizer } + >>> UniversalCheckpointer.save("/ckpt", checkpoint_state) + """ + UniversalCheckpointer.save(path, checkpoint_state, async_checkpoint=async_checkpoint) + + +def universal_load(path: str, checkpoint_state: CheckpointState, broadcast_checkpoint=False): + """ + Load a checkpoint from a given path + Args: + path: Defines the storage path for checkpoint. + checkpoint_state: A dictionary contains key-value pairs for model and optimizer. + - Model: Identified by 'model' key, value should be a model instance. + - Optimizer: Identified by 'optimizer' key, value should be an optimizer instance. + broadcast_checkpoint: A boolean value decides if load a model replica from one data parallel process group + then broadcast tensors to other data parallel process group using GPUs + to reduce the file system access + For example, when data parellel size = 2, + processes with data parallel rank = 0 load model from file system + then broadcast it to processes with data parallel rank = 1 + Example: + >>> checkpoint_state = { "model": nn.Module, "optimizer": HybridZeroOptimizer } + >>> UniversalCheckpointer.load("/ckpt", checkpoint_state) + """ + UniversalCheckpointer.load(path, checkpoint_state, broadcast_checkpoint=broadcast_checkpoint) diff --git a/internlm/checkpoint/universal_checkpoint/checkpointer.py b/internlm/checkpoint/universal_checkpoint/checkpointer.py new file mode 100644 index 000000000..16838ff61 --- /dev/null +++ b/internlm/checkpoint/universal_checkpoint/checkpointer.py @@ -0,0 +1,269 @@ +#!/usr/bin/env python +# -*- encoding: utf-8 -*- + +# adopted from https://github.com/volcengine/veScale/blob/main/vescale/checkpoint/api + +import atexit +import os +from concurrent.futures import Future, ProcessPoolExecutor +from typing import Dict, List + +import torch +import torch.distributed as dist +from torch.distributed.checkpoint.storage import WriteResult + +from internlm.core.context import ParallelMode +from internlm.core.context import global_context as gpc +from internlm.solver.optimizer import HybridZeroOptimizer +from internlm.utils.common import get_current_device +from internlm.utils.logger import get_logger +from internlm.utils.megatron_timers import megatron_timer as timer + +from .common import MODEL_STR, OPTIMIZER_STR, CheckpointState +from .load_state_dict import load_state_dict +from .planner import UniversalLoadPlanner, UniversalSavePlanner +from .save_state_dict import save_state_dict + +logger = get_logger(__file__) + +NUM_IO_WORKER = 1 +SUPPORTED_TYPES = {MODEL_STR, OPTIMIZER_STR} + + +class BaseCheckpointer: + """ + The Checkpointer class offers APIs that enable users to save and load state dictionarie. + It is designed for extension across various training frameworks. + """ + + # Async IO related members. + state_io_workers: Dict[str, ProcessPoolExecutor] = {} + state_write_futures: Dict[str, Future[List[WriteResult]]] = {} + + @classmethod + def save(cls, path: str, checkpoint_state: CheckpointState): + """ + A Method for saving checkpoint + Args: + path: Defines the storage path for checkpoint. + checkpoint_state: A dictionary contains key-value pairs for model and optimizer. + - Model: Identified by 'model' key, value should be a model instance. + - Optimizer: Identified by 'optimizer' key, value should be an optimizer instance. + + """ + raise NotImplementedError() + + @classmethod + def load(cls, path: str, checkpoint_state: CheckpointState): + """ + A Method for loading checkpoint + Args: + path: Defines the storage path for checkpoint. + checkpoint_state: A dictionary contains key-value pairs for model and optimizer. + - Model: Identified by 'model' key, value should be a model instance. + - Optimizer: Identified by 'optimizer' key, value should be an optimizer instance. + + """ + raise NotImplementedError() + + @classmethod + def _cleanup_futures(cls): + """ + Wait for all write futures to finish before exit, then do the cleanup works. + + WARNING: this method cannot be called by the users. + """ + for key in SUPPORTED_TYPES: + if key in cls.state_write_futures: + futures = cls.state_write_futures[key] + for fut in futures: + fut.result() + cls.state_write_futures[key] = [] + if cls.state_io_workers[key] is not None: + cls.state_io_workers[key].shutdown() + cls.state_io_workers[key] = None + + +class UniversalCheckpointer(BaseCheckpointer): + """ + The Checkpointer class for universal checkpoint, A PyTorch Native Auto Parallelism Framework + """ + + save_planner = UniversalSavePlanner() + load_planner = UniversalLoadPlanner() + + optim_ckpt_proces_group = None + for key in SUPPORTED_TYPES: + BaseCheckpointer.state_io_workers[key] = ProcessPoolExecutor(max_workers=NUM_IO_WORKER) + BaseCheckpointer.state_write_futures[key] = [] + + @classmethod + def save( + cls, + path: str, + checkpoint_state: CheckpointState, + async_checkpoint: bool = False, + ): + """ + async_checkpoint: A boolean value indicating if saving checkpoint asynchronously, + i.e. after dumping tensors from GPU memory to Host memory, + the training program can continue training immediately. + Then checkpoint will serialize tensors and dumping to + the persistent storage asynchronously. + """ + # Check if we support saving the components + for key in checkpoint_state.keys(): + if key not in SUPPORTED_TYPES: + raise ValueError(f"{key} is not supported by UniversalCheckpointer") + + # Preprocess saving path + if path.startswith("local:"): + path = path.split(":")[1] + assert ":" not in path, f"{path} is not valid for universal checkpoint!" + + # Start saving checkpoint + for key, value in checkpoint_state.items(): + if key == MODEL_STR: + # Get model path + model_path = os.path.join(path, MODEL_STR) + # Create a "model" folder on under root path + if gpc.get_global_rank() == 0: + os.makedirs(model_path, exist_ok=True) + dist.barrier() + # Save model. + timer("save-model").start() + _, new_write_futures = save_state_dict( + state_dict=value.state_dict(), + path=model_path, + process_group=None, + coordinator_rank=0, + no_dist=False, + planner=cls.save_planner, + async_io=async_checkpoint, + last_write_futures=cls.state_write_futures[MODEL_STR], + io_workers=cls.state_io_workers[MODEL_STR], + is_optimizer=False, + ) + # Record new write futures. + cls.state_write_futures[MODEL_STR] = new_write_futures + dist.barrier() + timer("save-model").stop() + elif key == OPTIMIZER_STR: + # adamW hybrid zero optim + assert isinstance(value, HybridZeroOptimizer), "unsupported optimizer for universal ckpt" + optimizer_state = value.state_dict() + # Create a "optimizer" folder on under root path + # to save different parts of optimizer + optimizer_path = os.path.join(path, OPTIMIZER_STR) + if gpc.get_global_rank() == 0: + os.makedirs(optimizer_path, exist_ok=True) + dist.barrier() + # Save optimizer + timer("save-optimizer").start() + _, new_write_futures = save_state_dict( + state_dict=optimizer_state["sharded_optimizer_state"], + path=optimizer_path, + process_group=None, + coordinator_rank=0, + no_dist=False, + planner=cls.save_planner, + async_io=async_checkpoint, + last_write_futures=cls.state_write_futures[OPTIMIZER_STR], + io_workers=cls.state_io_workers[OPTIMIZER_STR], + is_optimizer=True, + ) + # Record new write futures. + cls.state_write_futures[OPTIMIZER_STR] = new_write_futures + # Save the global part of optimizer state + optimizer_state.pop("sharded_optimizer_state") + if gpc.get_global_rank() == 0: + torch.save(optimizer_state, os.path.join(path, "global_optimizer_state.pt")) + dist.barrier() + timer("save-optimizer").stop() + + @classmethod + def load( + cls, + path: str, + checkpoint_state: CheckpointState, + broadcast_checkpoint: bool = False, + ): + """ + broadcast_checkpoint: A boolean value decides if load a model replica from one data parallel process group + then broadcast tensors to other data parallel process group using GPUs + to reduce the file system access + For example, when data parellel size = 2, + processes with data parallel rank = 0 load model from file system + then broadcast it to processes with data parallel rank = 1 + """ + # Check if we support loading the component. + for key in checkpoint_state.keys(): + if key not in SUPPORTED_TYPES: + raise ValueError(f"{key} is not supported by UniversalCheckpointer") + + # Preprocess loading path + if path.startswith("local:"): + path = path.split(":")[1] + assert ":" not in path, f"{path} is not valid for universal checkpoint!" + + # Start loading checkpoint + for key, value in checkpoint_state.items(): + if key == MODEL_STR: + # Get model path and state dictionary + model_path = os.path.join(path, MODEL_STR) + model_state = value.state_dict() + # Set process group + if broadcast_checkpoint: + model_load_process_group = gpc.get_group(ParallelMode.DATA) + else: + model_load_process_group = None + # Load model + load_state_dict( + state_dict=model_state, + path=model_path, + process_group=model_load_process_group, + coordinator_rank=0, + no_dist=False, + planner=cls.load_planner, + broadcast_tensors=broadcast_checkpoint, + ) + # Load back to model + value.load_state_dict(model_state) + elif key == OPTIMIZER_STR: + # Get optimizer path and state dictionary + optimizer_path = os.path.join(path, OPTIMIZER_STR) + optimizer_state = value.state_dict() + # Load optimizer state dictionary + load_state_dict( + state_dict=optimizer_state["sharded_optimizer_state"], + path=optimizer_path, + process_group=None, + coordinator_rank=0, + no_dist=False, + planner=cls.load_planner, + broadcast_tensors=False, + is_optimizer=True, + ) + # Load back to optimizer + global_optimizer_state = torch.load( + os.path.join(path, "global_optimizer_state.pt"), map_location=get_current_device() + ) + value.load_state_dict(optimizer_state["sharded_optimizer_state"], global_optimizer_state) + dist.barrier() + + @classmethod + def __cleanup(cls): + """ + Wait for all write futures to finish before exit, then do the cleanup works. + + WARNING: this method cannot be called by the users. + """ + cls.save_planner.clear_cache() + BaseCheckpointer._cleanup_futures() + + @classmethod + def _register_cleanup(cls): + atexit.register(UniversalCheckpointer.__cleanup) + + +UniversalCheckpointer._register_cleanup() diff --git a/internlm/checkpoint/universal_checkpoint/common.py b/internlm/checkpoint/universal_checkpoint/common.py new file mode 100644 index 000000000..a2e6d6bb3 --- /dev/null +++ b/internlm/checkpoint/universal_checkpoint/common.py @@ -0,0 +1,110 @@ +#!/usr/bin/env python +# -*- encoding: utf-8 -*- + +# adopted from https://github.com/volcengine/veScale/blob/main/vescale/checkpoint/planner + +import collections +import dataclasses +from collections import OrderedDict +from typing import Any, Dict, Hashable, List, Optional, Tuple, TypeVar + +from torch.distributed.checkpoint.metadata import Metadata, MetadataIndex +from torch.distributed.checkpoint.planner import SavePlan +from typing_extensions import Protocol, runtime_checkable + +from internlm.utils.logger import get_logger + +logger = get_logger(__file__) + + +MODEL_STR = "model" +OPTIMIZER_STR = "optimizer" +STATE_DICT_STR = "state_dict" +_MAX_CACHE_SIZE = 2 # model ckpt + optm ckpt + + +@runtime_checkable +class Stateful(Protocol): + def state_dict(self) -> Dict[str, Any]: + ... + + def load_state_dict(self, state_dict: Dict[str, Any], *args) -> None: + ... + + +T = TypeVar("T", bound=Stateful) +CheckpointState = Dict[str, T] + + +class PlanLRUCache: + """ + For saving cache. + """ + + def __init__(self) -> None: + self._cache: OrderedDict[Hashable, Tuple[SavePlan, Metadata]] = OrderedDict() # pylint: disable=E1136 + self._capacity = _MAX_CACHE_SIZE + + def get(self, key: Hashable) -> Optional[Tuple[SavePlan, Metadata]]: + if key in self._cache: + return self._cache[key] + else: + return None + + def put(self, key: Hashable, plan_value: SavePlan, metadata_value: Metadata) -> None: + if key in self._cache: + self._cache.move_to_end(key, last=False) + else: + self._cache[key] = (plan_value, metadata_value) + if len(self._cache) > self._capacity: + self._cache.popitem() + + def clear(self) -> None: + self._cache.clear() + self._capacity = _MAX_CACHE_SIZE + + def __repr__(self) -> str: + return f"PlanLURCache(capacity: {self._capacity}, keys: {tuple(self._cache.keys())})" + + +def custom_dedup_tensors(all_plans: List[SavePlan]) -> List[SavePlan]: + """ + A function to remove duplicate tensors to write + when creating global writing plan for saving checkpoint + During the deduplication, + we balance the workloads for duplicated tensors + """ + key_to_plan: Dict[MetadataIndex, List[int]] = {} + for plan_idx, plan in enumerate(all_plans): + for write_item in plan.items: + key_to_plan.setdefault(write_item.index, []).append(plan_idx) + + replicated_items = {k: v for k, v in key_to_plan.items() if len(v) > 1} + # Remove duplicates by always keeping the first entry (Not balance). + # Compute the per-rank remove set. + plan_to_keys: Dict[int, List[MetadataIndex]] = {} + # Record the number of non-duplicated tensors assigned to each rank + assigned_work_load = collections.defaultdict(int) + for plan_idx, plan in enumerate(all_plans): + for write_item in plan.items: + if write_item.index not in replicated_items: + assigned_work_load[plan_idx] += 1 + + for key, plans in replicated_items.items(): + # For duplicated tensors, select the rank assigned with minimum number tensors so far + writer_id = min(plans, key=lambda k: assigned_work_load[k]) + assigned_work_load[writer_id] += 1 + for plan_idx in plans: + # If the rank is not writer rank, remove the key in the rank's plan + if plan_idx != writer_id: + plan_to_keys.setdefault(plan_idx, []).append(key) + # logger.info("Duplicate keys to remove: %s", plan_to_keys) + + for plan_idx, keys in plan_to_keys.items(): + # Key Set contains keys to remove + key_set = set(keys) + # rewrite items and remove elements + new_items = [write_item for write_item in all_plans[plan_idx].items if write_item.index not in key_set] + all_plans[plan_idx] = dataclasses.replace(all_plans[plan_idx], items=new_items) + + return all_plans diff --git a/internlm/checkpoint/universal_checkpoint/filesystem.py b/internlm/checkpoint/universal_checkpoint/filesystem.py new file mode 100644 index 000000000..2db95137a --- /dev/null +++ b/internlm/checkpoint/universal_checkpoint/filesystem.py @@ -0,0 +1,690 @@ +#!/usr/bin/env python +# -*- encoding: utf-8 -*- + +# adopted from https://github.com/volcengine/veScale/blob/main/vescale/checkpoint/storage + +import collections +import dataclasses +import io +import os +import pickle +from abc import ABC, abstractmethod +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Tuple, Union, cast + +import torch +import torch.distributed as dist +from torch import Tensor +from torch._utils import _get_device_module +from torch.distributed._shard._utils import narrow_tensor_by_index +from torch.distributed.checkpoint.metadata import Metadata, MetadataIndex +from torch.distributed.checkpoint.planner import ( + LoadPlan, + LoadPlanner, + ReadItem, + SavePlan, + SavePlanner, + WriteItem, + WriteItemType, +) +from torch.distributed.checkpoint.storage import ( + StorageReader, + StorageWriter, + WriteResult, +) +from torch.distributed.checkpoint.utils import _create_file_view +from torch.futures import Future + +from internlm.core.context import ParallelMode +from internlm.core.context import global_context as gpc +from internlm.train.pipeline import map_fqn_global_to_local +from internlm.utils.common import get_current_device +from internlm.utils.logger import get_logger + +from .mem_checkpoint import ( + copy_gpu_tensor_to_cpu_pinned_mem_pool, + deallocate_cpu_tensor_in_pinned_mem_pool, +) + +logger = get_logger(__file__) + +__all__ = [ + "FileSystemWriter", + "FileSystemReader", +] + + +@dataclass +class _StorageInfo: + """ + This is the per entry storage info + """ + + relative_path: str + offset: int + length: int + + +@dataclass +class _StoragePrefix: + prefix: str + + +DEFAULT_SUFFIX = ".distcp" + + +def _result_from_write_item(item: WriteItem, size_in_bytes, storage_data) -> WriteResult: + return WriteResult(index=item.index, size_in_bytes=size_in_bytes, storage_data=storage_data) + + +class _TensorLoader(ABC): + """ + Abstract class + """ + + @abstractmethod + def add(self, fqn, size, obj): + pass + + @abstractmethod + def start_loading(self): + pass + + @abstractmethod + def values(self): + pass + + +class _SerialCpuLoader(_TensorLoader): + """ + Abstract class + currently no use + """ + + def __init__(self, resolve_fun, p2p_tensors_info=None): + self.resolve_fun = resolve_fun + self.items = [] + # For HybridZeroOptimizer, p2p_tensors_info is always none. + self.p2p_tensors_info = p2p_tensors_info + + def add(self, fqn, size, obj): + self.items.append((fqn, size, obj)) + + def start_loading(self): + pass + + def values(self): + for _, _, obj in self.items: + tensor = self.resolve_fun(obj).detach() + tensor = copy_gpu_tensor_to_cpu_pinned_mem_pool(tensor) + # Comment the original DCP code + # When dumping to pinned memory, + # the memory layout for tensor has been contiguous + # if tensor.storage().size() != tensor.numel(): + # tensor = tensor.clone() + yield ( + tensor, + obj, + ) + + +class _OverlappingCpuLoader(_TensorLoader): + """ + currently no use + """ + + def __init__( + self, + resolve_fun, + p2p_tensors_info=None, + stream=None, + inflight_threshhold=1_000_000, + ): + self.resolve_fun = resolve_fun + self.items = [] + self.inflight_threshhold = inflight_threshhold + self.in_flight_data = 0 + self.current_items: collections.deque = collections.deque() + self.idx = 0 + self.started = False + self.device_type = stream.device_type if stream else torch.device("cuda").type + self.device_module = _get_device_module(self.device_type) + # For HybridZeroOptimizer, p2p_tensors_info is always none. + self.p2p_tensors_info = p2p_tensors_info + self.stream = stream or self.device_module.current_stream() + if self.stream != self.device_module.current_stream(): + self.stream.wait_stream(self.device_module.current_stream()) + + @property + def _done(self): + return self.idx >= len(self.items) + + def _drain(self): + drained = [] + if self.in_flight_data >= self.inflight_threshhold: + self.stream.synchronize() + while self.in_flight_data >= self.inflight_threshhold: + val = self.current_items.popleft() + self.in_flight_data -= val[0].numel() * val[0].element_size() + drained.append(val) + return drained + + def _refill(self): + with self.device_module.stream(self.stream): + while not self._done and self.in_flight_data < self.inflight_threshhold: + _, _, obj = self.items[self.idx] + self.idx += 1 + tensor = self.resolve_fun(obj).detach() + if tensor.device.type == self.device_type: + tensor = copy_gpu_tensor_to_cpu_pinned_mem_pool(tensor, non_blocking=True) + + self.current_items.append( + ( + tensor, + obj, + ) + ) + self.in_flight_data += tensor.numel() * tensor.element_size() + + def _finish(self): + assert self._done + if len(self.current_items) > 0: + self.stream.synchronize() + return self.current_items + + def add(self, fqn, size, obj): + if self.started: + raise RuntimeError("cannot add items after loading started") + self.items.append((fqn, size, obj)) + + def start_loading(self): + if self.started: + return + self.started = True + self.items.sort(key=lambda x: x[1]) + self._refill() + + def values(self): + self.start_loading() + while not self._done: + drained = self._drain() + self._refill() + yield from drained + + yield from self._finish() + + +def _item_fqn(item: WriteItem) -> str: + return item.index.fqn + + +def _item_size(item: WriteItem) -> int: + size = 1 + assert item.tensor_data is not None + # can't use math.prod as PT needs to support older python + for s in item.tensor_data.size: + size *= s + + dtype = item.tensor_data.properties.dtype + return size * torch._utils._element_size(dtype) + + +def _split_by_size_and_type(bins, items: List[WriteItem]) -> List[List[WriteItem]]: + if bins == 1: + return [items] + + bytes_w = [wi for wi in items if wi.type == WriteItemType.BYTE_IO] + tensor_w = [wi for wi in items if wi.type != WriteItemType.BYTE_IO] + assert len(bytes_w) == 0, "currently no WriteItemType.BYTE_IO" + + buckets: List[List[WriteItem]] = [[] for _ in range(bins)] + bucket_sizes = [0 for _ in range(bins)] + + tensor_w.sort(key=_item_size, reverse=True) + + for i, wi in enumerate(bytes_w): + buckets[i % bins].append(wi) + + for wi in tensor_w: + idx = min(enumerate(bucket_sizes), key=lambda x: x[1])[0] + buckets[idx].append(wi) + bucket_sizes[idx] += _item_size(wi) + + return buckets + + +def _write_item(stream, data, write_item, storage_key): + offset = stream.tell() + assert isinstance(data, torch.Tensor) + assert data.device == torch.device("cpu") + torch.save(data, stream) + length = stream.tell() - offset + + return _result_from_write_item(write_item, length, _StorageInfo(storage_key, offset, length)) + + +def _write_files_from_queue( + file_name, + storage_key, + write_items, + planner: SavePlanner, + inflight_threshhold: int, + use_fsync: bool, + p2p_tensors_info=None, # currently no use +): + loader: _TensorLoader + + if torch.cuda.is_available() and inflight_threshhold > 0: + loader = _OverlappingCpuLoader( + lambda x: planner.resolve_data(x), # pylint: disable=W0108 + inflight_threshhold=inflight_threshhold, + p2p_tensors_info=p2p_tensors_info, + ) + else: + loader = _SerialCpuLoader( + lambda x: planner.resolve_data(x), p2p_tensors_info=p2p_tensors_info # pylint: disable=W0108 + ) + + tensor_w = list(write_items) + for write_item in tensor_w: + loader.add(_item_fqn(write_item), _item_size(write_item), write_item) + loader.start_loading() + + write_results = [] + stream = open(file_name, "wb") # pylint: disable=R1732 + + for tensor, write_item in loader.values(): + assert tensor.is_cpu + write_results.append(_write_item(stream, tensor, write_item, storage_key)) + # WARNING: Call deallocate_cpu_tensor_in_pinned_mem_pooltensor + # when the reference to CPU tensor goes to zero + # so the memory pool will reuse the memory if possbile + # Othterwise, the memory pool will allocate memory on the used memory range, + # leading to cuda error 712 cudaErrorHostMemoryAlreadyRegistered + deallocate_cpu_tensor_in_pinned_mem_pool(tensor) + + if use_fsync: + os.fsync(stream.fileno()) + + stream.close() + return write_results + + +def _serialize_tensor(tensor: torch.Tensor) -> bytes: + bio = io.BytesIO() + # NOTE: currently use torch.save() to do the serialization. + torch.save(tensor, bio) + return bio.getbuffer() + + +def _write_to_file(stream, content: bytes, write_item: WriteItem, storage_key: str) -> WriteResult: + offset = stream.tell() + stream.write(content) + length = stream.tell() - offset + return _result_from_write_item(write_item, length, _StorageInfo(storage_key, offset, length)) + + +def _write_files_per_proc_pipe( + file_path: Path, + storage_key: str, + byte_data_item: List[Tuple[io.BytesIO, WriteItem]], + tensor_data_item: List[Tuple[torch.Tensor, WriteItem]], + use_fsync: bool, +) -> List[WriteResult]: + write_futures = [] + write_results = [] + stream = open(file_path, "wb") # pylint: disable=R1732 + executor = ThreadPoolExecutor(max_workers=1) + # For byte data, directly write byte data. + assert len(byte_data_item) == 0 + for write_data, write_item in byte_data_item: + content = write_data.getbuffer() + write_futures.append( + executor.submit( + _write_to_file, + stream, + content, + write_item, + storage_key, + ) + ) + + # For tensor data, perform serialization in process then do saving in threadpool. + for write_data, write_item in tensor_data_item: + content = _serialize_tensor(write_data) + write_futures.append( + executor.submit( + _write_to_file, + stream, + content, + write_item, + storage_key, + ) + ) + + for fut in write_futures: + write_results.append(fut.result()) + if use_fsync: + os.fsync(stream.fileno()) + executor.shutdown(wait=False) + return write_results + + +class FileSystemWriter(StorageWriter): + """ + Basic implementation of StorageWriter using file IO. + + This implementation makes the following assumptions and simplifications: + + * The checkpoint path is an empty or non-existing directory. + * File creation is atomic + + The checkpoint consist of one file per write request plus + a `.metadata` file with the serialized metadata. + + """ + + def __init__( + self, + path: Union[str, os.PathLike], + single_file_per_rank: bool = True, + sync_files: bool = True, + worker_count: int = 1, + per_process_copy_ahead: int = 10_000_000, + ) -> None: + """ + Initialize the writer pointing to `path` + + Args: + path: directory where the checkpoint will be written to. + single_file_per_rank: Produce one file per rank instead of one file per tensor/blob. Default to True. + sync_files : force files to be synced to permanent storage. Default to True. + worker_count: Number of IO workers (processes) to use to write. Default to 1. + per_process_copy_ahead: How many bytes to copy from the GPU ahead of saving then. Default 10Mb. + + N. B. If sync_files is disabled, there's no guarantee that the checkpoint will be + consistent in the case of a failure. + """ + super().__init__() + self.path = Path(path) + self.single_file_per_rank = single_file_per_rank + self.sync_files = sync_files + self.worker_count = worker_count + self.per_process_copy_ahead = per_process_copy_ahead + + def set_up_storage_writer(self, is_coordinator: bool) -> None: + pass + + def prepare_local_plan(self, plan: SavePlan, p2p_tensors_info=None) -> SavePlan: + self.path.mkdir(parents=True, exist_ok=True) + # For HybridZeroOptimizer, p2p_tensors_info is always none. + self.p2p_tensors_info = p2p_tensors_info + return plan + + def prepare_global_plan(self, global_plan: List[SavePlan]) -> List[SavePlan]: # pylint: disable=W0237 + new_plans = [ + dataclasses.replace(plan, storage_data=_StoragePrefix(f"__{i}_")) for i, plan in enumerate(global_plan) + ] + return new_plans + + def prepare_write_data(self, tasks: List[Tuple[Path, str, List[WriteItem]]], planner: SavePlanner, is_optimizer): + """ + First stage of saving, Perform Copy data to CPU (D2H). + + Args: + tasks: partitoned tasks for workers to conduct serialization and the actual saving. + planner: save planner used to resolve the bytes and tensor data. + + NOTE: Currently we do D2H synchronously. + """ + + byte_data_item_writes: List[List[Tuple[io.BytesIO, WriteItem]]] = [] + tensor_data_item_writes: List[List[Tuple[torch.Tensor, WriteItem]]] = [] + file_path_names: List[Tuple[Path, str]] = [] + + # Perform D2H in copy stream. + for task in tasks: + file_path, file_name, write_items = task + byte_w = [wi for wi in write_items if wi.type == WriteItemType.BYTE_IO] + tensor_w = [wi for wi in write_items if wi.type != WriteItemType.BYTE_IO] + byte_data_item = [(planner.resolve_data(wi), wi) for wi in byte_w] + assert len(byte_data_item) == 0, "currentlu no WriteItemType.BYTE_IO" + tensor_data_item = [] + + item_list = [] + # Map global fqn to local fqn to search local model state data. + for item in tensor_w: + fqn = _item_fqn(item) + if not is_optimizer: + if fqn in map_fqn_global_to_local: # pylint: disable=R1715 + fqn = map_fqn_global_to_local[fqn] + + # Use tensor.clone() when saving slices to avoid unexpected memory issues. + tensor = planner.resolve_data(item, fqn).detach().clone() + tensor = copy_gpu_tensor_to_cpu_pinned_mem_pool(tensor, non_blocking=True) + tensor_data_item.append((tensor, item)) + item_list.append(item.index.fqn) + byte_data_item_writes.append(byte_data_item) + tensor_data_item_writes.append(tensor_data_item) + file_path_names.append((file_path, file_name)) + + # Deallocate pinned memory. + # NOTE: when prepare_write_data() is called next time, make sure the previous save event is completed. + # Otherwise, tensors in pinned memory pool may be overwritten. + for tensor_data_item in tensor_data_item_writes: + for tensor, _ in tensor_data_item: + assert tensor.is_cpu + deallocate_cpu_tensor_in_pinned_mem_pool(tensor) + + return byte_data_item_writes, tensor_data_item_writes, file_path_names + + def write_data( + self, plan: SavePlan, planner: SavePlanner, async_io: bool = False, io_workers=False, is_optimizer=False + ) -> Future[List[WriteResult]]: + storage_plan: _StoragePrefix = plan.storage_data + file_count = 0 + + def gen_file(): + nonlocal file_count + file_name = f"{storage_plan.prefix}{file_count}{DEFAULT_SUFFIX}" + file_count += 1 + return file_name + + tasks: List[Tuple[Path, str, List[WriteItem]]] = [] + # Generate K tasks where K is the number of worker_count. + if self.single_file_per_rank: + for bucket in _split_by_size_and_type(self.worker_count, plan.items): + file_name = gen_file() + tasks.append((self.path / file_name, file_name, bucket)) + # Generate K tasks where K is the number of write items. + else: + for item in plan.items: + file_name = gen_file() + tasks.append((self.path / file_name, file_name, [item])) + + futures = [] + if not io_workers: + executor = ProcessPoolExecutor(max_workers=self.worker_count) + else: + executor = io_workers + + # ProcessPool VERSION. + if isinstance(executor, ProcessPoolExecutor): + byte_data_item_writes, tensor_data_item_writes, file_path_names = self.prepare_write_data( + tasks, planner, is_optimizer + ) + for byte_data_item, tensor_data_item, file_path_name in zip( + byte_data_item_writes, tensor_data_item_writes, file_path_names + ): + file_path, storage_key = file_path_name + worker_args = (file_path, storage_key, byte_data_item, tensor_data_item, self.sync_files) + futures.append(executor.submit(_write_files_per_proc_pipe, *worker_args)) + if async_io: + return futures + else: + for fut in futures: + fut.result() + # fut.wait() + return futures + else: + # ThreadPool VERSION. + assert False, "unavailable version" + for task in tasks: + futures.append( + executor.submit( + _write_files_from_queue, + *task, + planner, + self.per_process_copy_ahead, + self.sync_files, + ) + ) + if async_io: + return futures + else: + for fut in futures: + fut.result() + return futures + + def finish(self, metadata: Metadata, results: List[List[WriteResult]]) -> None: + storage_md = dict() + for wr_list in results: + storage_md.update({wr.index: wr.storage_data for wr in wr_list}) + metadata.storage_data = storage_md + with (self.path / ".metadata.tmp").open("wb") as metadata_file: + pickle.dump(metadata, metadata_file) + os.fsync(metadata_file.fileno()) + + (self.path / ".metadata.tmp").rename(self.path / ".metadata") + + def reset(self, checkpoint_id): + pass + + @classmethod + def validate_checkpoint_id(cls, checkpoint_id): + pass + + +class FileSystemReader(StorageReader): + """ + Basic implementation of StorageReader using file IO. + """ + + def __init__( + self, + path: Union[str, os.PathLike], + broadcast_tensors=False, + data_parallel_process_group=None, + ) -> None: + super().__init__() + self.path = path + self.storage_data: Dict[MetadataIndex, _StorageInfo] = dict() + self.broadcast_tensors = broadcast_tensors + self.data_parallel_process_group = data_parallel_process_group + + # If broadcast_tensors is enabled, the data_parallel_process_group is not none + if self.broadcast_tensors: + assert self.data_parallel_process_group + + def _slice_file(self, file, sinfo: _StorageInfo): + return _create_file_view(file, sinfo.offset, sinfo.length) + + def _get_file_path(self, relative_path): + file_path = os.path.join(self.path, relative_path) + return file_path + + def read_data(self, plan: LoadPlan, planner: LoadPlanner) -> Future[None]: + # group requests by file + per_file: Dict[str, List[ReadItem]] = dict() + for read_item in plan.items: + item_md = self.storage_data[read_item.storage_index] + path = item_md.relative_path + per_file.setdefault(path, []).append(read_item) + + # If broadcasting model tensors is enabled, + # let processes with dp_rank=0 load models and broadcast them to other processes + if self.broadcast_tensors: + self.read_data_with_broadcast(per_file=per_file, planner=planner) + else: + # Otherwise, let all ranks load tensors from files + self.read_from_files(per_file=per_file, planner=planner) + + fut: Future = Future() + fut.set_result(None) + + return fut + + def read_from_files(self, per_file: Dict[str, List[ReadItem]], planner: LoadPlanner): + for relative_path, reqs in per_file.items(): + file_path = self._get_file_path(relative_path) + with open(file_path, "rb") as file: + reqs = sorted(reqs, key=lambda req: self.storage_data[req.storage_index].offset) + for req in reqs: + item_md = self.storage_data[req.storage_index] + file_slice = self._slice_file(file, item_md) + tensor = cast(Tensor, torch.load(file_slice, map_location="cpu")) # att + tensor = narrow_tensor_by_index(tensor, req.storage_offsets, req.lengths) + target_tensor = planner.resolve_tensor(req).detach() + + assert ( + target_tensor.size() == tensor.size() + ), f"req {req.storage_index} mismatch sizes {target_tensor.size()} vs {tensor.size()}" + + target_tensor.copy_(tensor) + planner.commit_tensor(req, target_tensor) + + def read_data_with_broadcast(self, per_file: Dict[str, List[ReadItem]], planner: LoadPlanner): + for relative_path, reqs in per_file.items(): + if gpc.get_local_rank(ParallelMode.DATA) == 0: + file_path = self._get_file_path(relative_path) + file = open(file_path, "rb") # pylint: disable=R1732 + dist.barrier(self.data_parallel_process_group) + reqs = sorted(reqs, key=lambda req: self.storage_data[req.storage_index].offset) + for req in reqs: + if gpc.get_local_rank(ParallelMode.DATA) == 0: + item_md = self.storage_data[req.storage_index] + file_slice = self._slice_file(file, item_md) + object_list = [cast(Tensor, torch.load(file_slice, map_location="cuda"))] + else: + object_list = [None] + dist.broadcast_object_list( + object_list, + src=dist.get_global_rank(self.data_parallel_process_group, 0), + group=self.data_parallel_process_group, + device=get_current_device(), + ) + tensor = object_list[0].cpu() + tensor = narrow_tensor_by_index(tensor, req.storage_offsets, req.lengths) + target_tensor = planner.resolve_tensor(req).detach() + + assert ( + target_tensor.size() == tensor.size() + ), f"req {req.storage_index} mismatch sizes {target_tensor.size()} vs {tensor.size()}" + target_tensor.copy_(tensor) + planner.commit_tensor(req, target_tensor) + + # Implementing the abstract function in StorageReader + def read_metadata(self) -> Metadata: + metadata_path = self._get_file_path(".metadata") + with open(metadata_path, "rb") as metadata_file: + metadata = pickle.load(metadata_file) + return metadata + + def set_up_storage_reader(self, metadata: Metadata, is_coordinator: bool) -> None: # pylint: disable=W0613 + self.storage_data = metadata.storage_data + assert self.storage_data is not None + + def prepare_local_plan(self, plan: LoadPlan) -> LoadPlan: + return plan + + def prepare_global_plan(self, global_plan: List[LoadPlan]) -> List[LoadPlan]: # pylint: disable=W0237 + return global_plan + + def reset(self, checkpoint_id): + pass + + @classmethod + def validate_checkpoint_id(cls, checkpoint_id): + pass diff --git a/internlm/checkpoint/universal_checkpoint/load_state_dict.py b/internlm/checkpoint/universal_checkpoint/load_state_dict.py new file mode 100644 index 000000000..65c9f413d --- /dev/null +++ b/internlm/checkpoint/universal_checkpoint/load_state_dict.py @@ -0,0 +1,85 @@ +################################################################################ +# Copyright (c) Meta Platforms, Inc. and affiliates +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +################################################################################ +# Modification Copyright 2023 ByteDance Ltd. and/or its affiliates. +################################################################################ + +from typing import Optional + +import torch.distributed as dist +from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner +from torch.distributed.checkpoint.planner import LoadPlanner +from torch.distributed.checkpoint.utils import _DistWrapper + +from internlm.utils.logger import get_logger + +from .filesystem import FileSystemReader +from .planner import UniversalLoadPlanner + +logger = get_logger(__file__) + +META_DATA_FILE = ".metadata" + + +def load_state_dict( + state_dict: dict, + path: str, + process_group: Optional[dist.ProcessGroup] = None, + coordinator_rank: int = 0, + no_dist: bool = False, + planner: Optional[LoadPlanner] = None, + broadcast_tensors=False, + is_optimizer=False, +) -> None: + """ + Loads a distributed ``state_dict`` in SPMD style. Fix sub-group storage. + """ + storage_reader = FileSystemReader( + path, + broadcast_tensors=broadcast_tensors, + data_parallel_process_group=process_group, + ) + + # Step 0: create distributed world based on process group and coordinator rank + distW = _DistWrapper(process_group, not no_dist, coordinator_rank) + if process_group: + distW.coordinator_rank = dist.get_global_rank(process_group, distW.coordinator_rank) + if planner is None: + planner = DefaultLoadPlanner() + + # Step 1: all processes create local read plan, + # then coordinator gathers all local plans and create global plan. + def local_step(): + assert planner is not None + metadata = storage_reader.read_metadata() + planner.set_up_planner(state_dict, metadata, distW.is_coordinator) + storage_reader.set_up_storage_reader(metadata, distW.is_coordinator) + + local_plan = planner.create_local_plan(is_optimizer=is_optimizer) + local_plan = storage_reader.prepare_local_plan(local_plan) + return local_plan + + def global_step(all_local_plans): + assert planner is not None + all_local_plans = planner.create_global_plan(all_local_plans) + all_local_plans = storage_reader.prepare_global_plan(all_local_plans) + return all_local_plans + + if isinstance(planner, UniversalLoadPlanner): + central_plan = distW.reduce_scatter("plan", local_step, global_step) + else: + raise AssertionError("Unsupported planner for saving checkpoint") + + # Step 2: all processes read data from the given path + def read_data(): # pylint: disable=R1711 + assert planner is not None + final_local_plan = planner.finish_plan(central_plan) + all_reads = storage_reader.read_data(final_local_plan, planner) + all_reads.wait() + return None + + _ = distW.all_gather("read", read_data) diff --git a/internlm/checkpoint/universal_checkpoint/mem_checkpoint.py b/internlm/checkpoint/universal_checkpoint/mem_checkpoint.py new file mode 100644 index 000000000..a5c4923eb --- /dev/null +++ b/internlm/checkpoint/universal_checkpoint/mem_checkpoint.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python +# -*- encoding: utf-8 -*- + +# adopted from https://github.com/volcengine/veScale/tree/main/vescale/checkpoint/utilities + +import io +import pickle +import threading +from typing import DefaultDict + +import torch + +from internlm.utils.logger import get_logger + +logger = get_logger(__file__) + +if hasattr(torch.storage, "TypedStorage"): + TypedStorage = torch.storage.TypedStorage +elif hasattr(torch.storage, "_TypedStorage"): + TypedStorage = torch.storage._TypedStorage + +# TypedStorage changes in pytorch 2. +if torch.__version__ >= "2": + + def untyped_storage(o): + return o.untyped_storage() + + def location_caster(o): + return o + +elif torch.__version__ >= "1.11": + + def untyped_storage(o): + return o.storage()._storage + + def location_caster(o): + return o._storage if isinstance(o, TypedStorage) else o + + +try: + lib = torch.cuda.cudart() +except Exception: + lib = None + + +class PinnedStoragePool: # pylint: disable=C0115 + def __init__(self): + self._l = threading.Lock() + self._m = DefaultDict(set) + + def allocate(self, nbytes: int): + with self._l: + # We don't really need storage to have the exact size. So in theory we can find a + # bigger storage that may suit here. But so far we keep everything simple here. + s = self._m[nbytes] + if not s: + t = torch.empty([nbytes], dtype=torch.uint8) + t = t.share_memory_() + if lib is not None and nbytes != 0: + err = lib.cudaHostRegister(t.data_ptr(), t.numel() * t.element_size(), 0) + assert err == 0, err + storage = untyped_storage(t) + s.add(storage) + return s.pop() + + def deallocate(self, s): + # WARNING: Call deallocate when the reference to CPU tensor goes to zero + # so the memory pool will reuse the memory if possbile + # Othterwise, the memory pool will allocate memory on the used memory range, + # leading to cuda error 712 cudaErrorHostMemoryAlreadyRegistered + with self._l: + self._m[s.nbytes()].add(s) + + +GLOBAL_POOL = PinnedStoragePool() + +TID = threading.get_ident() + + +def copy_gpu_tensor_to_cpu_pinned_mem_pool(tensor: torch.Tensor, non_blocking=False) -> torch.Tensor: + """ + Copy a tensor on GPU to pinned memory pool (host CPU memory). + The input tensor will not be modified + Args: + tensor: a tensor on cuda device + Return: + a tensor on cpu, whose data is the same as input tensor + """ + m = {} + _old_warning = getattr(torch.storage, "_warn_typed_storage_removal", None) + torch.storage._warn_typed_storage_removal = lambda *args, **kwags: None + + def persistent_id(o): + if torch.is_storage(o) or isinstance(o, TypedStorage): + storage = o + if storage._cdata in m: + return storage._cdata + if storage.device.type != "cpu": + copied = GLOBAL_POOL.allocate(storage.nbytes()) + copied.copy_(storage, non_blocking=non_blocking) + if isinstance(storage, TypedStorage): + copied = storage._new_wrapped_storage(copied) + else: + copied = storage.clone() + m[storage._cdata] = copied + return storage._cdata + return + + b = io.BytesIO() + p = pickle.Pickler(b) + p.persistent_id = persistent_id + p.dump(tensor) + b.seek(0) + up = pickle.Unpickler(b) + up.persistent_load = lambda i: m[i] + cpu_tensor = up.load() + """ + assert type(tensor) == torch.Tensor + storage_obj = tensor.storage() + cpu_storage = GLOBAL_POOL.allocate(storage_obj.nbytes()) + + cpu_storage.copy_(storage_obj, non_blocking=non_blocking) + cpu_tensor = torch.tensor(cpu_storage) + """ + torch.storage._warn_typed_storage_removal = _old_warning + return cpu_tensor + + +def deallocate_cpu_tensor_in_pinned_mem_pool(tensor: torch.Tensor): + "Deallocate CPU tensor in the global pinned memory pool" + GLOBAL_POOL.deallocate(tensor.untyped_storage()) diff --git a/internlm/checkpoint/universal_checkpoint/planner.py b/internlm/checkpoint/universal_checkpoint/planner.py new file mode 100644 index 000000000..b9dc5bece --- /dev/null +++ b/internlm/checkpoint/universal_checkpoint/planner.py @@ -0,0 +1,150 @@ +#!/usr/bin/env python +# -*- encoding: utf-8 -*- + +# adopted from https://github.com/volcengine/veScale/tree/main/vescale/checkpoint/planner/vescale + +import dataclasses +import io +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +from torch.distributed.checkpoint.default_planner import ( + DefaultLoadPlanner, + DefaultSavePlanner, +) +from torch.distributed.checkpoint.metadata import Metadata, MetadataIndex +from torch.distributed.checkpoint.planner import ( + LoadPlan, + ReadItem, + SavePlan, + WriteItem, + WriteItemType, +) + +from internlm.train.pipeline import map_fqn_local_to_global +from internlm.utils.logger import get_logger + +from .common import STATE_DICT_STR, PlanLRUCache, custom_dedup_tensors +from .planner_helpers import ( + _create_read_items, + _create_write_items, + find_state_dict_object, +) + +logger = get_logger(__file__) +__all__ = [ + "UniversalSavePlanner", + "UniversalLoadPlanner", + "create_default_local_load_plan", + "create_default_local_save_plan", +] + + +class UniversalLoadPlanner(DefaultLoadPlanner): + """ + A planner class for loading checkpoint using PyTorch DCP + """ + + def __init__(self): + super().__init__() + + def create_local_plan(self, is_optimizer=False) -> LoadPlan: + return create_default_local_load_plan(self.state_dict, self.metadata, is_optimizer) + + def resolve_tensor(self, read_item: ReadItem): + tensor = self.lookup_tensor(read_item.dest_index) + return self.transform_tensor(read_item, tensor) + + def lookup_tensor(self, index: MetadataIndex) -> torch.Tensor: + """ + This is an extension from the planner interface to make it easy to extend the default planner + """ + return find_state_dict_object(self.state_dict, index) + + +def create_default_local_load_plan(state_dict: Dict[str, Any], metadata: Metadata, is_optimizer) -> LoadPlan: + """ + A function for creating local loading plan for loading checkpoint + """ + requests = [] + for fqn, obj in state_dict.items(): + global_fqn = fqn + # For local model state_dict, the default is to use local fqn. + # We need to map it to global fqn for pipeline parallel. + # As saving ckpt is always using global fqn. + if not is_optimizer: + if fqn in map_fqn_local_to_global: # pylint: disable=R1715 + global_fqn = map_fqn_local_to_global[fqn] + + md = metadata.state_dict_metadata[global_fqn] + item = _create_read_items(fqn, global_fqn, md, obj) + requests += item + return LoadPlan(requests) + + +class UniversalSavePlanner(DefaultSavePlanner): + """ + A planner class for saving checkpoint using PyTorch DCP + """ + + def __init__(self): + super().__init__() + self._plan_cache = PlanLRUCache() + + def resolve_data(self, write_item: WriteItem, fqn=None) -> Union[torch.Tensor, io.BytesIO]: + assert write_item.type != WriteItemType.BYTE_IO + local_object = self.lookup_object(write_item.index, fqn) + return self.transform_object(write_item, local_object) + + def create_local_plan(self, is_optimizer=False) -> Tuple[SavePlan, None]: + plan, p2p_tensors_info = create_default_local_save_plan(self.state_dict, self.is_coordinator, is_optimizer) + if self.flatten_state_dict: + plan = dataclasses.replace(plan, planner_data=self.mappings) + self.plan = plan + return self.plan, p2p_tensors_info + + def lookup_object(self, index: MetadataIndex, fqn=None) -> Any: + return find_state_dict_object(self.state_dict, index, fqn) + + def lookup_plan_meta(self) -> Optional[Tuple[SavePlan, Metadata]]: + if not hasattr(self, STATE_DICT_STR): + return None + else: + plan_key = hash((frozenset(self.state_dict.keys()), self.is_coordinator)) + return self._plan_cache.get(plan_key) + + def cache_plan_meta(self, new_plan: SavePlan, new_metadata: Metadata) -> None: + plan_key = hash((frozenset(self.state_dict.keys()), self.is_coordinator)) + self._plan_cache.put(plan_key, new_plan, new_metadata) + + def clear_cache(self) -> None: + self._plan_cache.clear() + + def create_dedup_global_plan(self, all_plans: List[SavePlan]) -> Tuple[List[SavePlan], Metadata]: + # Disable DCP's dedup replicated tensors function + self.dedup_replicated_tensors = False + rst_value = super().create_global_plan(all_plans) + return rst_value + + def create_global_plan(self, all_plans: List[SavePlan]) -> Tuple[List[SavePlan], Metadata]: + # Disable DCP's dedup replicated tensors function + self.dedup_replicated_tensors = False + # Use customized deduplicate function for load balance + all_plans = custom_dedup_tensors(all_plans) + rst_value = super().create_global_plan(all_plans) + return rst_value + + +def create_default_local_save_plan( + state_dict: Dict[str, Any], is_coordinator: bool, is_optimizer=False # pylint: disable=W0613 +) -> SavePlan: + """ + A function for creating local saving plan for saving checkpoint. + """ + requests = [] + for fqn, obj in state_dict.items(): + assert isinstance(obj, (torch.Tensor)) + item = _create_write_items(fqn, obj, is_optimizer=is_optimizer) + requests += item + + return SavePlan(requests), None diff --git a/internlm/checkpoint/universal_checkpoint/planner_helpers.py b/internlm/checkpoint/universal_checkpoint/planner_helpers.py new file mode 100644 index 000000000..076d21bce --- /dev/null +++ b/internlm/checkpoint/universal_checkpoint/planner_helpers.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python +# -*- encoding: utf-8 -*- + +# adopted from https://github.com/volcengine/veScale/tree/main/vescale/checkpoint/planner/vescale + +from typing import Any, List + +import torch +from torch.distributed._shard.sharded_tensor import TensorProperties +from torch.distributed.checkpoint.metadata import ( + STORAGE_TYPES, + ChunkStorageMetadata, + MetadataIndex, + TensorStorageMetadata, +) +from torch.distributed.checkpoint.planner import ( + LoadItemType, + ReadItem, + TensorWriteData, + WriteItem, + WriteItemType, +) +from torch.distributed.checkpoint.resharding import ( + _check_shard_metadata_pair_overlap, + _shards_get_overlap_region_wrt_saved_tensor, +) + +from internlm.train.pipeline import map_fqn_local_to_global, map_layer_attr + + +def _create_write_item_for_tensor(fqn: str, tensor: torch.Tensor, is_optimizer=False) -> WriteItem: + offsets = torch.Size([0] * len(tensor.size())) + size = tensor.size() + + # We always save ckpt using global fqn. + # For optim state_dict, it originally uses global fqn. + if not is_optimizer: + if fqn in map_fqn_local_to_global: # pylint: disable=R1715 + fqn = map_fqn_local_to_global[fqn] + + map_fqn = fqn + if map_fqn not in map_layer_attr: + # Deal with exp_avg and exp_avg_sq in base optim + map_fqn = fqn.rsplit(".", 1)[0] + assert map_fqn in map_layer_attr + offsets = torch.Size(map_layer_attr[map_fqn]["offset"]) + size = torch.Size(map_layer_attr[map_fqn]["complete_size"]) + + return WriteItem( + index=MetadataIndex(fqn, offsets), + type=WriteItemType.SHARD, + tensor_data=TensorWriteData( + chunk=ChunkStorageMetadata(offsets=offsets, sizes=tensor.size()), + properties=TensorProperties.create_from_tensor(tensor), + size=size, + ), + ) + + +def _create_write_items(fqn: str, object: Any, is_optimizer=False) -> List[WriteItem]: # pylint: disable=W0622 + assert isinstance(object, torch.Tensor) + return [_create_write_item_for_tensor(fqn, object, is_optimizer=is_optimizer)] + + +def _create_read_item_for_tensor(dest_index, dest_offsets, storage_index, storage_offsets, lengths): + return ReadItem( + type=LoadItemType.TENSOR, + dest_index=dest_index, + dest_offsets=torch.Size(dest_offsets), + storage_index=storage_index, + storage_offsets=torch.Size(storage_offsets), + lengths=torch.Size(lengths), + ) + + +def create_read_items_for_chunk_list( + fqn: str, + global_fqn: str, + checkpoint_md: TensorStorageMetadata, + local_chunks: List[ChunkStorageMetadata], +) -> List[ReadItem]: + """ + Creates a list of ``ReadItem`` based on the checkpoint and local chunks. + + This applies the resharding algorithm and computes the reads needed + to satisfy ``local_chunks`` with a checkpoint described by ``checkpoint_md``. + + Args: + fqn (str): The local state_dict FQN to pass to ``ReadItem``. + global_fqn (str): The global FQN in the checkpoint. + checkpoint_md (TensorStorageMetadata): metadata for a given tensor + from a checkpoint. + local_chunks (List[ChunkStorageMetadata]): Local chunks that needs to be + loaded. + + Returns: + A list of ``ReadItem`` that will satisfy all input chunks. + """ + read_items = [] + # this is a naive quadratic algo that can be optimized later + for idx, shard in enumerate(local_chunks): + for storage_idx, storage_md in enumerate(checkpoint_md.chunks): + if not _check_shard_metadata_pair_overlap(shard, storage_md): + continue + + storage_offsets = [] + dest_offsets = [] + lengths = [] + for ( + _, + offset_for_saved_tensor, + offset_for_current_tensor, + length, + ) in _shards_get_overlap_region_wrt_saved_tensor(saved_shard=storage_md, current_shard=shard): + storage_offsets.append(offset_for_saved_tensor) + dest_offsets.append(offset_for_current_tensor) + lengths.append(length) + + read_items.append( + _create_read_item_for_tensor( + dest_index=MetadataIndex(fqn, shard.offsets, idx), + dest_offsets=dest_offsets, + storage_index=MetadataIndex(global_fqn, storage_md.offsets, storage_idx), + storage_offsets=storage_offsets, + lengths=lengths, + ) + ) + return read_items + + +def _create_chunk_from_tensor(global_fqn, tensor: torch.Tensor) -> ChunkStorageMetadata: + if global_fqn not in map_layer_attr: + # Deal with exp_avg and exp_avg_sq in base optim + global_fqn = global_fqn.rsplit(".", 1)[0] + assert global_fqn in map_layer_attr, f"{global_fqn}" + offsets = torch.Size(map_layer_attr[global_fqn]["offset"]) + + return ChunkStorageMetadata(offsets=offsets, sizes=tensor.size()) + + +def _create_read_items(fqn: str, global_fqn, md: STORAGE_TYPES, obj: Any) -> List[ReadItem]: + assert isinstance(obj, torch.Tensor) + local_chunks = [_create_chunk_from_tensor(global_fqn, obj)] + return create_read_items_for_chunk_list(fqn, global_fqn, md, local_chunks) + + +def find_state_dict_object(state_dict: dict, index: MetadataIndex, fqn=None) -> Any: + # Called when real writing happened + # The filesystem writer calls resolve_data , then it will + # call find_state_dict_object + if fqn is None: + fqn = index.fqn + if fqn not in state_dict: + raise ValueError(f"Could not find FQN: '{fqn}'") + obj = state_dict[fqn] + + return obj diff --git a/internlm/checkpoint/universal_checkpoint/save_state_dict.py b/internlm/checkpoint/universal_checkpoint/save_state_dict.py new file mode 100644 index 000000000..8552c4acf --- /dev/null +++ b/internlm/checkpoint/universal_checkpoint/save_state_dict.py @@ -0,0 +1,163 @@ +################################################################################ +# Copyright (c) Meta Platforms, Inc. and affiliates +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +################################################################################ +# Modification Copyright 2023 ByteDance Ltd. and/or its affiliates. +################################################################################ + +import os +import pickle +import time +from concurrent.futures import Future +from typing import List, Optional, Tuple + +import torch.distributed as dist +from torch.distributed.checkpoint.default_planner import DefaultSavePlanner +from torch.distributed.checkpoint.metadata import Metadata +from torch.distributed.checkpoint.planner import SavePlanner +from torch.distributed.checkpoint.storage import WriteResult +from torch.distributed.checkpoint.utils import _DistWrapper + +from internlm.core.context import global_context as gpc +from internlm.utils.logger import get_logger + +from .filesystem import FileSystemWriter +from .planner import UniversalSavePlanner + +logger = get_logger(__file__) + + +def save_state_dict( + state_dict: dict, + path: str, + process_group: Optional[dist.ProcessGroup] = None, + coordinator_rank: int = 0, + no_dist: bool = False, + planner: Optional[SavePlanner] = None, + async_io: bool = True, + last_write_futures: Future[List[WriteResult]] = None, # pylint: disable=E1136 + io_workers=None, + is_optimizer=False, +) -> Tuple[Metadata, Future[List[WriteResult]]]: # pylint: disable=E1136 + """ + Saves a model in SPMD style. Fix sub-group storage. + """ + # Step 0: create distributed world based on process group and coordinator rank + distW = _DistWrapper(process_group, not no_dist, coordinator_rank) + if process_group: + distW.coordinator_rank = dist.get_global_rank(process_group, distW.coordinator_rank) + if planner is None: + planner = DefaultSavePlanner() + assert planner is not None + + global_metatadata = None + + storage_writer = FileSystemWriter(path) + + # Step 1: all processes create local write plan, + # then coordinator gathers all local plans and create global plan. + def local_step(): + if isinstance(planner, UniversalSavePlanner): + local_plan, p2p_tensors_info = planner.create_local_plan(is_optimizer=is_optimizer) + local_plan = storage_writer.prepare_local_plan(local_plan, p2p_tensors_info) + else: + raise AssertionError("Unsupported planner for planning") + return local_plan + + def global_step(all_local_plans): + nonlocal global_metatadata + all_local_plans, global_metatadata = planner.create_global_plan(all_local_plans) + all_local_plans = storage_writer.prepare_global_plan(all_local_plans) + return all_local_plans + + # Step 2: all processes write data from GPUs to pinned memory pool, then dump to local path + # then coordinator write meta-data to local path. + def write_data(async_io: bool = False, io_workers=io_workers): + final_local_plan = planner.finish_plan(central_plan) + if isinstance(planner, UniversalSavePlanner): + # Use pinned memory pool and mult_processing for dumping ckpt to local directory efficiently + all_write_futures = storage_writer.write_data(final_local_plan, planner, async_io, io_workers, is_optimizer) + if async_io: + return all_write_futures + else: + # Gather write results. + values = [] + for fut in all_write_futures: + # values += fut.get() + values += fut.result() + return values + else: + raise AssertionError("Unsupported planner for writing data") + + def finish_checkpoint(all_results): + assert global_metatadata is not None, f"rank: {distW.get_rank()} has no global_metadata" + storage_writer.finish(metadata=global_metatadata, results=all_results) + return global_metatadata + + assert planner is not None + planner.set_up_planner(state_dict, distW.is_coordinator) + storage_writer.set_up_storage_writer(distW.is_coordinator) + + # Wait for last write futures to finish. + if last_write_futures: + if gpc.is_rank_for_log(): + logger.info("Start waiting for last write events.") + last_write_start_time = time.time() + for fut in last_write_futures: + fut.result() + last_write_time = time.time() - last_write_start_time + if gpc.is_rank_for_log(): + logger.info(f"Finish waiting for last write events. Time cost: {last_write_time}s") + + # Each worker bypass the `reduce_scatter()` and `all_reduce()` if finding cached central_plan and metadata. + # NOTE: it fails when the plans of partial workers change while others keep the same. + cached_data = None + if isinstance(planner, UniversalSavePlanner): + cached_data = planner.lookup_plan_meta() + if cached_data: + if gpc.is_rank_for_log(): + logger.info("Plan cache hit. Reuse existing plan") + central_plan, _ = cached_data + # _ = local_step() # backup for origin + else: + if gpc.is_rank_for_log(): + logger.info("Plan cache miss. The model/optimizer appears for the first time.") + central_plan = distW.reduce_scatter("plan", local_step, global_step) + else: + raise AssertionError("Unsupported planner for saving checkpoint") + + write_futures = [] + if isinstance(planner, UniversalSavePlanner): + if cached_data: + if gpc.is_rank_for_log(): + logger.info("Metdata cache hit. Reuse existing metadata") + _, final_storage_metadata = cached_data + write_results = write_data(async_io=async_io) + # Be sure to write cache metadata to .metadata file + # Otherwises only the first checkpoint has .metadata + # which leads to error when loading other checkpoints + if distW.is_coordinator: + with (storage_writer.path / ".metadata.tmp").open("wb") as metadata_file: + pickle.dump(final_storage_metadata, metadata_file) + os.fsync(metadata_file.fileno()) + + (storage_writer.path / ".metadata.tmp").rename(storage_writer.path / ".metadata") + + if async_io: + write_futures = write_results + else: + if gpc.is_rank_for_log(): + logger.info("Metadata cache miss. The model/optimizer appears for the first time.") + # First time do synchronous storing to get final_storage_metatdata. + # Determine which communication topology to use. + final_storage_metadata = distW.all_reduce("write", write_data, finish_checkpoint) + assert central_plan is not None + assert final_storage_metadata is not None + planner.cache_plan_meta(central_plan, final_storage_metadata) # attn + else: + raise AssertionError("Unsupported planner for writing data and metadata") + + return final_storage_metadata, write_futures diff --git a/internlm/initialize/launch.py b/internlm/initialize/launch.py index fc63b8a23..8578eb514 100644 --- a/internlm/initialize/launch.py +++ b/internlm/initialize/launch.py @@ -260,6 +260,9 @@ def args_sanity_check(): # to auto-load latest checkpoint. ckpt._add_item("auto_resume", True) + if "universal_ckpt" not in ckpt: + ckpt._add_item("universal_ckpt", dict(enable=False, aysnc_save=False, broadcast_load=False)) + if gpc.is_rank_for_log(): logger.info("+" * 15 + " Ckpt Info " + "+" * 15) # pylint: disable=W1201 logger.info(f"is enable save ckpt: {ckpt.enable_save_ckpt}") diff --git a/internlm/model/modules/embedding.py b/internlm/model/modules/embedding.py index 93fcd6b23..8e0632f16 100644 --- a/internlm/model/modules/embedding.py +++ b/internlm/model/modules/embedding.py @@ -47,14 +47,16 @@ def __init__( self.vocab_parallel = vocab_parallel parallel_size = gpc.weight_parallel_size if is_using_isp() else gpc.tensor_parallel_size + rank = gpc.get_local_rank(ParallelMode.WEIGHT) if is_using_isp() else gpc.get_local_rank(ParallelMode.TENSOR) if vocab_parallel: assert num_embeddings % parallel_size == 0, f"{num_embeddings} is not divisible by {parallel_size}" self.num_embeddings_per_partition = num_embeddings // parallel_size self.embed_dim_per_partition = embedding_dim - self.vocab_start_index = gpc.get_local_rank(ParallelMode.TENSOR) * self.num_embeddings_per_partition + self.vocab_start_index = rank * self.num_embeddings_per_partition self.vocab_end_index = self.vocab_start_index + self.num_embeddings_per_partition + self.offset = [self.vocab_start_index, 0] else: assert embedding_dim % parallel_size == 0, f"{embedding_dim} is not divisible by {parallel_size}" @@ -62,12 +64,15 @@ def __init__( self.embed_dim_per_partition = embedding_dim // parallel_size self.vocab_start_index = 0 self.vocab_end_index = self.num_embeddings_per_partition + self.offset = [0, self.embed_dim_per_partition * rank] self.weight = nn.Parameter( torch.empty((self.num_embeddings_per_partition, self.embed_dim_per_partition), dtype=dtype) ) - + self.complete_size = [num_embeddings, embedding_dim] setattr(self.weight, "is_embedding_param", True) + setattr(self.weight, "offset", self.offset) + setattr(self.weight, "complete_size", [num_embeddings, embedding_dim]) def forward(self, input_: Tensor) -> Tensor: if self.vocab_parallel and not is_using_isp(): diff --git a/internlm/model/modules/linear.py b/internlm/model/modules/linear.py index 29070b429..7759757bb 100644 --- a/internlm/model/modules/linear.py +++ b/internlm/model/modules/linear.py @@ -597,6 +597,7 @@ def __init__( world_size = gpc.get_world_size(parallel_mode) rank = gpc.get_local_rank(parallel_mode) + self.offset = None if split_mode != "none": split_features = out_features if split_mode == "column" else in_features @@ -609,11 +610,17 @@ def __init__( if split_mode == "column": super().__init__(in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype) + self.offset = [rank * local_multiple * multiple_of, 0] elif split_mode == "row": super().__init__(local_multiple * multiple_of, out_features, bias=bias, device=device, dtype=dtype) + self.offset = [0, rank * local_multiple * multiple_of] else: super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) + self.complete_size = [out_features, in_features] + setattr(self.weight, "offset", self.offset) + setattr(self.weight, "complete_size", [out_features, in_features]) + def forward(self, input: torch.Tensor, batch_sizes: torch.Tensor = None) -> torch.Tensor: # pylint: disable=W0622 _class_name = self.__class__.__name__ assert self._communicator is not None, f"{_class_name} should register with a communicator first." diff --git a/internlm/model/modules/mha.py b/internlm/model/modules/mha.py index a8ef77bc1..704da3072 100644 --- a/internlm/model/modules/mha.py +++ b/internlm/model/modules/mha.py @@ -65,22 +65,6 @@ def _qkv_pre_load_convert(module: "GQA", state_dict, prefix: str, *args, **kwarg ) -def _qkv_save_convert(module: "GQA", state_dict, prefix: str, *args, **kwargs) -> Dict: # pylint: disable=W0613 - wq_name, wk_name, wv_name, fused_name = ( - f"{prefix}wq.weight", - f"{prefix}wk.weight", - f"{prefix}wv.weight", - f"{prefix}wqkv.weight", - ) - - if module.enable_qkv_fusion: - state_dict[wq_name], state_dict[wk_name], state_dict[wv_name] = split_fused_wqkv_weight( - state_dict.pop(fused_name), *args, **kwargs - ) - - return state_dict - - class MHA(nn.Module): """ Multi-head self-attention and cross-attention. @@ -462,15 +446,15 @@ def __init__( if enable_qkv_fusion: assert bias is False, "Fuesd wqkv only support bias is False." self.wqkv = new_linear("wqkv", embed_dim, q_dim + 2 * self.kv_dim, bias, **factory_kwargs) - self._register_load_state_dict_pre_hook( - partial(_qkv_pre_load_convert, q_dim=q_dim, kv_dim=self.kv_dim), with_module=True - ) - self._register_state_dict_hook(partial(_qkv_save_convert, q_dim=q_dim, kv_dim=self.kv_dim)) else: self.wq = new_linear("wq", embed_dim, q_dim, bias, **factory_kwargs) self.wk = new_linear("wk", embed_dim, self.kv_dim, bias, **factory_kwargs) self.wv = new_linear("wv", embed_dim, self.kv_dim, bias, **factory_kwargs) + self._register_load_state_dict_pre_hook( + partial(_qkv_pre_load_convert, q_dim=q_dim, kv_dim=self.kv_dim), with_module=True + ) + self.inner_attn = SelfAttention( causal=causal, softmax_scale=softmax_scale, attention_dropout=dropout, layer_idx=layer_idx ) diff --git a/internlm/solver/optimizer/hybrid_zero_optim.py b/internlm/solver/optimizer/hybrid_zero_optim.py index 49f3fbcf0..dc9ae7c57 100644 --- a/internlm/solver/optimizer/hybrid_zero_optim.py +++ b/internlm/solver/optimizer/hybrid_zero_optim.py @@ -40,6 +40,7 @@ release_param_grad, split_half_float_double, sync_param, + unflatten, ) from internlm.utils.common import get_current_device from internlm.utils.logger import get_logger @@ -290,7 +291,6 @@ def _partition_param_list(self, group_id, param_group): logger.info( # pylint: disable=W1203 f"Number of elements on ranks: {numel_per_rank}, rank:{gpc.get_global_rank()}" ) - return params_per_rank, set(no_params_ranks) def _is_moe_group(self, param_group): @@ -961,59 +961,178 @@ def clip_grad_norm(self, model, max_norm): def state_dict(self): states = {} + optim_states = self.optim.state_dict() grad_scaler = self.grad_scaler.state_dict() states["grad_scaler"] = grad_scaler - optim_states = self.optim.state_dict() - states["base_optim_states"] = optim_states + if not gpc.config.ckpt.universal_ckpt.enable: + # original ckpt + states["base_optim_states"] = optim_states + flat_fp32_weights = {} + for group_id, param in self._fp32_flat_param_groups_of_current_rank.items(): + if self._zero_local_rank[group_id] not in self.param_group_no_params_ranks[group_id]: + assert param.grad is None + flat_fp32_weights[group_id] = param + states["flat_fp32_weights"] = flat_fp32_weights + states["zero_devide_optim_plan"] = self.params_per_rank_id_dict + else: + # universal ckpt. INFO: currently only adapt to AdamW + # empty_states is used as the initialization of state_dict when loading ckpt + empty_states = False + if len(optim_states["state"]) == 0: + empty_states = True + optim_states["state"] = {0: {"step": 0}} + + # To save tensor that needs to be sharded + sharded_optimizer_state = {} + for group_id, flatten_fp32_param in self._fp32_flat_param_groups_of_current_rank.items(): + rank = self._zero_local_rank[group_id] + if rank not in self.param_group_no_params_ranks[group_id]: + tensor_list = self._param_store.get_fp16_params_by_rank_group(rank, group_id) + if len(tensor_list) > 0: + unflatten_tensor_list = unflatten(flatten_fp32_param, tensor_list) + # notice: we assume that one param group corresponds to one flattened tensor. + # we will save unflattened otimizer state for universal ckpt. + if not empty_states: + flatten_exp_avg = optim_states["state"][group_id]["exp_avg"] + flatten_exp_avg_sq = optim_states["state"][group_id]["exp_avg_sq"] + assert flatten_exp_avg.shape == flatten_fp32_param.shape == flatten_exp_avg_sq.shape + unflatten_exp_avg = unflatten(flatten_exp_avg, tensor_list) + unflatten_exp_avg_sq = unflatten(flatten_exp_avg_sq, tensor_list) + assert ( + len(unflatten_tensor_list) + == len(tensor_list) + == len(unflatten_exp_avg) + == len(unflatten_exp_avg_sq) + ) + + from internlm.train.pipeline import map_fqn_local_to_global + + for i in range(len(tensor_list)): + assert tensor_list[i].fqn not in sharded_optimizer_state + fqn = tensor_list[i].fqn + # For optim ckpt, we directly save global_fqn + if fqn in map_fqn_local_to_global: # pylint: disable=consider-using-get + fqn = map_fqn_local_to_global[fqn] + + sharded_optimizer_state[fqn] = unflatten_tensor_list[i] + if not empty_states: + sharded_optimizer_state[f"{fqn}.exp_avg"] = unflatten_exp_avg[i] + sharded_optimizer_state[f"{fqn}.exp_avg_sq"] = unflatten_exp_avg_sq[i] + else: + sharded_optimizer_state[f"{fqn}.exp_avg"] = torch.empty_like(unflatten_tensor_list[i]) + sharded_optimizer_state[f"{fqn}.exp_avg_sq"] = torch.empty_like( + unflatten_tensor_list[i] + ) - flat_fp32_weights = {} - for group_id, param in self._fp32_flat_param_groups_of_current_rank.items(): - if self._zero_local_rank[group_id] not in self.param_group_no_params_ranks[group_id]: - assert param.grad is None - flat_fp32_weights[group_id] = param - states["flat_fp32_weights"] = flat_fp32_weights - states["zero_devide_optim_plan"] = self.params_per_rank_id_dict + states["step"] = optim_states["state"][0]["step"] + states["param_groups"] = optim_states["param_groups"] + states["sharded_optimizer_state"] = sharded_optimizer_state return states - def load_state_dict(self, states): - # TODO: Need to take into account the change in the number of DP. - assert "grad_scaler" in states, "Not found grad_scaler state!" - grad_scaler = states["grad_scaler"] - self.grad_scaler.load_state_dict(grad_scaler) - optim_states = states["base_optim_states"] - - if gpc.config.get("only_load_lr", False): - if gpc.is_rank_for_log(): - logger.info("Only load lr in param_groups, skip loading weights in optimizer...") - for pg1, pg2 in zip(self.optim.param_groups, optim_states["param_groups"]): - pg1["lr"] = pg2["lr"] - return - - self.optim.load_state_dict(optim_states) + def load_state_dict(self, states, global_optimizer_state=None): + if not gpc.config.ckpt.universal_ckpt.enable: + # original ckpt + # TODO: Need to take into account the change in the number of DP. + assert "grad_scaler" in states, "Not found grad_scaler state!" + grad_scaler = states["grad_scaler"] + self.grad_scaler.load_state_dict(grad_scaler) + optim_states = states["base_optim_states"] + + if gpc.config.get("only_load_lr", False): + if gpc.is_rank_for_log(): + logger.info("Only load lr in param_groups, skip loading weights in optimizer...") + for pg1, pg2 in zip(self.optim.param_groups, optim_states["param_groups"]): + pg1["lr"] = pg2["lr"] + return + + self.optim.load_state_dict(optim_states) + + # load fp32 model weight. + flat_fp32_weights = states["flat_fp32_weights"] + assert set(flat_fp32_weights.keys()) == set(self._fp32_flat_param_groups_of_current_rank) + for group_id, param in flat_fp32_weights.items(): + if self._zero_local_rank[group_id] not in self.param_group_no_params_ranks[group_id]: + self_param = self._fp32_flat_param_groups_of_current_rank[group_id] + assert ( + self_param.shape == param.shape + ), f"The loaded parameter shape is inconsistent, {self_param.shape} != {param.shape}" + self_param.data.copy_(param.data) + + # Load the fp16 model weights. + for group_id in range(len(self._fp16_param_groups)): + if self._zero_local_rank[group_id] not in self.param_group_no_params_ranks[group_id]: + fp16_param = self._param_store.get_flat_fp16_param_by_rank_group( + rank=self._zero_local_rank[group_id], group_id=group_id + ) + fp32_param = self._fp32_flat_param_groups_of_current_rank[group_id] + fp16_param.data.copy_(fp32_param) - # load fp32 model weight. - flat_fp32_weights = states["flat_fp32_weights"] - assert set(flat_fp32_weights.keys()) == set(self._fp32_flat_param_groups_of_current_rank) - for group_id, param in flat_fp32_weights.items(): - if self._zero_local_rank[group_id] not in self.param_group_no_params_ranks[group_id]: - self_param = self._fp32_flat_param_groups_of_current_rank[group_id] - assert ( - self_param.shape == param.shape - ), f"The loaded parameter shape is inconsistent, {self_param.shape} != {param.shape}" - self_param.data.copy_(param.data) + if "zero_devide_optim_plan" in states: + self.params_per_rank_id_dict = states["zero_devide_optim_plan"] + else: + # universal ckpt + assert global_optimizer_state is not None + assert "grad_scaler" in global_optimizer_state, "Not found grad_scaler state!" + assert "step" in global_optimizer_state, "Not found step state!" + assert "param_groups" in global_optimizer_state, "Not found param_groups state!" + + grad_scaler = global_optimizer_state["grad_scaler"] + self.grad_scaler.load_state_dict(grad_scaler) + + step = global_optimizer_state["step"] + param_groups = global_optimizer_state["param_groups"] + + if gpc.config.get("only_load_lr", False): + if gpc.is_rank_for_log(): + logger.info("Only load lr in param_groups, skip loading weights in optimizer...") + for pg1, pg2 in zip(self.optim.param_groups, param_groups): + pg1["lr"] = pg2["lr"] + return + + optim_states = {"state": {}, "param_groups": param_groups} + for group_id, self_flatten_fp32_param in self._fp32_flat_param_groups_of_current_rank.items(): + rank = self._zero_local_rank[group_id] + if rank not in self.param_group_no_params_ranks[group_id]: + # self fp16 unflatten param list + self_tensor_list = self._param_store.get_fp16_params_by_rank_group(rank, group_id) + + if len(self_tensor_list) > 0: + optim_states["state"][group_id] = {"step": step} + ckpt_fp32_params = [] + ckpt_exp_avg_list = [] + ckpt_exp_avg_sq_list = [] + + for tensor in self_tensor_list: + fqn = tensor.fqn + from internlm.train.pipeline import map_fqn_local_to_global + + if fqn in map_fqn_local_to_global: # pylint: disable=consider-using-get + fqn = map_fqn_local_to_global[fqn] + assert ( + tensor.shape + == states[fqn].shape + == states[f"{fqn}.exp_avg"].shape + == states[f"{fqn}.exp_avg_sq"].shape + ) + ckpt_fp32_params.append(states[fqn]) + ckpt_exp_avg_list.append(states[f"{fqn}.exp_avg"]) + ckpt_exp_avg_sq_list.append(states[f"{fqn}.exp_avg_sq"]) + + ckpt_flatten_fp32_param = flatten(ckpt_fp32_params) + ckpt_flatten_exp_avg = flatten(ckpt_exp_avg_list) + ckpt_flatten_exp_avg_sq = flatten(ckpt_exp_avg_sq_list) + + assert self_flatten_fp32_param.shape == ckpt_flatten_fp32_param.shape, ( + "The loaded parameter shape is inconsistent," + f"{self_flatten_fp32_param.shape} != {ckpt_flatten_fp32_param.shape}" + ) - # Load the fp16 model weights. - for group_id in range(len(self._fp16_param_groups)): - if self._zero_local_rank[group_id] not in self.param_group_no_params_ranks[group_id]: - fp16_param = self._param_store.get_flat_fp16_param_by_rank_group( - rank=self._zero_local_rank[group_id], group_id=group_id - ) - fp32_param = self._fp32_flat_param_groups_of_current_rank[group_id] - fp16_param.data.copy_(fp32_param) + self_flatten_fp32_param.data.copy_(ckpt_flatten_fp32_param.data) + optim_states["state"][group_id]["exp_avg"] = ckpt_flatten_exp_avg + optim_states["state"][group_id]["exp_avg_sq"] = ckpt_flatten_exp_avg_sq - if "zero_devide_optim_plan" in states: - self.params_per_rank_id_dict = states["zero_devide_optim_plan"] + self.optim.load_state_dict(optim_states) def reload_zero_fp32_buff(self): # If we use AMP optimizer, we need to update its fp32 buffer as newly loaded weights value. diff --git a/internlm/train/pipeline.py b/internlm/train/pipeline.py index 5907a4e30..f34ced17b 100644 --- a/internlm/train/pipeline.py +++ b/internlm/train/pipeline.py @@ -116,17 +116,33 @@ logger = get_logger(__file__) internlm_accelerator = get_accelerator() +# For universal checkpoint +# record offset and complete_size of param in each layer +map_layer_attr = {} +map_fqn_local_to_global = {} +map_fqn_global_to_local = {} + def set_param_unique_tracking_name(model): + uc_enable = gpc.config.ckpt.universal_ckpt.enable for chunk_id, chunk in enumerate(unwrap_naive_amp(model)): # Important: only works for llama-class models childrens = chunk.named_children() - for _, children in childrens: + for children_name, children in childrens: if isinstance(children, nn.ModuleList): for idx, block in enumerate(children): for name, child in block.named_modules(): + if name == "": + continue + + full_name = f"{chunk_id}.{idx}.{name}" + name_parts = f"{full_name}.weight".split(".", 2) + # global_id for pipeline parallel case + global_id = model.first_layer + idx + local_fqn = f"{children_name}." + ".".join(name_parts[1:]) + global_fqn = f"{children_name}.{global_id}." + ".".join(name_parts[2:]) + if isinstance(child, (ParallelLinearWithCommExt)): - full_name = f"{chunk_id}.{idx}.{name}" setattr( child.weight, "tracking_name", @@ -138,19 +154,80 @@ def set_param_unique_tracking_name(model): "tracking_name", f"{full_name}.bias", ) + + if uc_enable: + setattr( + child.weight, + "fqn", + f"{local_fqn}", + ) + if child.bias is not None: + setattr( + child.bias, + "fqn", + f"{local_fqn}", + ) + + assert hasattr(child, "offset"), f"{child}" + map_fqn_local_to_global[local_fqn] = global_fqn + map_fqn_global_to_local[global_fqn] = local_fqn + + assert global_fqn not in map_layer_attr, f"{map_layer_attr} exists" + map_layer_attr[global_fqn] = { + "offset": getattr(child, "offset", [0] * len(child.weight.size())), + "complete_size": getattr(child, "complete_size", list(child.weight.size())), + } + + elif isinstance(child, (RMSNorm)) and uc_enable: + map_fqn_local_to_global[local_fqn] = global_fqn + map_fqn_global_to_local[global_fqn] = local_fqn + setattr( + child.weight, + "fqn", + f"{local_fqn}", + ) + map_layer_attr[global_fqn] = { + "offset": getattr(child, "offset", [0] * len(child.weight.size())), + "complete_size": getattr(child, "complete_size", list(child.weight.size())), + } + else: + full_name = f"{chunk_id}.{children_name}" + local_fqn = f"{children_name}.weight" + assert getattr(children, "bias", None) is None if isinstance(children, Embedding1D): setattr( children.weight, "tracking_name", - f"{chunk_id}_embedding.weight", + f"{chunk_id}_embeddings.weight", ) + assert local_fqn not in map_layer_attr, f"{map_layer_attr} exists" else: setattr( children.weight, "tracking_name", - f"{chunk_id}_head.weight", + f"{full_name}.weight", + ) + assert local_fqn not in map_layer_attr, f"{map_layer_attr} exists" + + if uc_enable: + setattr( + children.weight, + "fqn", + f"{local_fqn}", ) + if getattr(children, "bias", None) is not None: + if children.bias is not None: + setattr( + children.bias, + "fqn", + f"{local_fqn}", + ) + + map_layer_attr[local_fqn] = { + "offset": getattr(children, "offset", [0] * len(children.weight.size())), + "complete_size": getattr(children, "complete_size", list(children.weight.size())), + } def set_fp32_attr_for_model(model: Union[nn.Module, nn.ModuleList]): diff --git a/tests/test_training/train_CI.py b/tests/test_training/train_CI.py index b33cf4c38..be8e9847a 100644 --- a/tests/test_training/train_CI.py +++ b/tests/test_training/train_CI.py @@ -20,7 +20,7 @@ from internlm.checkpoint import CheckpointManager # noqa: E402 from internlm.core.context import ParallelMode # noqa: E402 from internlm.core.context import global_context as gpc # noqa: E402 -from internlm.core.trainer import TrainState, Trainer # noqa: E402 +from internlm.core.trainer import Trainer, TrainState # noqa: E402 from internlm.data import ( # noqa: E402 build_train_loader_with_data_type, build_valid_loader_with_data_type, @@ -70,7 +70,8 @@ def check_model_weights(model, ckpt_path, total_equal=False): model2_dict[key.replace("wqkv", "Wqkv")] = model2_dict.pop(key) key = key.replace("wqkv", "Wqkv") if key not in model1_dict: - assert False, f"Error: The key {key} for current model dose not exist in standard ckpt!" + if "Wqkv" not in key: + assert False, f"Error: The key {key} for current model dose not exist in standard ckpt!" for key in model1_dict.keys(): if key in model2_dict: diff --git a/tests/test_utils/test_model_checkpoint.py b/tests/test_utils/test_model_checkpoint.py index 5fe8b3c49..06e364fff 100644 --- a/tests/test_utils/test_model_checkpoint.py +++ b/tests/test_utils/test_model_checkpoint.py @@ -1,7 +1,4 @@ import multiprocessing - -backup_ForkingPickler = multiprocessing.reduction.ForkingPickler -backup_dump = multiprocessing.reduction.dump import os from functools import partial @@ -25,6 +22,9 @@ reset_singletons, ) +backup_ForkingPickler = multiprocessing.reduction.ForkingPickler +backup_dump = multiprocessing.reduction.dump + # (TOTAL_STEP, CKPT_EVERY, SNPASHOT_EVERY) step_info_list = [(8, 4, 2), (3, 4, 2), (1, 6, 3)] ckpt_config_list = [ @@ -201,8 +201,8 @@ def return_latest_save_path(save_ckpt_folder, total_step, snapshot_freq, ckpt_fr @pytest.mark.parametrize("step_info", step_info_list) @pytest.mark.parametrize("ckpt_config", ckpt_config_list) def test_ckpt_mm(step_info, ckpt_config, init_dist_and_model): # noqa # pylint: disable=unused-import - from internlm.core.context import global_context as gpc from internlm.checkpoint.checkpoint_manager import CheckpointLoadMask + from internlm.core.context import global_context as gpc ckpt_config = Config(ckpt_config) total_step, checkpoint_every, oss_snapshot_freq = step_info @@ -222,6 +222,8 @@ def test_ckpt_mm(step_info, ckpt_config, init_dist_and_model): # noqa # pylint: ) model, opim = init_dist_and_model + gpc.config._add_item("ckpt", dict()) + gpc.config.ckpt._add_item("universal_ckpt", dict(enable=False, aysnc_save=True, broadcast_load=False)) train_state = TrainState(gpc.config, None) if isinstance(opim, HybridZeroOptimizer): print("Is HybridZeroOptimizer!", flush=True) @@ -297,9 +299,9 @@ def test_ckpt_mm(step_info, ckpt_config, init_dist_and_model): # noqa # pylint: def query_quit_file(rank, world_size=2): + from internlm.checkpoint.checkpoint_manager import CheckpointSaveType from internlm.core.context import global_context as gpc from internlm.initialize import initialize_distributed_env - from internlm.checkpoint.checkpoint_manager import CheckpointSaveType ckpt_config = Config( dict( @@ -348,8 +350,6 @@ def query_quit_file(rank, world_size=2): def test_quit_siganl_handler(): # noqa # pylint: disable=unused-import - import multiprocessing - # we do hack here to workaround the bug of 3rd party library dill, which only occurs in this unittest: # https://github.com/uqfoundation/dill/issues/380 multiprocessing.reduction.ForkingPickler = backup_ForkingPickler