wehub-resource-sync
176 lines
6.2 KiB
Python
176 lines
6.2 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
"""
|
|
Functionality of swapping tensors to/from (NVMe) storage devices.
|
|
"""
|
|
import torch
|
|
|
|
from deepspeed import comm as dist
|
|
from deepspeed.utils.logging import logger
|
|
from deepspeed.runtime.swap_tensor.utils import swap_out_tensors, SwapBuffer
|
|
from deepspeed.accelerator import get_accelerator
|
|
|
|
INVALID_BUFFER_INDEX = -1
|
|
ASYNC_SWAPPER_WAIT_TIMER = 'async_swap_gradient_wait'
|
|
|
|
|
|
class AsyncTensorSwapper(object):
|
|
|
|
def __init__(self, aio_handle, numel_alignment, timers):
|
|
self.free_buffer_index = []
|
|
self.swapping_buffer_index = []
|
|
self.ready_buffer_index = []
|
|
self.current_buffer_index = INVALID_BUFFER_INDEX
|
|
self.all_buffers = []
|
|
self.aio_handle = aio_handle
|
|
self.numel_alignment = numel_alignment
|
|
self.max_numel = 0
|
|
self.num_pending_swaps = 0
|
|
self.timers = timers
|
|
self.timer_names = set()
|
|
self.num_elements_swapped = 0
|
|
self.dtype = None
|
|
|
|
def has_buffers(self):
|
|
return len(self.all_buffers) > 0
|
|
|
|
def add_buffers(self, buffer_list):
|
|
assert len(self.all_buffers) == 0
|
|
assert all([get_accelerator().is_pinned(buffer) for buffer in buffer_list])
|
|
dtype = buffer_list[0].dtype
|
|
assert all([buffer.dtype == dtype for buffer in buffer_list])
|
|
|
|
self.dtype = dtype
|
|
self.all_buffers = [SwapBuffer(buffer) for buffer in buffer_list]
|
|
self.free_buffer_index += [i for i in range(len(self.all_buffers))]
|
|
self.max_numel = max([buffer.numel() for buffer in buffer_list])
|
|
self.timer_names = set()
|
|
|
|
def get_timer_names(self):
|
|
return list(self.timer_names)
|
|
|
|
def release_buffers(self):
|
|
self._report_statistics('Swapped out[Before flush]')
|
|
self._flush_buffers_until_complete()
|
|
self._report_statistics('Swapped out[After flush]')
|
|
|
|
pinned_buffers = [buf.buffer for buf in self.all_buffers]
|
|
self.all_buffers = []
|
|
self.free_buffer_index = []
|
|
self.current_buffer_index = INVALID_BUFFER_INDEX
|
|
self.num_elements_swapped = 0
|
|
self.dtype = None
|
|
|
|
return pinned_buffers
|
|
|
|
def swap_out_tensors(self, tensor_list, path_list):
|
|
for tensor, swap_path in zip(tensor_list, path_list):
|
|
self._swap_out_tensor(tensor, swap_path)
|
|
|
|
def _report_statistics(self, message):
|
|
if dist.get_rank() == 0:
|
|
element_size = torch.tensor([], dtype=self.dtype).element_size()
|
|
swapped_GB = (self.num_elements_swapped * element_size) / (1024**3)
|
|
logger.debug(f'{message} num_elems = {self.num_elements_swapped}, {swapped_GB:5.2f} GB')
|
|
|
|
def _swap_out_tensor(self, tensor, swap_path):
|
|
assert len(self.all_buffers) > 0
|
|
|
|
aligned_numel = self._io_aligned_numel(tensor.numel())
|
|
assert aligned_numel <= self.max_numel
|
|
|
|
self._make_swap_space(aligned_numel)
|
|
assert self.current_buffer_index != INVALID_BUFFER_INDEX
|
|
|
|
swap_buffer = self._get_current_buffer()
|
|
swap_buffer.insert_tensor(tensor, swap_path, aligned_numel)
|
|
|
|
def _make_swap_space(self, numel):
|
|
if self.current_buffer_index == INVALID_BUFFER_INDEX:
|
|
self._allocate_buffer()
|
|
return
|
|
|
|
if not self._get_current_buffer().has_space(numel):
|
|
if len(self.free_buffer_index) > 0:
|
|
self._flush_ready_buffers()
|
|
else:
|
|
self._flush_buffers_until_complete()
|
|
self._allocate_buffer()
|
|
|
|
def _io_aligned_numel(self, numel):
|
|
remainder = numel % self.numel_alignment
|
|
return numel if remainder == 0 else (numel + self.numel_alignment - remainder)
|
|
|
|
def _allocate_buffer(self):
|
|
assert self.current_buffer_index == INVALID_BUFFER_INDEX
|
|
assert len(self.all_buffers) > 0
|
|
assert len(self.free_buffer_index) > 0
|
|
self.current_buffer_index = self.free_buffer_index[-1]
|
|
self.free_buffer_index = self.free_buffer_index[:-1]
|
|
|
|
def _flush_ready_buffers(self):
|
|
if self.current_buffer_index != INVALID_BUFFER_INDEX:
|
|
self.ready_buffer_index.append(self.current_buffer_index)
|
|
self.current_buffer_index = INVALID_BUFFER_INDEX
|
|
|
|
self._swap_out_ready_buffers()
|
|
|
|
def _flush_buffers_until_complete(self):
|
|
self._flush_ready_buffers()
|
|
assert len(self.ready_buffer_index) == 0
|
|
|
|
self._wait_for_swap_complete()
|
|
assert len(self.swapping_buffer_index) == 0
|
|
assert len(self.free_buffer_index) == len(self.all_buffers)
|
|
|
|
def _swap_out_ready_buffers(self):
|
|
for buffer_index in self.ready_buffer_index:
|
|
buffer = self._get_buffer(buffer_index)
|
|
swap_tensors = buffer.get_swap_tensors()
|
|
swap_paths = buffer.get_swap_paths()
|
|
self.num_pending_swaps += len(swap_tensors)
|
|
swap_out_tensors(self.aio_handle, swap_tensors, swap_paths)
|
|
|
|
self.swapping_buffer_index += self.ready_buffer_index
|
|
self.ready_buffer_index = []
|
|
|
|
def _wait_for_swap_complete(self):
|
|
assert len(self.swapping_buffer_index) > 0
|
|
|
|
self._start_timer(ASYNC_SWAPPER_WAIT_TIMER)
|
|
assert self.aio_handle.wait() == self.num_pending_swaps
|
|
self._stop_timer(ASYNC_SWAPPER_WAIT_TIMER)
|
|
self.timer_names.add(ASYNC_SWAPPER_WAIT_TIMER)
|
|
|
|
self.num_pending_swaps = 0
|
|
|
|
for buffer_index in self.swapping_buffer_index:
|
|
buffer = self._get_buffer(buffer_index)
|
|
self.num_elements_swapped += buffer.get_num_elem()
|
|
buffer.reset()
|
|
|
|
self.free_buffer_index += self.swapping_buffer_index
|
|
assert len(self.free_buffer_index) <= len(self.all_buffers)
|
|
self.swapping_buffer_index = []
|
|
|
|
def _get_buffer(self, index):
|
|
assert index != INVALID_BUFFER_INDEX
|
|
return self.all_buffers[index]
|
|
|
|
def _get_current_buffer(self):
|
|
return self._get_buffer(self.current_buffer_index)
|
|
|
|
def _start_timer(self, name):
|
|
if self.timers:
|
|
self.timers(name).start()
|
|
|
|
def _stop_timer(self, name):
|
|
if self.timers:
|
|
self.timers(name).stop()
|
|
|
|
def _log_timers(self, name_list, force=False):
|
|
if self.timers and force:
|
|
self.timers.log(name_list)
|