def _swap_out_unpinned_tensors(self,
aio_handle,
unpinned_tensors,
dest_paths,
pinned_buffers):
swap_buffer_count = len(pinned_buffers)
unpinned_tensor_count = len(unpinned_tensors)
for i in range(0, unpinned_tensor_count, swap_buffer_count):
swap_tensor_count = min((unpinned_tensor_count - i), swap_buffer_count)
src_tensors = unpinned_tensors[i:(i + swap_tensor_count)]
compute_lengths = [t.numel() for t in src_tensors]
compute_buffers = get_sized_buffers(pinned_buffers, compute_lengths)
for dst, src in zip(compute_buffers, src_tensors):
dst.data.copy_(src.data)
swap_lengths = [self._io_aligned_numel(t.numel()) for t in src_tensors]
swap_buffers = get_sized_buffers(pinned_buffers, swap_lengths)
swap_paths = dest_paths[i:(i + swap_tensor_count)]
swap_out_tensors(aio_handle, swap_buffers, swap_paths)
assert aio_handle.wait() == swap_tensor_count
def _swap_in_parameter(self, aio_handle, parameter, dest_buffers):
swap_info = self._get_param_swap_info(parameter)
if swap_info is None:
return
assert len(swap_info.tensors) <= len(dest_buffers)
swap_lengths = [self._io_aligned_numel(swap_info.numel())] * len(
swap_info.tensors)
swap_buffers = get_sized_buffers(dest_buffers, swap_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.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)
compute_lengths = [swap_info.numel()] * len(swap_info.tensors)
compute_buffers = get_sized_buffers(dest_buffers, compute_lengths)
for t, buffer in zip(swap_info.tensors, compute_buffers):
t.data = buffer.data
self._log_timers([READ_TIMER, WAIT_TIMER])
if DEBUG_MODE and torch.distributed.get_rank() == 0:
logger.info(
f'optimizer_param_swap_in: {(swap_bytes/(1024**3)):5.2f} GB')
def set_swap_buffers(self, buffers):
compute_lengths = [self.numel()] * len(self.tensors)
compute_buffers = get_sized_buffers(buffers, compute_lengths)
for t, buffer in zip(self.tensors, compute_buffers):
t.data = buffer.data