File size: 1,257 Bytes
b8c8910 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | # Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
import os
import abc
from abc import ABC
from dataclasses import dataclass
@dataclass
class CheckpointCommitInfo(object):
tag: str
save_dir: str
save_latest: bool
class CheckpointEngine(ABC):
# init checkpoint engine for save/load
def __init__(self, config_params=None):
self.name = None
@abc.abstractmethod
def create(self, info: CheckpointCommitInfo):
# create checkpoint on give tag for save/load.
...
@abc.abstractmethod
def save(self, state_dict, path: str):
...
def makedirs(self, path, exist_ok=False):
os.makedirs(path, exist_ok=exist_ok)
@abc.abstractmethod
def load(self, path: str, map_location=None):
...
@abc.abstractmethod
def commit(self, info: CheckpointCommitInfo):
# to tell checkpoint services if all files are ready.
...
def is_data_parallel_writer(self, dp_rank):
return dp_rank == 0
def is_decoupled(self):
return False
def set_commit_info(self, info: CheckpointCommitInfo):
pass
def get_commit_info(self):
return None
def cleanup(self):
pass
|