chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,37 @@
|
||||
# Checkpoint Engine
|
||||
|
||||
|
||||
The `CheckpointEngine` was designed to modularized the checkpoint serialization. In this way, we can simply replace/refine the checkpoint serialization methods.
|
||||
|
||||
### Interface for `CheckpointEngine`
|
||||
|
||||
Basically, for checkpoint management(save/load by deepspeed with the given tag), the `CheckpointEngine` will:
|
||||
|
||||
1. To make preliminaries ready by call `create(tag)`. For `torch`, we can just log some extra info as `torch` can directly call `save/load` without other preparation.
|
||||
|
||||
2. After the `create(tag)`, deepspeed can call `save/load` to persist files into disk/memory/etc.
|
||||
|
||||
3. When all the files for a tag are ready, deepspeed engine will call `commit()` to tell the checkpoint engine current checkpoint is complete. For original torch, it also plays the role of logger.
|
||||
|
||||
|
||||
```python
|
||||
class CheckpointEngine(object):
|
||||
# init checkpoint engine for save/load
|
||||
def __init__(self, config_params=None):
|
||||
pass
|
||||
|
||||
def create(self, info:CheckpointCommitInfo):
|
||||
# create checkpoint on give tag for save/load.
|
||||
pass
|
||||
|
||||
def save(self, state_dict, path: str):
|
||||
pass
|
||||
|
||||
def load(self, path: str, map_location=None):
|
||||
pass
|
||||
|
||||
def commit(self, info:CheckpointCommitInfo):
|
||||
# to tell checkpoint services if all files are readys.
|
||||
pass
|
||||
|
||||
```
|
||||
@@ -0,0 +1,12 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
'''Copyright The Microsoft DeepSpeed Team'''
|
||||
|
||||
from .fast_checkpoint_engine import FastCheckpointEngine
|
||||
from .torch_checkpoint_engine import TorchCheckpointEngine
|
||||
from .decoupled_checkpoint_engine import DecoupledCheckpointEngine
|
||||
from .checkpoint_engine import CheckpointCommitInfo
|
||||
from .datastates_checkpoint_engine import DataStatesCheckpointEngine
|
||||
from .utils import create_checkpoint_engine
|
||||
@@ -0,0 +1,63 @@
|
||||
# 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
|
||||
|
||||
def preserves_storage_sharing(self):
|
||||
return True
|
||||
@@ -0,0 +1,59 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Apache-2.0 License Copyright (c) UChicago Argonne LLC, operator of Argonne National Laboratory.
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from deepspeed.runtime.checkpoint_engine.checkpoint_engine import \
|
||||
CheckpointEngine, CheckpointCommitInfo
|
||||
|
||||
ENGINE_NAME = "DataStatesCheckpointEngine"
|
||||
|
||||
|
||||
class DataStatesCheckpointEngine(CheckpointEngine):
|
||||
|
||||
def __init__(self, deepspeed_config, rank):
|
||||
super().__init__(deepspeed_config)
|
||||
self.commit_info = None
|
||||
self.ckpt_engine = None
|
||||
try:
|
||||
from datastates import CheckpointEngine as DataStatesEngine
|
||||
self.ckpt_engine = DataStatesEngine(deepspeed_config, rank)
|
||||
except ImportError:
|
||||
raise RuntimeError("Please install DataStates from https://github.com/DataStates/datastates-llm.")
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"An error occurred while initializing DataStates Checkpoint Engine: {e}")
|
||||
|
||||
def __del__(self):
|
||||
self.cleanup()
|
||||
|
||||
def create(self, info: CheckpointCommitInfo):
|
||||
self.commit_info = info
|
||||
return None
|
||||
|
||||
def save(self, state_dict, path: str):
|
||||
return self.ckpt_engine.save(state_dict, path)
|
||||
|
||||
def load(self, path: str, map_location=None):
|
||||
return self.ckpt_engine.load(path, map_location)
|
||||
|
||||
def commit(self, info: CheckpointCommitInfo):
|
||||
if info is None:
|
||||
return
|
||||
assert info == self.commit_info
|
||||
self.ckpt_engine.wait(persist=True)
|
||||
self.commit_info = None
|
||||
return True
|
||||
|
||||
def cleanup(self):
|
||||
self.commit(self.commit_info)
|
||||
if self.ckpt_engine:
|
||||
self.ckpt_engine.wait(persist=True)
|
||||
del self.ckpt_engine
|
||||
|
||||
def is_decoupled(self):
|
||||
return True
|
||||
|
||||
def preserves_storage_sharing(self):
|
||||
return False
|
||||
@@ -0,0 +1,243 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
import torch.multiprocessing as mp
|
||||
from deepspeed.runtime.checkpoint_engine.checkpoint_engine import \
|
||||
CheckpointEngine, CheckpointCommitInfo
|
||||
from deepspeed.runtime.checkpoint_engine.fast_checkpoint_engine import FastCheckpointEngine
|
||||
from deepspeed import comm as dist
|
||||
from deepspeed.runtime.utils import get_checkpoint_folder_size
|
||||
from deepspeed.utils import logger
|
||||
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class DecoupledEvent(Enum):
|
||||
SAVE_EVENT = 1
|
||||
COMMIT_EVENT = 2
|
||||
EXIT_EVENT = 3
|
||||
|
||||
|
||||
class CheckpointSize(object):
|
||||
|
||||
def __init__(self):
|
||||
self._pre = None
|
||||
self._post = None
|
||||
self._gigabytes = None
|
||||
|
||||
def gb_size(self):
|
||||
return self._gigabytes
|
||||
|
||||
def set_pre_size(self, size):
|
||||
self._pre = size
|
||||
|
||||
def set_post_size(self, size):
|
||||
self._post = size
|
||||
self._gigabytes = (self._post - self._pre) / (1024**3)
|
||||
|
||||
|
||||
def init_decoupled_checkpoint(config_params, dp_writer_config, save_event, save_queue, optimize_dp_state):
|
||||
try:
|
||||
checkpoint_engine = FastCheckpointEngine(config_params, dp_writer_config, optimize_dp_state)
|
||||
print('Created FastCheckpointEngine for Decoupled Checkpointing')
|
||||
save_path_list = []
|
||||
while True:
|
||||
(save_info, event_type) = save_queue.get()
|
||||
if event_type == DecoupledEvent.SAVE_EVENT and save_info is not None:
|
||||
state_dict, save_path = save_info
|
||||
# print(f'Received decoupled checkpoint request for {save_path=}')
|
||||
save_path_list.append(save_path)
|
||||
checkpoint_engine.save(state_dict, save_path)
|
||||
del state_dict
|
||||
# print(f'Completed decoupled checkpoint request for {save_path=}')
|
||||
|
||||
if event_type == DecoupledEvent.COMMIT_EVENT:
|
||||
# print(f'Recieved commit request for {save_path_list=}')
|
||||
save_path_list = []
|
||||
save_event.set()
|
||||
|
||||
if event_type == DecoupledEvent.EXIT_EVENT:
|
||||
# print(f'Received decoupled exit request')
|
||||
break
|
||||
except Exception as e:
|
||||
print(f'[{ENGINE_NAME}] Checkpoint subprocess crashed with error: {e}')
|
||||
raise
|
||||
|
||||
|
||||
ENGINE_NAME = "DecoupledCheckpointEngine"
|
||||
|
||||
# Default timeout for checkpoint operations (5 minutes)
|
||||
DEFAULT_CHECKPOINT_TIMEOUT_SECONDS = 300
|
||||
# Interval for checking process health while waiting
|
||||
PROCESS_HEALTH_CHECK_INTERVAL_SECONDS = 10
|
||||
|
||||
|
||||
class DecoupledCheckpointEngine(CheckpointEngine):
|
||||
|
||||
def __init__(self, config_params, dp_writer_config, optimize_dp_state):
|
||||
# Set spawn method if not already set (needed for CUDA tensor sharing)
|
||||
try:
|
||||
mp.set_start_method('spawn')
|
||||
except RuntimeError:
|
||||
pass # Already set, ignore
|
||||
super().__init__(config_params)
|
||||
self.name = ENGINE_NAME
|
||||
self.dp_writer_config = dp_writer_config
|
||||
self.commit_info = None
|
||||
self.checkpoint_size = CheckpointSize()
|
||||
self.global_rank = dist.get_rank()
|
||||
self.optimize_dp_state = optimize_dp_state
|
||||
self._cleanup_called = False
|
||||
if dp_writer_config is None:
|
||||
self.save_event = None
|
||||
self.save_queue = None
|
||||
self.ckpt_process = None
|
||||
self.local_rank = None
|
||||
print(
|
||||
f'[{ENGINE_NAME}]: No checkpoint process self.global_rank={self.global_rank} self.dp_writer_config={self.dp_writer_config}'
|
||||
)
|
||||
else:
|
||||
self.save_event = mp.Event()
|
||||
self.save_queue = mp.SimpleQueue()
|
||||
engine_args = (config_params, dp_writer_config, self.save_event, self.save_queue, self.optimize_dp_state)
|
||||
self.ckpt_process = mp.Process(target=init_decoupled_checkpoint, args=engine_args)
|
||||
self.ckpt_process.start()
|
||||
self.local_rank = dp_writer_config.local_rank
|
||||
print(
|
||||
f'[{ENGINE_NAME}]: Create checkpoint process self.global_rank={self.global_rank} self.ckpt_process.pid={self.ckpt_process.pid} self.dp_writer_config={self.dp_writer_config}'
|
||||
)
|
||||
|
||||
def __del__(self):
|
||||
try:
|
||||
self.cleanup()
|
||||
except Exception:
|
||||
# Suppress exceptions in destructor to avoid crashes during shutdown
|
||||
pass
|
||||
|
||||
def _check_process_alive(self):
|
||||
"""Check if the checkpoint process is still alive.
|
||||
|
||||
Note: Only call this when self.ckpt_process is not None.
|
||||
Some ranks don't have a checkpoint process by design (see Figure 6 in paper).
|
||||
"""
|
||||
return self.ckpt_process.is_alive()
|
||||
|
||||
def _wait_for_event_with_timeout(self, timeout_seconds=DEFAULT_CHECKPOINT_TIMEOUT_SECONDS):
|
||||
"""Wait for save_event with timeout and process health checks.
|
||||
|
||||
Returns True if event was set, raises RuntimeError if process died or timeout occurred.
|
||||
"""
|
||||
elapsed = 0
|
||||
while elapsed < timeout_seconds:
|
||||
if self.save_event.wait(timeout=PROCESS_HEALTH_CHECK_INTERVAL_SECONDS):
|
||||
return True
|
||||
elapsed += PROCESS_HEALTH_CHECK_INTERVAL_SECONDS
|
||||
|
||||
# Check if process is still alive
|
||||
if not self._check_process_alive():
|
||||
raise RuntimeError(f"[{ENGINE_NAME}] Checkpoint process died unexpectedly. "
|
||||
f"Check logs for OOM or other errors in the checkpoint subprocess.")
|
||||
|
||||
raise RuntimeError(f"[{ENGINE_NAME}] Checkpoint commit timed out after {timeout_seconds} seconds. "
|
||||
f"Process alive: {self._check_process_alive()}")
|
||||
|
||||
def create(self, info: CheckpointCommitInfo):
|
||||
self.commit_info = info
|
||||
if self.checkpoint_size.gb_size() is None:
|
||||
pre_size = get_checkpoint_folder_size(info.save_dir, info.tag, self.local_rank)
|
||||
self.checkpoint_size.set_pre_size(pre_size)
|
||||
|
||||
def load(self, path: str, map_location=None):
|
||||
sd = torch.load(path, map_location=map_location)
|
||||
return sd
|
||||
|
||||
def save(self, state_dict, path: str):
|
||||
if self.ckpt_process is None:
|
||||
return
|
||||
|
||||
# Check process health before attempting to save
|
||||
if not self._check_process_alive():
|
||||
return
|
||||
|
||||
save_info = (state_dict, path)
|
||||
self.save_queue.put((save_info, DecoupledEvent.SAVE_EVENT))
|
||||
|
||||
def commit(self, info: CheckpointCommitInfo):
|
||||
# Use proper validation instead of assert (assert is disabled with python -O)
|
||||
if info != self.commit_info:
|
||||
raise ValueError(f"[{ENGINE_NAME}] Checkpoint commit info mismatch: "
|
||||
f"expected {self.commit_info}, got {info}")
|
||||
|
||||
if self.ckpt_process is not None:
|
||||
# Check process health before waiting
|
||||
if not self._check_process_alive():
|
||||
raise RuntimeError(f"[{ENGINE_NAME}] Cannot commit checkpoint: checkpoint process is not running.")
|
||||
|
||||
self.save_queue.put((None, DecoupledEvent.COMMIT_EVENT))
|
||||
# Wait with timeout and health checks instead of blocking forever
|
||||
self._wait_for_event_with_timeout()
|
||||
self.save_event.clear()
|
||||
|
||||
self.commit_info = None
|
||||
|
||||
if self.checkpoint_size.gb_size() is None:
|
||||
dist.barrier()
|
||||
post_size = get_checkpoint_folder_size(info.save_dir, info.tag, self.local_rank)
|
||||
self.checkpoint_size.set_post_size(post_size)
|
||||
|
||||
assert self.checkpoint_size.gb_size() is not None, "Checkpoint size should be set after commit"
|
||||
|
||||
if self.global_rank == 0:
|
||||
print(
|
||||
f'{self.name} self.global_rank={self.global_rank} created checkpoint of {round(self.checkpoint_size.gb_size(), 2)} GB'
|
||||
)
|
||||
|
||||
return True
|
||||
|
||||
def get_commit_info(self):
|
||||
# print(f'getting commit info {self.commit_info=}')
|
||||
return self.commit_info
|
||||
|
||||
def is_decoupled(self):
|
||||
return True
|
||||
|
||||
def cleanup(self):
|
||||
# Prevent multiple cleanup calls (especially from __del__)
|
||||
if self._cleanup_called:
|
||||
return
|
||||
self._cleanup_called = True
|
||||
|
||||
try:
|
||||
if self.get_commit_info() is not None:
|
||||
self.commit(self.commit_info)
|
||||
except Exception as e:
|
||||
logger.warning(f"[{ENGINE_NAME}] Error during commit in cleanup: {e}")
|
||||
|
||||
if self.ckpt_process is not None:
|
||||
try:
|
||||
self.save_queue.put((None, DecoupledEvent.EXIT_EVENT))
|
||||
except Exception:
|
||||
pass # Queue may be broken if process died
|
||||
|
||||
# Join with timeout to avoid hanging forever
|
||||
self.ckpt_process.join(timeout=DEFAULT_CHECKPOINT_TIMEOUT_SECONDS)
|
||||
|
||||
# If process didn't exit, terminate it forcefully
|
||||
if self.ckpt_process.is_alive():
|
||||
logger.warning(
|
||||
f"[{ENGINE_NAME}] Checkpoint process did not exit within timeout, terminating forcefully.")
|
||||
self.ckpt_process.terminate()
|
||||
self.ckpt_process.join(timeout=5) # Brief wait after terminate
|
||||
|
||||
# Last resort: kill
|
||||
if self.ckpt_process.is_alive():
|
||||
self.ckpt_process.kill()
|
||||
|
||||
self.ckpt_process = None
|
||||
self.save_queue = None
|
||||
|
||||
def is_data_parallel_writer(self, dp_rank):
|
||||
return self.ckpt_process is not None
|
||||
@@ -0,0 +1,50 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
from deepspeed.runtime.checkpoint_engine.checkpoint_engine import \
|
||||
CheckpointEngine, CheckpointCommitInfo
|
||||
from deepspeed.runtime.model_checkpointing import (
|
||||
CHECKPOINT_WRITER,
|
||||
CHECKPOINT_SERIALIZATION,
|
||||
CheckpointWriterFactory,
|
||||
)
|
||||
|
||||
|
||||
class FastCheckpointEngine(CheckpointEngine):
|
||||
|
||||
def __init__(self, config_params, dp_writer_config, optimize_dp_state):
|
||||
super().__init__(config_params)
|
||||
self.name = 'FastCheckpointEngine'
|
||||
self.serialization_enabled = config_params.checkpoint_config[CHECKPOINT_SERIALIZATION]
|
||||
self.optimize_dp_state = optimize_dp_state
|
||||
if dp_writer_config is None:
|
||||
self._writer = None
|
||||
else:
|
||||
self._writer = CheckpointWriterFactory(writer_config=config_params.checkpoint_config[CHECKPOINT_WRITER],
|
||||
aio_config=config_params.aio_config,
|
||||
dp_writer_config=dp_writer_config)
|
||||
|
||||
def create(self, info: CheckpointCommitInfo):
|
||||
pass
|
||||
|
||||
def save(self, state_dict, path: str):
|
||||
if self._writer is None:
|
||||
return
|
||||
|
||||
torch.save(obj=state_dict,
|
||||
f=self._writer.create_writer(path, self.optimize_dp_state),
|
||||
_use_new_zipfile_serialization=self.serialization_enabled)
|
||||
self._writer.release_writer()
|
||||
|
||||
def load(self, path: str, map_location=None):
|
||||
sd = torch.load(path, map_location=map_location)
|
||||
return sd
|
||||
|
||||
def commit(self, info: CheckpointCommitInfo):
|
||||
return True
|
||||
|
||||
def is_data_parallel_writer(self, dp_rank):
|
||||
return self._writer is not None
|
||||
@@ -0,0 +1,108 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import os
|
||||
import torch
|
||||
import torch_nebula
|
||||
|
||||
from deepspeed.runtime.checkpoint_engine.checkpoint_engine import \
|
||||
CheckpointEngine, CheckpointCommitInfo
|
||||
from deepspeed.utils import logger, log_dist
|
||||
from deepspeed.nebula.constants import *
|
||||
|
||||
|
||||
def _get_tag_from_path(path):
|
||||
return os.path.basename(os.path.dirname(path))
|
||||
|
||||
|
||||
class NebulaCheckpointEngine(CheckpointEngine):
|
||||
|
||||
def __init__(self, config_params=None):
|
||||
super().__init__(config_params)
|
||||
self.name = "NebulaCheckpointEngine"
|
||||
self.checkpoint = None
|
||||
self.tag_flag = None
|
||||
self.enable_nebula_load = config_params.enable_nebula_load
|
||||
self.nebula_load_path = config_params.load_path
|
||||
if self.nebula_load_path is None:
|
||||
self.nebula_load_path = config_params.persistent_storage_path
|
||||
|
||||
nebula_config_params = {
|
||||
NEBULA_PERSISTENT_STORAGE_PATH: config_params.persistent_storage_path,
|
||||
NEBULA_PERSISTENT_TIME_INTERVAL: config_params.persistent_time_interval,
|
||||
NEBULA_NUM_OF_VERSION_IN_RETENTION: config_params.num_of_version_in_retention,
|
||||
}
|
||||
torch_nebula.init(**nebula_config_params)
|
||||
|
||||
def create(self, info: CheckpointCommitInfo):
|
||||
log_dist(f"[Nebula] Start Checkpoint for tag:{info.tag}", ranks=[0])
|
||||
# -2 means: customer needs to explicitly tell nebula
|
||||
# current checkpoint is complete by commit methond.
|
||||
self.checkpoint = torch_nebula.Checkpoint(info.tag, -2)
|
||||
|
||||
def save(self, state_dict, path: str):
|
||||
log_dist("[Nebula] Create dummy files for loading.")
|
||||
torch.save("", path)
|
||||
|
||||
tag = _get_tag_from_path(path)
|
||||
partititon_name = os.path.basename(path)
|
||||
logger.info(f"[Nebula] Saving {partititon_name} under tag {tag}...")
|
||||
self.checkpoint.save(partititon_name, state_dict)
|
||||
logger.info(f"[Nebula] Saved {partititon_name} under tag {tag}.")
|
||||
|
||||
def load(self, path: str, map_location=None):
|
||||
tag = _get_tag_from_path(path)
|
||||
first_load_flag = self.tag_flag is None or self.tag_flag == tag
|
||||
if not self.enable_nebula_load and first_load_flag:
|
||||
self.tag_flag = tag
|
||||
logger.info(f"[Nebula] Disable nebula load. Loading checkpoint from {path} ...")
|
||||
partition = torch.load(path, map_location=map_location, weights_only=False)
|
||||
logger.info(f"[Nebula] Disable nebula load. Loaded checkpoint from {path} .")
|
||||
return partition
|
||||
|
||||
partition_name = os.path.basename(path)
|
||||
logger.info(f"[Nebula] Loading {path} under tag {tag} from nebula path {self.nebula_load_path}...")
|
||||
|
||||
checkpoint = None
|
||||
if tag in (None, 'latest', 'latest_universal'):
|
||||
# In some cases, there is the inconsistent tag between deepspeed metadata (latest file)
|
||||
# and nebula metadata, will lead to the failure on loading with deepspeed tag. Then we
|
||||
# will try to load the valid latest checkpoint from nebula(tier3 > tier1). So, in summary
|
||||
# when met failure loading for given tag, the loading priority would be like:
|
||||
# nebula tier3 latest > nebula tier1 latest.
|
||||
checkpoint = torch_nebula.get_latest_checkpoint(persist_path=self.nebula_load_path)
|
||||
else:
|
||||
checkpoint = torch_nebula.get_checkpoint(tag=tag, persist_path=self.nebula_load_path)
|
||||
|
||||
if checkpoint is None or (checkpoint is not None and checkpoint.tag == ''):
|
||||
logger.info(
|
||||
f"Unable to find valid checkpoint tag:{tag} from Nebula, try to get latest checkpoint again from nebula {self.nebula_load_path} path!"
|
||||
)
|
||||
# nebula tier3 latest
|
||||
checkpoint = torch_nebula.get_latest_checkpoint(persist_path=self.nebula_load_path)
|
||||
if checkpoint is None or (checkpoint is not None and checkpoint.tag == ''):
|
||||
logger.info(
|
||||
"Unable to find latest checkpoint from Nebula tier3, try to get latest checkpoint again from nebula tier1 path!"
|
||||
)
|
||||
# nebula tier1 latest
|
||||
checkpoint = torch_nebula.get_latest_checkpoint()
|
||||
logger.warning(f"Unable to find valid checkpoint from Nebula under tag:{tag}.")
|
||||
return None
|
||||
|
||||
tag = checkpoint.tag
|
||||
self.tag_flag = -1
|
||||
partition = checkpoint.load(partition_name, map_location=map_location)
|
||||
logger.info(f"[Nebula] Loaded {path} under tag {tag} from {self.nebula_load_path}.")
|
||||
return partition
|
||||
|
||||
def commit(self, info: CheckpointCommitInfo):
|
||||
tag = info.tag
|
||||
# nebula commit will be call when all files under give tag are ready to be persisted in the async way.
|
||||
logger.info(f"[Nebula] all files for {tag} are saved in tier1. It is ready to start persisting")
|
||||
commit_rls = self.checkpoint.commit()
|
||||
if not commit_rls:
|
||||
logger.error("[Nebula] failed to commit the checkpoint, please check the log.")
|
||||
return False
|
||||
return commit_rls
|
||||
@@ -0,0 +1,43 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
from deepspeed.utils import log_dist
|
||||
from deepspeed.runtime.checkpoint_engine.checkpoint_engine import \
|
||||
CheckpointEngine, CheckpointCommitInfo
|
||||
from deepspeed.runtime.model_checkpointing import CHECKPOINT_SERIALIZATION
|
||||
|
||||
ENGINE_NAME = "TorchCheckpointEngine"
|
||||
|
||||
|
||||
class TorchCheckpointEngine(CheckpointEngine):
|
||||
|
||||
def __init__(self, config_params=None):
|
||||
super().__init__(config_params)
|
||||
self.name = ENGINE_NAME
|
||||
if config_params is None:
|
||||
self.zipfile_serialization = False
|
||||
else:
|
||||
self.zipfile_serialization = config_params.checkpoint_config[CHECKPOINT_SERIALIZATION]
|
||||
log_dist(f'[{ENGINE_NAME}] Initialized with serialization = {self.zipfile_serialization}', ranks=[0])
|
||||
|
||||
def create(self, info: CheckpointCommitInfo):
|
||||
log_dist(f"[Torch] Checkpoint {info.tag} is about to be saved!", ranks=[0])
|
||||
pass
|
||||
|
||||
def save(self, state_dict, path: str):
|
||||
# log_dist(f"[Torch] Saving [begin] {path}... {self.zipfile_serialization=}", ranks=[0])
|
||||
torch.save(state_dict, path, _use_new_zipfile_serialization=self.zipfile_serialization)
|
||||
# log_dist(f"[Torch] Saving [end] {path}... {self.zipfile_serialization=}", ranks=[0])
|
||||
|
||||
def load(self, path: str, map_location=None):
|
||||
log_dist(f"[Torch] Begin Load checkpoint from {path}...", ranks=[0])
|
||||
partition = torch.load(path, map_location=map_location, weights_only=False)
|
||||
log_dist(f"[Torch] End Load checkpoint from {path}...", ranks=[0])
|
||||
return partition
|
||||
|
||||
def commit(self, info: CheckpointCommitInfo):
|
||||
#logger.info(f"[Torch] Checkpoint {tag} is ready now!")
|
||||
return True
|
||||
@@ -0,0 +1,48 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from deepspeed.runtime.model_checkpointing.constants import *
|
||||
from deepspeed.runtime.model_checkpointing.utils import create_data_parallel_writer_config
|
||||
from deepspeed.utils import logger
|
||||
from deepspeed import comm as dist
|
||||
from .decoupled_checkpoint_engine import DecoupledCheckpointEngine
|
||||
from .fast_checkpoint_engine import FastCheckpointEngine
|
||||
from .torch_checkpoint_engine import TorchCheckpointEngine
|
||||
|
||||
|
||||
def create_checkpoint_engine(config_params, groups, zero_stage, has_moe_layers, optimize_dp_state):
|
||||
if config_params is not None:
|
||||
if config_params.checkpoint_config[CHECKPOINT_WRITER] is not None:
|
||||
writer_config = config_params.checkpoint_config[CHECKPOINT_WRITER]
|
||||
dp_writer_config = create_data_parallel_writer_config(
|
||||
groups=groups,
|
||||
parallel_unit=writer_config[CHECKPOINT_DATA_PARALLEL],
|
||||
zero_stage=zero_stage,
|
||||
has_moe_layers=has_moe_layers)
|
||||
if writer_config[CHECKPOINT_WRITER_DECOUPLED]:
|
||||
return DecoupledCheckpointEngine(config_params, dp_writer_config, optimize_dp_state)
|
||||
else:
|
||||
return FastCheckpointEngine(config_params, dp_writer_config, optimize_dp_state)
|
||||
|
||||
if config_params is not None and config_params.nebula_config.enabled:
|
||||
try:
|
||||
from .nebula_checkpoint_engine import NebulaCheckpointEngine
|
||||
except ImportError as err:
|
||||
logger.error(f"No torch_nebula was found! Will fall back to torch.save. Details: {err}")
|
||||
return TorchCheckpointEngine(config_params)
|
||||
else:
|
||||
return NebulaCheckpointEngine(config_params=config_params.nebula_config)
|
||||
|
||||
if config_params.datastates_config.enabled:
|
||||
try:
|
||||
from .datastates_checkpoint_engine import DataStatesCheckpointEngine
|
||||
return DataStatesCheckpointEngine(deepspeed_config=config_params, rank=dist.get_rank())
|
||||
except ImportError as err:
|
||||
logger.error(
|
||||
f"No datastates engine found! Install from https://github.com/DataStates/datastates-llm. Will fall back to torch.save. Details: {err}"
|
||||
)
|
||||
return TorchCheckpointEngine(config_params)
|
||||
|
||||
return TorchCheckpointEngine(config_params)
|
||||
Reference in New Issue
Block a user