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(aio_config[AIO_BLOCK_SIZE],
aio_config[AIO_QUEUE_DEPTH],
aio_config[AIO_SINGLE_SUBMIT],
aio_config[AIO_OVERLAP_EVENTS],
aio_config[AIO_THREAD_COUNT])
# 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 torch.distributed.get_rank() == 0:
print_object(obj=self,
name='PartitionedOptimizerSwapper',
exclude_list=self.print_exclude_list)
def __init__(self, swap_config, aio_config, base_folder, optimizer,
largest_numel, device, dtype, timers):
super(PipelinedOptimizerSwapper,
self).__init__(swap_config, aio_config, base_folder, optimizer,
largest_numel, device, dtype, timers)
aio_op = AsyncIOBuilder().load()
self.write_aio_handle = aio_op.aio_handle(
aio_config[AIO_BLOCK_SIZE], aio_config[AIO_QUEUE_DEPTH],
aio_config[AIO_SINGLE_SUBMIT], aio_config[AIO_OVERLAP_EVENTS],
aio_config[AIO_THREAD_COUNT])
self.read_aio_handle = aio_op.aio_handle(
aio_config[AIO_BLOCK_SIZE], aio_config[AIO_QUEUE_DEPTH],
aio_config[AIO_SINGLE_SUBMIT], aio_config[AIO_OVERLAP_EVENTS],
aio_config[AIO_THREAD_COUNT])
# Overlap gradient swap out
self.gradient_swapper = AsyncTensorSwapper(
aio_handle=self.write_aio_handle,
numel_alignment=self.numel_alignment,
timers=self.timers)
self.async_swap_in = swap_config[OFFLOAD_OPTIMIZER_PIPELINE_READ]
self.async_swap_out = swap_config[OFFLOAD_OPTIMIZER_PIPELINE_WRITE]
self.swap_ops = {
SYNC_SWAP_IN: None,
ASYNC_SWAP_IN: None,
SYNC_SWAP_OUT: None,
ASYNC_SWAP_OUT: None
}
self.print_exclude_list += [
'gradient_swapper', 'read_aio_handle', 'write_aio_handle',
'swap_ops', 'print_exclude_list'
]
if torch.distributed.get_rank() == 0:
print_object(obj=self,
name='PipelinedOptimizerSwapper',
exclude_list=self.print_exclude_list)