wehub-resource-sync
230 lines
10 KiB
Python
230 lines
10 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
"""
|
|
Functionality of swapping optimizer tensors to/from (NVMe) storage devices.
|
|
"""
|
|
|
|
from deepspeed.utils.logging import logger
|
|
from deepspeed.ops.op_builder import AsyncIOBuilder
|
|
from deepspeed import comm as dist
|
|
|
|
from deepspeed.runtime.swap_tensor.constants import *
|
|
from deepspeed.runtime.swap_tensor.utils import swap_in_tensors, swap_out_tensors, print_object, \
|
|
get_sized_buffers
|
|
from deepspeed.runtime.swap_tensor.async_swapper import AsyncTensorSwapper
|
|
from deepspeed.runtime.swap_tensor.optimizer_utils import OptimizerSwapper
|
|
from deepspeed.accelerator import get_accelerator
|
|
|
|
DEBUG_MODE = False
|
|
|
|
SWAP_IN_PARAM_TIMER = 'swap_in_param'
|
|
SWAP_OUT_PARAM_TIMER = 'swap_out_param'
|
|
SWAP_IN_GRADIENT_TIMER = 'swap_in_gradient'
|
|
|
|
|
|
class PartitionedOptimizerSwapper(OptimizerSwapper):
|
|
|
|
def __init__(self, swap_config, aio_config, base_folder, optimizer, largest_numel, device, dtype, timers):
|
|
super(PartitionedOptimizerSwapper, self).__init__(swap_config, aio_config, base_folder, optimizer,
|
|
largest_numel, device, dtype, timers)
|
|
|
|
aio_op = AsyncIOBuilder().load()
|
|
self.aio_handle = aio_op.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])
|
|
|
|
# Overlap swapping out
|
|
self.gradient_swapper = AsyncTensorSwapper(aio_handle=self.aio_handle,
|
|
numel_alignment=self.numel_alignment,
|
|
timers=self.timers)
|
|
|
|
self.print_exclude_list += ['aio_handle', 'gradient_swapper', 'print_exclude_list']
|
|
|
|
if dist.get_rank() == 0:
|
|
print_object(obj=self, name='PartitionedOptimizerSwapper', exclude_list=self.print_exclude_list)
|
|
|
|
def initialize_parameters(self, parameters, src_tensors):
|
|
self._initialize_parameters(parameters=parameters, src_tensors=src_tensors, aio_handle=self.aio_handle)
|
|
|
|
def initialize_from_swapped_fp16_params(self, fp16_partitions_info, fp16_num_elems, fp16_pinned_buffers,
|
|
fp32_parameters):
|
|
self._initialize_from_swapped_fp16_params(aio_handle=self.aio_handle,
|
|
fp16_partitions_info=fp16_partitions_info,
|
|
fp16_num_elems=fp16_num_elems,
|
|
fp16_pinned_buffers=fp16_pinned_buffers,
|
|
fp32_parameters=fp32_parameters)
|
|
|
|
def flush_gradients(self):
|
|
self._flush_gradient_swapper(self.gradient_swapper)
|
|
|
|
def release_swap_buffers(self, parameter):
|
|
swap_info = self._get_param_swap_info(parameter)
|
|
if swap_info is None:
|
|
return
|
|
swap_info.release_memory()
|
|
|
|
self.swap_buffer_manager.free(swap_info.swap_buffers)
|
|
swap_info.swap_buffers = []
|
|
|
|
def swap_in_optimizer_state(self, parameter, async_parameter=None):
|
|
swap_info = self._get_param_swap_info(parameter)
|
|
if swap_info is None:
|
|
return
|
|
|
|
self._flush_gradient_swapper(self.gradient_swapper)
|
|
|
|
required_buffer_count = swap_info.num_tensors() + (1 if swap_info.has_gradients() else 0)
|
|
aligned_numel = self._io_aligned_numel(swap_info.numel())
|
|
pinned_buffers = self.swap_buffer_manager.allocate(num_elems=aligned_numel,
|
|
count=required_buffer_count,
|
|
dtype=parameter.dtype)
|
|
assert pinned_buffers is not None
|
|
swap_info.swap_buffers = pinned_buffers.copy()
|
|
|
|
self._start_timer(SWAP_IN_PARAM_TIMER)
|
|
self._swap_in_parameter(aio_handle=self.aio_handle,
|
|
parameter=parameter,
|
|
dest_buffers=pinned_buffers[:swap_info.num_tensors()])
|
|
self._stop_timer(SWAP_IN_PARAM_TIMER)
|
|
self.timer_names.add(SWAP_IN_PARAM_TIMER)
|
|
|
|
if swap_info.has_gradients():
|
|
self._start_timer(SWAP_IN_GRADIENT_TIMER)
|
|
self._swap_in_gradients(aio_handle=self.aio_handle, parameter=parameter, dest_buffer=pinned_buffers[-1])
|
|
self._stop_timer(SWAP_IN_GRADIENT_TIMER)
|
|
self.timer_names.add(SWAP_IN_GRADIENT_TIMER)
|
|
|
|
def _swap_out_optimizer_state(self, swap_info):
|
|
pinned_tensors, pinned_paths = swap_info.get_swap_buffers_and_paths(True)
|
|
WRITE_TIMER = 'swap_submit_write'
|
|
self._start_timer(WRITE_TIMER)
|
|
|
|
swap_out_tensors(self.aio_handle, pinned_tensors, pinned_paths)
|
|
assert self.aio_handle.wait() == len(pinned_tensors)
|
|
|
|
unpinned_tensors, unpinned_paths = swap_info.get_swap_buffers_and_paths(False)
|
|
if len(unpinned_tensors) > 0:
|
|
pinned_buffers = self.swap_buffer_manager.allocate_all(num_elems=self.largest_numel, dtype=self.dtype)
|
|
self._swap_out_unpinned_tensors(aio_handle=self.aio_handle,
|
|
unpinned_tensors=unpinned_tensors,
|
|
dest_paths=unpinned_paths,
|
|
pinned_buffers=pinned_buffers)
|
|
swap_info.swap_buffers += pinned_buffers.copy()
|
|
|
|
self._stop_timer(WRITE_TIMER)
|
|
self._log_timers([WRITE_TIMER])
|
|
|
|
def writeback_optimizer_state_and_gradients(self, parameter, write_opt_state, write_gradients):
|
|
swap_info = self._get_param_swap_info(parameter=parameter)
|
|
|
|
if swap_info is None:
|
|
return
|
|
|
|
if write_opt_state:
|
|
self._swap_out_optimizer_state(swap_info)
|
|
|
|
if write_gradients and swap_info.has_gradients():
|
|
param_gradients = swap_info.swapped_gradients.values()
|
|
swap_buffers = [parameter.grad.narrow(0, grad.offset, grad.length) for grad in param_gradients]
|
|
swap_paths = [grad.path for grad in param_gradients]
|
|
swap_out_tensors(self.aio_handle, swap_buffers, swap_paths)
|
|
assert len(swap_buffers) == self.aio_handle.wait()
|
|
if swap_info.unswapped_gradients:
|
|
swap_info.write_unswapped_gradients(src_buffer=parameter.grad)
|
|
|
|
self.release_swap_buffers(parameter)
|
|
|
|
def swap_out_optimizer_state(self, parameter, async_swap=False):
|
|
swap_info = self._get_param_swap_info(parameter=parameter)
|
|
|
|
if swap_info is None:
|
|
return
|
|
|
|
swap_bytes = sum(
|
|
[self._io_aligned_numel(t.numel()) * t.element_size() for t in swap_info.get_compute_tensors()])
|
|
|
|
self._start_timer(SWAP_OUT_PARAM_TIMER)
|
|
self._swap_out_optimizer_state(swap_info)
|
|
self.release_swap_buffers(parameter)
|
|
self._stop_timer(SWAP_OUT_PARAM_TIMER)
|
|
self.timer_names.add(SWAP_OUT_PARAM_TIMER)
|
|
|
|
if DEBUG_MODE and dist.get_rank() == 0:
|
|
logger.info(f'optimizer_param_swap_out: {(swap_bytes/(1024**3)):5.2f} GB')
|
|
|
|
def swap_out_gradients(self, parameter, gradient_offsets, gradient_tensors):
|
|
self._swap_out_gradients(parameter=parameter,
|
|
gradient_offsets=gradient_offsets,
|
|
gradient_tensors=gradient_tensors,
|
|
gradient_swapper=self.gradient_swapper)
|
|
|
|
def _swap_in_parameter(self, aio_handle, parameter, dest_buffers):
|
|
swap_info = self._get_param_swap_info(parameter)
|
|
if swap_info is None:
|
|
return
|
|
|
|
num_swap_tensors = swap_info.num_tensors()
|
|
assert num_swap_tensors <= len(dest_buffers)
|
|
|
|
swap_lengths = [self._io_aligned_numel(swap_info.numel())] * num_swap_tensors
|
|
swap_buffers = get_sized_buffers(dest_buffers, swap_lengths)
|
|
|
|
compute_lengths = [swap_info.numel()] * num_swap_tensors
|
|
compute_buffers = get_sized_buffers(dest_buffers, compute_lengths)
|
|
|
|
READ_TIMER = 'swap_submit_read_param'
|
|
WAIT_TIMER = 'swap_wait_read_param'
|
|
|
|
self._start_timer(READ_TIMER)
|
|
swap_in_tensors(aio_handle, swap_buffers, swap_info.get_swap_paths())
|
|
self._stop_timer(READ_TIMER)
|
|
|
|
swap_bytes = sum([buffer.numel() * buffer.element_size() for buffer in swap_buffers])
|
|
|
|
self._start_timer(WAIT_TIMER)
|
|
aio_handle.wait()
|
|
self._stop_timer(WAIT_TIMER)
|
|
|
|
swap_info.set_swap_buffers(dest_buffers, self._io_aligned_numel(swap_info.numel()))
|
|
|
|
self._log_timers([READ_TIMER, WAIT_TIMER])
|
|
if DEBUG_MODE and dist.get_rank() == 0:
|
|
logger.info(f'optimizer_param_swap_in: {(swap_bytes/(1024**3)):5.2f} GB')
|
|
|
|
def _swap_in_pinned_gradients(self, aio_handle, parameter, gradient_tensor):
|
|
swap_info = self.swap_params_info[OptimizerSwapper.parameter_id(parameter)]
|
|
param_gradients = swap_info.swapped_gradients.values()
|
|
swap_buffers = [gradient_tensor.narrow(0, grad.offset, grad.length) for grad in param_gradients]
|
|
swap_paths = [grad.path for grad in param_gradients]
|
|
SWAP_READ_GRADIENTS = 'swap_submit_read_gradient'
|
|
SWAP_WAIT_GRADIENTS = 'swap_submit_wait_gradient'
|
|
self._start_timer(SWAP_READ_GRADIENTS)
|
|
swap_in_tensors(aio_handle, swap_buffers, swap_paths)
|
|
self._stop_timer(SWAP_READ_GRADIENTS)
|
|
|
|
self._start_timer(SWAP_WAIT_GRADIENTS)
|
|
assert len(swap_buffers) == aio_handle.wait()
|
|
self._stop_timer(SWAP_WAIT_GRADIENTS)
|
|
|
|
self._log_timers([SWAP_READ_GRADIENTS, SWAP_WAIT_GRADIENTS])
|
|
|
|
def _swap_in_gradients(self, aio_handle, parameter, dest_buffer):
|
|
swap_info = self.swap_params_info.get(OptimizerSwapper.parameter_id(parameter), None)
|
|
if not (swap_info and swap_info.has_gradients()):
|
|
return
|
|
|
|
assert get_accelerator().is_pinned(dest_buffer)
|
|
assert parameter.numel() <= dest_buffer.numel()
|
|
|
|
parameter.grad = dest_buffer.narrow(0, 0, parameter.numel())
|
|
|
|
if swap_info.swapped_gradients:
|
|
self._swap_in_pinned_gradients(aio_handle, parameter, parameter.grad)
|
|
|
|
if swap_info.unswapped_gradients:
|
|
self._retrieve_unswapped_grad_partitions(swap_info=swap_info, dest_buffer=parameter.grad)
|