wehub-resource-sync
96 lines
4.8 KiB
Python
96 lines
4.8 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import torch
|
|
from deepspeed.ops.op_builder import AsyncIOBuilder, GDSBuilder
|
|
from deepspeed.io import MockFileWriter, PyFileWriter, FastFileWriter, FastFileWriterConfig
|
|
from deepspeed.runtime.swap_tensor.constants import *
|
|
from .constants import *
|
|
from deepspeed.accelerator import get_accelerator
|
|
|
|
|
|
class CheckpointWriterFactory(object):
|
|
|
|
def __init__(self, writer_config, aio_config, dp_writer_config):
|
|
self._type = writer_config[CHECKPOINT_WRITER_TYPE]
|
|
self._io_buffer_size = writer_config[CHECKPOINT_IO_BUFFER_SIZE]
|
|
self._io_buffer_double = writer_config[CHECKPOINT_IO_BUFFER_DOUBLE]
|
|
self._data_parallel_writer = dp_writer_config
|
|
self._io_multiplier = writer_config[CHECKPOINT_IO_MULTIPLIER]
|
|
if self._data_parallel_writer.pure_dp:
|
|
self._show_statistics = writer_config[CHECKPOINT_IO_STATISTICS] and self._data_parallel_writer is not None
|
|
else:
|
|
self._show_statistics = writer_config[CHECKPOINT_IO_STATISTICS] and self._data_parallel_writer is not None
|
|
self._io_buffer = None
|
|
self._dnvme_handle = None
|
|
self._writer = None
|
|
self._use_gds = False
|
|
|
|
if self._type == CheckpointWriterType.FAST:
|
|
self._use_gds = aio_config[AIO_USE_GDS]
|
|
if self._use_gds:
|
|
self._setup_for_gds(aio_config)
|
|
else:
|
|
self._setup_for_aio(aio_config)
|
|
print(
|
|
f'WriterFactory: self._data_parallel_writer={self._data_parallel_writer} self._show_statistics={self._show_statistics}'
|
|
)
|
|
|
|
def create_writer(self, file_path, optimize_dp_state):
|
|
assert self._writer is None, \
|
|
f'Cannot create checkpoint writer for {file_path} because writer is currently used for {self._writer.file_path()}.\
|
|
Must call writer.release() before reusing to avoid this error.'
|
|
|
|
if self._type == CheckpointWriterType.MOCK:
|
|
self._writer = MockFileWriter(file_path)
|
|
elif self._type == CheckpointWriterType.PYTHON:
|
|
self._writer = PyFileWriter(file_path)
|
|
else:
|
|
if optimize_dp_state:
|
|
num_parallel_writers = self._data_parallel_writer.world_size * self._io_multiplier
|
|
writer_rank = self._data_parallel_writer.rank
|
|
file_path = f'{file_path}-{writer_rank}.{num_parallel_writers}'
|
|
# print(f'create_dp_writer: {self._data_parallel_writer.global_rank=} {writer_rank=} {num_parallel_writers=} {file_path=}')
|
|
else:
|
|
num_parallel_writers = 1
|
|
writer_rank = 0
|
|
# print(f'create_rank0_writer: {self._data_parallel_writer.global_rank=} {writer_rank=} {num_parallel_writers=} {file_path=}')
|
|
|
|
config = FastFileWriterConfig(dnvme_handle=self._dnvme_handle,
|
|
pinned_tensor=self._io_buffer,
|
|
double_buffer=self._io_buffer_double,
|
|
num_parallel_writers=num_parallel_writers,
|
|
writer_rank=writer_rank,
|
|
global_rank=self._data_parallel_writer.global_rank)
|
|
self._writer = FastFileWriter(file_path=file_path, config=config)
|
|
|
|
return self._writer
|
|
|
|
def release_writer(self):
|
|
self._writer.close()
|
|
if self._show_statistics:
|
|
self._writer._dump_state()
|
|
self._writer = None
|
|
|
|
def _setup_for_aio(self, aio_config):
|
|
self._io_buffer = torch.zeros(self._io_buffer_size, dtype=torch.uint8, device='cpu').pin_memory()
|
|
self._dnvme_handle = AsyncIOBuilder().load().aio_handle(
|
|
block_size=aio_config[AIO_BLOCK_SIZE],
|
|
queue_depth=aio_config[AIO_QUEUE_DEPTH],
|
|
single_submit=aio_config[AIO_SINGLE_SUBMIT],
|
|
overlap_events=aio_config[AIO_OVERLAP_EVENTS],
|
|
intra_op_parallelism=aio_config[AIO_INTRA_OP_PARALLELISM])
|
|
|
|
def _setup_for_gds(self, aio_config):
|
|
self._io_buffer = torch.zeros(self._io_buffer_size,
|
|
dtype=torch.uint8,
|
|
device=get_accelerator().current_device_name())
|
|
self._dnvme_handle = GDSBuilder().load().gds_handle(block_size=aio_config[AIO_BLOCK_SIZE],
|
|
queue_depth=aio_config[AIO_QUEUE_DEPTH],
|
|
single_submit=aio_config[AIO_SINGLE_SUBMIT],
|
|
overlap_events=aio_config[AIO_OVERLAP_EVENTS],
|
|
intra_op_parallelism=aio_config[AIO_INTRA_OP_PARALLELISM])
|
|
self._dnvme_handle.pin_device_tensor(self._io_buffer)
|