wehub-resource-sync
1143 lines
43 KiB
Python
1143 lines
43 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
"""
|
|
Use to partition the activations stored for backward propagation
|
|
Therefore reduces the memory consumption
|
|
Also implements CPU checkpointing and contiguous memory checkpointing
|
|
Reduces memory consumption and memory fragmentation
|
|
|
|
Code for rng checkpointing taken from NVIDIA Megatron-LM mpu/random.py
|
|
b886b7bb972afe72bac0f5de4f42a4a7bae8ebef
|
|
"""
|
|
|
|
# Parts of the code here are adapted from PyTorch
|
|
# repo: https://github.com/pytorch/pytorch
|
|
import copy
|
|
import torch
|
|
import contextlib
|
|
from deepspeed import comm as dist
|
|
import weakref
|
|
|
|
import mmap
|
|
from torch import _C
|
|
|
|
from deepspeed.runtime.config import DeepSpeedConfig
|
|
from deepspeed.utils import logger
|
|
from deepspeed.runtime.utils import copy_to_device, move_to_device, see_memory_usage
|
|
from deepspeed.utils.timer import SynchronizedWallClockTimer as Timers, FORWARD_GLOBAL_TIMER
|
|
from deepspeed.utils.bwc import bwc_tensor_model_parallel_rank
|
|
from deepspeed.accelerator import get_accelerator
|
|
from deepspeed.runtime import compiler
|
|
|
|
# DeepSpeed Checkpointing Enabled or Disabled
|
|
deepspeed_checkpointing_enabled = False
|
|
|
|
# MP parameters
|
|
mpu = None
|
|
|
|
#set default values
|
|
mp_rank = 0
|
|
mp_size = 1
|
|
mp_group = None
|
|
|
|
# Model Parameters
|
|
num_layers = None
|
|
|
|
# Checkpointing buffers
|
|
contiguous_data_buffers = []
|
|
data_offsets = []
|
|
|
|
contiguous_size_buffers = []
|
|
size_offsets = []
|
|
|
|
timers = None
|
|
|
|
# optimization flags
|
|
PARTITION_ACTIVATIONS = False
|
|
CPU_CHECKPOINT = False
|
|
CONTIGUOUS_CHECKPOINTING = False
|
|
SYNCHRONIZE = False
|
|
PROFILE_TIME = False
|
|
|
|
# Default name for the model parallel rng tracker.
|
|
_MODEL_PARALLEL_RNG_TRACKER_NAME = 'model-parallel-rng'
|
|
|
|
|
|
def detach_variable(inputs, device=None):
|
|
if isinstance(inputs, tuple):
|
|
out = []
|
|
for inp in inputs:
|
|
if not isinstance(inp, torch.Tensor):
|
|
out.append(inp)
|
|
continue
|
|
|
|
requires_grad = inp.requires_grad
|
|
|
|
if device is not None:
|
|
x = inp.to(device=device)
|
|
else:
|
|
x = inp
|
|
|
|
x = x.detach()
|
|
x.requires_grad = requires_grad
|
|
out.append(x)
|
|
return tuple(out)
|
|
else:
|
|
raise RuntimeError("Only tuple of tensors is supported. Got Unsupported input type: ", type(inputs).__name__)
|
|
|
|
|
|
def _set_cuda_rng_state(new_state, device=-1):
|
|
"""Sets the random number generator state of the current GPU.
|
|
|
|
Arguments:
|
|
new_state (torch.ByteTensor): The desired state
|
|
This function is adapted from PyTorch repo (torch.cuda.set_rng_state) #ignore-cuda
|
|
with a single change: the input state is not cloned. Cloning caused
|
|
major performance issues for +4 GPU cases.
|
|
"""
|
|
if hasattr(_C, '_cuda_setRNGState') and callable(_C._cuda_setRNGState):
|
|
# older PyTorch
|
|
def cb():
|
|
with get_accelerator().device(device):
|
|
_C._cuda_setRNGState(new_state)
|
|
else:
|
|
# newer PyTorch
|
|
if device == -1:
|
|
device = torch.device(get_accelerator().device_name())
|
|
elif isinstance(device, str):
|
|
device = torch.device(device)
|
|
elif isinstance(device, int):
|
|
device = torch.device(get_accelerator().device_name(), device)
|
|
|
|
def cb():
|
|
idx = device.index
|
|
if idx is None:
|
|
idx = get_accelerator().current_device()
|
|
default_generator = get_accelerator().default_generator(idx)
|
|
default_generator.set_state(new_state)
|
|
|
|
get_accelerator().lazy_call(cb)
|
|
|
|
|
|
class CudaRNGStatesTracker:
|
|
"""Tracker for the cuda RNG states.
|
|
|
|
Using the `add` method, a cuda rng state is initialized based on
|
|
the input `seed` and is assigned to `name`. Later, by forking the
|
|
rng state, we can perform operations and return to our starting
|
|
cuda state.
|
|
"""
|
|
|
|
def __init__(self):
|
|
# Map from a string name to the cuda rng state.
|
|
self.states_ = {}
|
|
# Seeds are just for book keeping and ensure no seed is set twice.
|
|
self.seeds_ = set()
|
|
|
|
def reset(self):
|
|
"""Set to the initial state (no tracker)."""
|
|
self.states_ = {}
|
|
self.seeds_ = set()
|
|
|
|
def get_states(self):
|
|
"""Get rng states. Copy the dictionary so we have direct
|
|
pointers to the states, not just a pointer to the dictionary."""
|
|
return copy.copy(self.states_)
|
|
|
|
def set_states(self, states):
|
|
"""Set the rng states. For efficiency purposes, we do not check
|
|
the size of seed for compatibility."""
|
|
self.states_ = states
|
|
|
|
def add(self, name, seed):
|
|
"""Track the rng state."""
|
|
# Check seed is not already used.
|
|
if seed in self.seeds_:
|
|
raise Exception('seed {} already exists'.format(seed))
|
|
self.seeds_.add(seed)
|
|
# Check that state is not already defined.
|
|
if name in self.states_:
|
|
raise Exception('cuda rng state {} already exists'.format(name))
|
|
# Get the current rng state.
|
|
orig_rng_state = get_accelerator().get_rng_state()
|
|
# Set the new state and store it.
|
|
get_accelerator().manual_seed(seed)
|
|
self.states_[name] = get_accelerator().get_rng_state()
|
|
# Reset rng state to what it was.
|
|
_set_cuda_rng_state(orig_rng_state)
|
|
|
|
@contextlib.contextmanager
|
|
def fork(self, name=_MODEL_PARALLEL_RNG_TRACKER_NAME):
|
|
"""Fork the cuda rng state, perform operations, and exit with
|
|
the original state."""
|
|
# Check if we have added the state
|
|
if name not in self.states_:
|
|
raise Exception('cuda rng state {} is not added'.format(name))
|
|
# Store current rng state.
|
|
orig_cuda_rng_state = get_accelerator().get_rng_state()
|
|
# Set rng state to the desired one
|
|
_set_cuda_rng_state(self.states_[name])
|
|
# Do the stuff we wanted to do.
|
|
try:
|
|
yield
|
|
finally:
|
|
# Update the current rng state for later use.
|
|
self.states_[name] = get_accelerator().get_rng_state()
|
|
# And set the state to the original state we started with.
|
|
_set_cuda_rng_state(orig_cuda_rng_state)
|
|
|
|
|
|
# RNG tracker object.
|
|
_CUDA_RNG_STATE_TRACKER = CudaRNGStatesTracker()
|
|
|
|
|
|
def get_cuda_rng_tracker():
|
|
"""Get cuda rng tracker."""
|
|
return _CUDA_RNG_STATE_TRACKER
|
|
|
|
|
|
def model_parallel_cuda_manual_seed(seed):
|
|
"""Initialize model parallel cuda seed.
|
|
|
|
This function should be called after the model parallel is
|
|
initialized. Also, no get_accelerator().manual_seed should be called
|
|
after this function. Basically, this is replacement for that
|
|
function.
|
|
Two set of RNG states are tracked:
|
|
default state: This is for data parallelism and is the same among a
|
|
set of model parallel GPUs but different across
|
|
different model parallel groups. This is used for
|
|
example for dropout in the non-model-parallel regions.
|
|
model-parallel state: This state is different among a set of model
|
|
parallel GPUs, but the same across data parallel
|
|
groups. This is used for example for dropout in
|
|
model parallel regions.
|
|
"""
|
|
global mpu
|
|
|
|
tp_rank = bwc_tensor_model_parallel_rank(mpu)
|
|
|
|
# 2718 is just for fun and any POSITIVE value will work.
|
|
offset = seed + 2718
|
|
model_parallel_seed = offset + tp_rank
|
|
# Data parallel gets the original seed.
|
|
data_parallel_seed = seed
|
|
|
|
if dist.get_rank() == 0:
|
|
logger.info(
|
|
'> initializing model parallel cuda seeds on global rank {}, '
|
|
'model parallel rank {}, and data parallel rank {} with '
|
|
'model parallel seed: {} and data parallel seed: {}'.format(dist.get_rank(), tp_rank,
|
|
mpu.get_data_parallel_rank(),
|
|
model_parallel_seed, data_parallel_seed), )
|
|
_CUDA_RNG_STATE_TRACKER.reset()
|
|
# Set the default state.
|
|
get_accelerator().manual_seed(data_parallel_seed)
|
|
# and model parallel state.
|
|
_CUDA_RNG_STATE_TRACKER.add(_MODEL_PARALLEL_RNG_TRACKER_NAME, model_parallel_seed)
|
|
|
|
|
|
def model_parallel_reconfigure_tp_seed(seed):
|
|
global mpu
|
|
tp_rank = bwc_tensor_model_parallel_rank(mpu)
|
|
model_parallel_seed = seed + 2718 + tp_rank
|
|
with _CUDA_RNG_STATE_TRACKER.fork():
|
|
get_accelerator().manual_seed(model_parallel_seed)
|
|
|
|
|
|
def get_partition_start(item):
|
|
global mp_rank, mp_size, mp_group
|
|
size = item.numel()
|
|
partition_size = size / mp_size
|
|
start = partition_size * mp_rank
|
|
return int(start)
|
|
|
|
|
|
def get_partition_size(item):
|
|
global mp_rank, mp_size, mp_group
|
|
size = item.numel()
|
|
assert size % mp_size == 0, "Doesn't handle if partition activation if item is not divisible by mp size"
|
|
partition_size = size / mp_size
|
|
return int(partition_size)
|
|
|
|
|
|
def gather_partitioned_activations(tensors, device=None):
|
|
global mp_rank, mp_size, mp_group
|
|
assert len(tensors) % 2 == 0, f'Expected even count of tensors, instead got {len(tensors)}'
|
|
inputs = []
|
|
num_args = int(len(tensors) / 2)
|
|
for i in range(num_args):
|
|
|
|
item = tensors[2 * i]
|
|
size = tensors[2 * i + 1]
|
|
|
|
if not is_activation_to_checkpoint(item):
|
|
inputs.append(item)
|
|
continue
|
|
|
|
# don't need to do all_gather if model parallel is not enabled
|
|
if mp_group is None or mp_size == 1:
|
|
item = item.view(list(size.numpy()))
|
|
if device is not None:
|
|
item = item.to(device)
|
|
inputs.append(item)
|
|
continue
|
|
|
|
partition_size = item.numel()
|
|
tensor_size = partition_size * mp_size
|
|
if device is not None:
|
|
flat_tensor = torch.zeros([tensor_size], dtype=item.dtype, device=device)
|
|
else:
|
|
flat_tensor = torch.zeros([tensor_size], dtype=item.dtype, device=item.device)
|
|
part = flat_tensor.narrow(0, partition_size * mp_rank, partition_size)
|
|
part.copy_(item)
|
|
dist.all_gather_into_tensor(flat_tensor, part, group=mp_group)
|
|
input_tensor = flat_tensor.view(list(size.numpy()))
|
|
item.data = input_tensor.data
|
|
|
|
inputs.append(item)
|
|
|
|
return tuple(inputs)
|
|
|
|
|
|
def extract_tensors(all_objects):
|
|
"""
|
|
Separate objects in list/tuple into tensors and non-tensors and create a mapping to enable re-aggregation.
|
|
The order of tensors and non-tensors is preserved in their respective output groups.
|
|
|
|
Parameters:
|
|
all_objects (list/tuple): Objects containing tensors and non-tensors to be split.
|
|
|
|
Returns:
|
|
tuple: Containing tensors, non-tensors, and bools of whether each position in original list/tuple was a tensor.
|
|
|
|
"""
|
|
tensor_objects = [v for v in all_objects if torch.is_tensor(v)]
|
|
non_tensor_objects = [v for v in all_objects if not torch.is_tensor(v)]
|
|
tensor_flags = [torch.is_tensor(v) for v in all_objects]
|
|
if type(all_objects) is tuple:
|
|
return tuple(tensor_objects), tuple(non_tensor_objects), tuple(tensor_flags)
|
|
return tensor_objects, non_tensor_objects, tensor_flags
|
|
|
|
|
|
def merge_tensors(tensor_objects, non_tensor_objects, tensor_flags):
|
|
"""
|
|
Merge two lists (or tuples) of tensors and non-tensors using a mapping of positions in merged list (or tuple).
|
|
|
|
Parameters:
|
|
tensor_objects (list/tuple): Tensors to merge.
|
|
non_tensor_objects (list/tuple): Non-tensors to merge.
|
|
tensor_flags (list/tuple): Indicates whether each position in output is a tensor.
|
|
|
|
Returns:
|
|
tuple: Merge of tensors and non-tensors
|
|
"""
|
|
merged_objects = []
|
|
tensor_idx = 0
|
|
non_tensor_idx = 0
|
|
|
|
real_tensor_flags = None
|
|
|
|
# remove the flags that are assigned to the size of the flattened tensors
|
|
if PARTITION_ACTIVATIONS:
|
|
real_tensor_flags = []
|
|
previous_flag = False
|
|
for flag in tensor_flags:
|
|
if previous_flag:
|
|
previous_flag = False
|
|
continue
|
|
previous_flag = flag
|
|
real_tensor_flags.append(flag)
|
|
else:
|
|
real_tensor_flags = tensor_flags
|
|
|
|
for is_tensor in real_tensor_flags:
|
|
if is_tensor:
|
|
merged_objects.append(tensor_objects[tensor_idx])
|
|
tensor_idx += 1
|
|
else:
|
|
merged_objects.append(non_tensor_objects[non_tensor_idx])
|
|
non_tensor_idx += 1
|
|
|
|
return tuple(merged_objects)
|
|
|
|
|
|
def is_activation_to_checkpoint(item):
|
|
"""
|
|
Is an activation to be checkpointed
|
|
"""
|
|
global mp_size
|
|
extra_flag = (not hasattr(item, 'no_checkpointing')) or (hasattr(item, 'no_checkpointing')
|
|
and item.no_checkpointing == False)
|
|
return torch.is_tensor(item) and item.is_floating_point() and item.numel() >= mp_size and extra_flag
|
|
|
|
|
|
def partition_activations(args, cpu_checkpoint, contiguous_checkpoint):
|
|
global contiguous_data_buffers, data_offsets
|
|
|
|
inputs = []
|
|
num_non_fp_tensors = 0
|
|
|
|
for arg_index, item in enumerate(args):
|
|
if not is_activation_to_checkpoint(item):
|
|
inputs.append(item)
|
|
num_non_fp_tensors += 1
|
|
continue
|
|
|
|
i = arg_index - num_non_fp_tensors
|
|
partition_size = get_partition_size(item)
|
|
partition = item.detach().contiguous().view(-1).narrow(0, get_partition_start(item), partition_size).clone()
|
|
|
|
buffer_device = torch.device('cpu') if cpu_checkpoint else partition.device
|
|
|
|
if contiguous_checkpoint:
|
|
if i >= len(contiguous_data_buffers):
|
|
tensor_list = [
|
|
torch.tensor(()).new_empty([partition_size], dtype=partition.dtype, device=buffer_device)
|
|
for _ in range(num_layers)
|
|
]
|
|
contiguous_data_buffers.append(tensor_list)
|
|
data_offsets.append(0)
|
|
elif contiguous_data_buffers[i] is None:
|
|
tensor_list = [
|
|
torch.tensor(()).new_empty([partition_size], dtype=partition.dtype, device=buffer_device)
|
|
for _ in range(num_layers)
|
|
]
|
|
contiguous_data_buffers[i] = tensor_list
|
|
data_offsets[i] = 0
|
|
|
|
# Because the 'new_empty' returns uninitialized pages,
|
|
# the pages need to be populated during the cudaMemcpy time
|
|
# which increases the data copy time. To avoid this, we
|
|
# pre-populate these pages by simply writing 0 ahead of
|
|
# the actual cudaMemcpy operation time. Due to the
|
|
# previously launched GPU kernels, there is a small
|
|
# window of time here for CPUs to populate pages asynchronously.
|
|
contiguous_data_buffers[i][data_offsets[i]].data[range(
|
|
0, contiguous_data_buffers[i][data_offsets[i]].data.shape[0],
|
|
int(mmap.PAGESIZE / contiguous_data_buffers[i][data_offsets[i]].data.element_size()))] = 0
|
|
|
|
contiguous_partition = contiguous_data_buffers[i][data_offsets[i]].data.copy_(partition.data)
|
|
data_offsets[i] = data_offsets[i] + 1
|
|
inputs.append(contiguous_partition)
|
|
else:
|
|
partition = partition.cpu() if CPU_CHECKPOINT else partition
|
|
inputs.append(partition)
|
|
|
|
return inputs
|
|
|
|
|
|
def get_partitioned_activations_for_backward(args, inputs, contiguous_checkpoint):
|
|
global contiguous_size_buffers, size_offsets
|
|
|
|
new_args = []
|
|
num_non_fp_tensors = 0
|
|
|
|
for arg_index, (arg, inp) in enumerate(zip(args, inputs)):
|
|
size = torch.tensor(arg.size()) if torch.is_tensor(arg) else None
|
|
if not is_activation_to_checkpoint(arg):
|
|
new_args.append(arg)
|
|
new_args.append(size)
|
|
num_non_fp_tensors += 1
|
|
continue
|
|
|
|
arg.data = torch.empty([], device=arg.device).data
|
|
arg.saved_data = inp.data
|
|
|
|
new_args.append(arg)
|
|
i = arg_index - num_non_fp_tensors
|
|
|
|
if contiguous_checkpoint:
|
|
numel = size.numel()
|
|
if i >= len(contiguous_size_buffers):
|
|
tmp = torch.tensor(())
|
|
contiguous_size_buffers.append(
|
|
tmp.new_empty([numel * num_layers], dtype=size.dtype, device=size.device))
|
|
size_offsets.append(0)
|
|
elif contiguous_size_buffers[i] is None:
|
|
tmp = torch.tensor(())
|
|
contiguous_size_buffers[i] = tmp.new_empty([numel * num_layers], dtype=size.dtype, device=size.device)
|
|
size_offsets[i] = 0
|
|
|
|
contiguous_size = contiguous_size_buffers[i].narrow(0, size_offsets[i], numel).data.copy_(size.data)
|
|
contiguous_size = contiguous_size.view_as(size)
|
|
size_offsets[i] = size_offsets[i] + numel
|
|
new_args.append(contiguous_size)
|
|
else:
|
|
new_args.append(size)
|
|
|
|
return new_args
|
|
|
|
|
|
def get_cpu_activations_for_backward(args, inputs):
|
|
new_args = []
|
|
for i, (arg, inp) in enumerate(zip(args, inputs)):
|
|
if not is_activation_to_checkpoint(arg):
|
|
new_args.append(arg)
|
|
continue
|
|
|
|
arg.data = torch.empty([], device=arg.device).data
|
|
arg.saved_data = inp.data
|
|
new_args.append(arg)
|
|
|
|
return new_args
|
|
|
|
|
|
class CheckpointFunction(torch.autograd.Function):
|
|
"""This function is adapted from torch.utils.checkpoint with
|
|
two main changes:
|
|
1) torch.cuda.set_rng_state is replaced with `_set_cuda_rng_state` #ignore-cuda
|
|
2) the states in the model parallel tracker are also properly
|
|
tracked/set/reset.
|
|
3) Performance activation partitioning, contiguous memory optimization
|
|
4) CPU Checkpointing
|
|
5) Profile forward and backward functions
|
|
"""
|
|
|
|
@staticmethod
|
|
def forward(ctx, run_function, all_outputs, *args):
|
|
global mpu, timers, SYNCHRONIZE, PROFILE_TIME
|
|
|
|
def save_args_for_backward(*all_args):
|
|
tensor_args, non_tensor_args, tensor_flags = extract_tensors(all_objects=all_args)
|
|
ctx.deepspeed_saved_tensors = tensor_args
|
|
ctx.non_tensor_args = non_tensor_args
|
|
ctx.tensor_flags = tensor_flags
|
|
|
|
if SYNCHRONIZE:
|
|
get_accelerator().synchronize()
|
|
|
|
if timers is None and PROFILE_TIME:
|
|
timers = Timers()
|
|
|
|
if PROFILE_TIME:
|
|
timers(FORWARD_GLOBAL_TIMER).start()
|
|
|
|
ctx.run_function = run_function
|
|
global num_layers
|
|
global mp_rank, mp_size, mp_group
|
|
global contiguous_data_buffers, contiguous_size_buffers
|
|
global data_offsets, size_offsets
|
|
global PARTITION_ACTIVATIONS, buffer_0, buffer_1, buffer_0_offset, buffer_1_offset
|
|
|
|
cuda_device = get_accelerator().current_device_name()
|
|
transport_stream = get_accelerator().Stream(device=cuda_device)
|
|
|
|
if PARTITION_ACTIVATIONS:
|
|
inputs = partition_activations(args, CPU_CHECKPOINT, CONTIGUOUS_CHECKPOINTING)
|
|
elif CPU_CHECKPOINT:
|
|
inputs = copy_to_device(args, device=torch.device('cpu'), criterion_func=is_activation_to_checkpoint)
|
|
|
|
# just in case something funky is happening such as reuse of inputs
|
|
inputs_cuda = copy_to_device(args, device=cuda_device, criterion_func=is_activation_to_checkpoint)
|
|
|
|
# Copy the rng states.
|
|
ctx.fwd_cpu_rng_state = torch.get_rng_state()
|
|
ctx.fwd_cuda_rng_state = get_accelerator().get_rng_state()
|
|
ctx.fwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states()
|
|
|
|
see_memory_usage("Before running forward on the layer", force=False)
|
|
# ctx.save_for_backward(*args)
|
|
with torch.no_grad():
|
|
outputs = run_function(*inputs_cuda)
|
|
|
|
see_memory_usage("After running forward on the layer", force=False)
|
|
del inputs_cuda
|
|
|
|
if PARTITION_ACTIVATIONS:
|
|
new_args = get_partitioned_activations_for_backward(args, inputs, CONTIGUOUS_CHECKPOINTING)
|
|
assert len(new_args) % 2 == 0, f'save_for_backward called with odd number of args, {len(new_args)}'
|
|
save_args_for_backward(*new_args)
|
|
elif CPU_CHECKPOINT:
|
|
new_args = get_cpu_activations_for_backward(args, inputs)
|
|
save_args_for_backward(*new_args)
|
|
else:
|
|
save_args_for_backward(*args)
|
|
|
|
if PROFILE_TIME:
|
|
timers(FORWARD_GLOBAL_TIMER).stop()
|
|
timers.log([FORWARD_GLOBAL_TIMER])
|
|
if SYNCHRONIZE:
|
|
get_accelerator().synchronize()
|
|
|
|
# Tensors returned from forward() may not be differentiable.
|
|
if torch.is_tensor(outputs):
|
|
non_grad_outputs = [outputs] if not outputs.is_floating_point() else []
|
|
else:
|
|
non_grad_outputs = [o for o in outputs if torch.is_tensor(o) and not o.is_floating_point()]
|
|
ctx.mark_non_differentiable(*non_grad_outputs)
|
|
|
|
if torch.is_tensor(outputs):
|
|
all_outputs += [outputs]
|
|
return outputs
|
|
else:
|
|
all_outputs += outputs
|
|
outputs, _, _ = extract_tensors(all_objects=outputs)
|
|
return tuple(outputs)
|
|
|
|
@staticmethod
|
|
def backward(ctx, *grads):
|
|
global timers
|
|
see_memory_usage("In backward", force=False)
|
|
# removing pointers to the contiguous buffer memory
|
|
# so that they can be garbage collected once the checkpoints
|
|
# have been used
|
|
if SYNCHRONIZE:
|
|
get_accelerator().synchronize()
|
|
if PROFILE_TIME:
|
|
timers('backward').start()
|
|
|
|
if CONTIGUOUS_CHECKPOINTING:
|
|
global data_offsets, size_offsets
|
|
global contiguous_data_buffers, contiguous_size_buffers
|
|
|
|
for buffers in contiguous_data_buffers:
|
|
buffers = []
|
|
|
|
# frees up all the pointers to the checkpoints except for the ones
|
|
# stored by save for backward
|
|
contiguous_data_buffers = []
|
|
contiguous_size_buffers = []
|
|
data_offsets = []
|
|
size_offsets = []
|
|
|
|
see_memory_usage("In backward checkpointing code", force=False)
|
|
if not torch.autograd._is_checkpoint_valid():
|
|
raise RuntimeError("Checkpointing is not compatible with .grad(), "
|
|
"please use .backward() if possible")
|
|
|
|
global PARTITION_ACTIVATIONS
|
|
cuda_device = get_accelerator().current_device_name()
|
|
transport_stream = get_accelerator().Stream(device=cuda_device)
|
|
# Rebuild deepspeed_saved_tensors
|
|
for t in ctx.deepspeed_saved_tensors:
|
|
if t is not None and hasattr(t, 'saved_data') and t.saved_data is not None:
|
|
t.data = t.saved_data.to(t.device)
|
|
t.saved_data = None
|
|
|
|
if PARTITION_ACTIVATIONS:
|
|
# with get_accelerator().stream(transport_stream):
|
|
inputs = gather_partitioned_activations(ctx.deepspeed_saved_tensors,
|
|
device=cuda_device if CPU_CHECKPOINT else None)
|
|
detached_inputs = detach_variable(inputs)
|
|
elif CPU_CHECKPOINT:
|
|
inputs = move_to_device(ctx.deepspeed_saved_tensors, cuda_device, is_activation_to_checkpoint)
|
|
detached_inputs = detach_variable(inputs)
|
|
else:
|
|
inputs = ctx.deepspeed_saved_tensors
|
|
detached_inputs = detach_variable(inputs)
|
|
|
|
# Add non tensor input args
|
|
detached_inputs = merge_tensors(tensor_objects=detached_inputs,
|
|
non_tensor_objects=ctx.non_tensor_args,
|
|
tensor_flags=ctx.tensor_flags)
|
|
|
|
# Store the current states.
|
|
bwd_cpu_rng_state = torch.get_rng_state()
|
|
bwd_cuda_rng_state = get_accelerator().get_rng_state()
|
|
bwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states()
|
|
|
|
# Set the states to what it used to be before the forward pass.
|
|
torch.set_rng_state(ctx.fwd_cpu_rng_state)
|
|
_set_cuda_rng_state(ctx.fwd_cuda_rng_state)
|
|
get_cuda_rng_tracker().set_states(ctx.fwd_cuda_rng_state_tracker)
|
|
|
|
# if PARTITION_ACTIVATIONS:
|
|
# current_stream=get_accelerator().current_stream()
|
|
# current_stream.wait_stream(transport_stream)
|
|
|
|
see_memory_usage("In backward checkpointing code before forward", force=False)
|
|
|
|
with torch.enable_grad():
|
|
outputs = ctx.run_function(*detached_inputs)
|
|
|
|
see_memory_usage("In backward checkpointing code after forward", force=False)
|
|
# Set the states back to what it was at the start of this function.
|
|
torch.set_rng_state(bwd_cpu_rng_state)
|
|
_set_cuda_rng_state(bwd_cuda_rng_state)
|
|
get_cuda_rng_tracker().set_states(bwd_cuda_rng_state_tracker)
|
|
|
|
if isinstance(outputs, torch.Tensor):
|
|
outputs = (outputs, )
|
|
|
|
# Filter out non tensor outputs
|
|
outputs, _, _ = extract_tensors(all_objects=outputs)
|
|
|
|
# Construct arguments to autograd.backward().
|
|
# This is usually just outputs and grads, but forward() can return tensors that
|
|
# are not differentiable.
|
|
output_tensors = []
|
|
grad_tensors = []
|
|
for out, grad in zip(outputs, grads):
|
|
if out.requires_grad:
|
|
output_tensors.append(out)
|
|
grad_tensors.append(grad)
|
|
|
|
see_memory_usage("In backward checkpointing code before backward", force=False)
|
|
|
|
torch.autograd.backward(output_tensors, grad_tensors)
|
|
|
|
# Force clear our stashed tensors to prevent a memory leak in certain scenarios
|
|
ctx.deepspeed_saved_tensors = None
|
|
ctx.non_tensor_args = None
|
|
ctx.tensor_flags = None
|
|
|
|
see_memory_usage("After backward checkpointing code after backward", force=False)
|
|
|
|
if PROFILE_TIME:
|
|
timers('backward').stop()
|
|
timers.log(['backward'])
|
|
if SYNCHRONIZE:
|
|
get_accelerator().synchronize()
|
|
ret_list = [None, None] # first None for ctx
|
|
for inp in detached_inputs:
|
|
if torch.is_tensor(inp):
|
|
ret_list.append(inp.grad)
|
|
else:
|
|
ret_list.append(None)
|
|
|
|
return tuple(ret_list)
|
|
|
|
|
|
def non_reentrant_checkpoint(function, *args):
|
|
"""This function is union of `torch.utils.checkpoint._checkpoint_without_reentrant` and `CheckpointFunction` in this module
|
|
|
|
This function is aim to solve the back propagation error raised from all input requires no grad.
|
|
* has already been implemented in pytorch for a while, the solution is stable at most time except for jit module mode.
|
|
* can help to solve the issue which is hacked by `deepspeed.runtime.pipe.module.PipelineModule._is_checkpointable`
|
|
|
|
Main modifications compared to the implementation of torch:
|
|
1. adapt to the signature of `checkpoint` function in this module
|
|
2. solve the non-deterministic by random state management consistent with deepspeed `CheckpointFunction`
|
|
3. when there is partition or cpu checkpointing, gather them in the unpack_hook during back propagation
|
|
4. make all after backward blocks in the hook which will executed after all leaf nodes backward execution.
|
|
5. above 4. is inspired by `torch.autograd.graph.register_multi_grad_hook`, which is only implemented after 2.0.0
|
|
"""
|
|
global mpu, timers, SYNCHRONIZE, PROFILE_TIME
|
|
|
|
deepspeed_saved_tensors = None
|
|
non_tensor_args = None
|
|
tensor_flags = None
|
|
|
|
def save_args_for_backward(*all_args):
|
|
"""keep this function to reduce the modification from original implementation"""
|
|
nonlocal deepspeed_saved_tensors, non_tensor_args, tensor_flags
|
|
tensor_args, non_tensor_args, tensor_flags = extract_tensors(all_objects=all_args)
|
|
deepspeed_saved_tensors = tensor_args
|
|
non_tensor_args = non_tensor_args
|
|
tensor_flags = tensor_flags
|
|
|
|
if SYNCHRONIZE:
|
|
get_accelerator().synchronize()
|
|
|
|
if timers is None and PROFILE_TIME:
|
|
timers = Timers()
|
|
|
|
if PROFILE_TIME:
|
|
timers(FORWARD_GLOBAL_TIMER).start()
|
|
|
|
global num_layers
|
|
global mp_rank, mp_size, mp_group
|
|
global contiguous_data_buffers, contiguous_size_buffers
|
|
global data_offsets, size_offsets
|
|
global PARTITION_ACTIVATIONS, buffer_0, buffer_1, buffer_0_offset, buffer_1_offset
|
|
|
|
cuda_device = get_accelerator().current_device_name()
|
|
transport_stream = get_accelerator().Stream(device=cuda_device)
|
|
|
|
if PARTITION_ACTIVATIONS:
|
|
inputs = partition_activations(args, CPU_CHECKPOINT, CONTIGUOUS_CHECKPOINTING)
|
|
elif CPU_CHECKPOINT:
|
|
inputs = copy_to_device(args, device=torch.device('cpu'), criterion_func=is_activation_to_checkpoint)
|
|
|
|
# just in case something funky is happening such as reuse of inputs
|
|
inputs_cuda = copy_to_device(args, device=cuda_device, criterion_func=is_activation_to_checkpoint)
|
|
|
|
# Copy the rng states.
|
|
fwd_cpu_rng_state = torch.get_rng_state()
|
|
fwd_cuda_rng_state = get_accelerator().get_rng_state()
|
|
fwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states()
|
|
|
|
if PARTITION_ACTIVATIONS:
|
|
new_args = get_partitioned_activations_for_backward(args, inputs, CONTIGUOUS_CHECKPOINTING)
|
|
assert len(new_args) % 2 == 0, f'save_for_backward called with odd number of args, {len(new_args)}'
|
|
save_args_for_backward(*new_args)
|
|
elif CPU_CHECKPOINT:
|
|
new_args = get_cpu_activations_for_backward(args, inputs)
|
|
save_args_for_backward(*new_args)
|
|
else:
|
|
save_args_for_backward(*args)
|
|
|
|
class Holder():
|
|
"""the place holder object used as activations to save memory"""
|
|
pass
|
|
|
|
# weakref seems utilized to discover the tensor deletion before a whole
|
|
# forward backward pair loop finished
|
|
storage: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()
|
|
weak_holder_list = []
|
|
leaf_tensors = []
|
|
backward_visited_leaf_nodes = 0
|
|
|
|
def checkpoint_pack(tensor_from_forward):
|
|
"""used to record the activation order in the `weak_holder_list`
|
|
|
|
the activation order in holder list is consistent between the first forward and recomputing forward.
|
|
* the jit compiled forward will break the order consistency *
|
|
"""
|
|
res = Holder()
|
|
weak_holder_list.append(weakref.ref(res))
|
|
|
|
# if this is a leaf tensor, save it for backward progression trace
|
|
# leaf tensor used to be input or parameters, which is not activations and
|
|
# has no memory overhead
|
|
if tensor_from_forward.requires_grad and tensor_from_forward.is_leaf:
|
|
leaf_tensors.append(tensor_from_forward)
|
|
return res
|
|
|
|
def checkpoint_unpack(holder_from_backward):
|
|
"""retrieve the activations from recompute"""
|
|
nonlocal deepspeed_saved_tensors, non_tensor_args, tensor_flags
|
|
|
|
# if this is the first step of backward propagation, recompute the graph and save
|
|
# all the activations with the same order as `checkpoint_pack` does
|
|
if len(storage) == 0:
|
|
unpack_counter = 0
|
|
|
|
def replay_pack(tensor_from_replay):
|
|
"""save recompute activations"""
|
|
nonlocal unpack_counter
|
|
unpack_counter += 1
|
|
|
|
if weak_holder_list[unpack_counter - 1]() is None:
|
|
return
|
|
|
|
detached_activations = tensor_from_replay.detach()
|
|
storage[weak_holder_list[unpack_counter - 1]()] = detached_activations
|
|
|
|
return
|
|
|
|
def replay_unpack(none_value):
|
|
"""recompute graph need not to backward"""
|
|
raise RuntimeError("You are calling backwards on a tensor that is never exposed.")
|
|
|
|
global timers
|
|
see_memory_usage("In backward", force=False)
|
|
# removing pointers to the contiguous buffer memory
|
|
# so that they can be garbage collected once the checkpoints
|
|
# have been used
|
|
if SYNCHRONIZE:
|
|
get_accelerator().synchronize()
|
|
if PROFILE_TIME:
|
|
timers('backward').start()
|
|
|
|
if CONTIGUOUS_CHECKPOINTING:
|
|
global data_offsets, size_offsets
|
|
global contiguous_data_buffers, contiguous_size_buffers
|
|
|
|
for buffers in contiguous_data_buffers:
|
|
buffers = []
|
|
|
|
# frees up all the pointers to the checkpoints except for the ones
|
|
# stored by save for backward
|
|
contiguous_data_buffers = []
|
|
contiguous_size_buffers = []
|
|
data_offsets = []
|
|
size_offsets = []
|
|
|
|
see_memory_usage("In backward checkpointing code", force=False)
|
|
if not torch.autograd._is_checkpoint_valid():
|
|
raise RuntimeError("Checkpointing is not compatible with .grad(), "
|
|
"please use .backward() if possible")
|
|
|
|
global PARTITION_ACTIVATIONS
|
|
cuda_device = get_accelerator().current_device_name()
|
|
transport_stream = get_accelerator().Stream(device=cuda_device)
|
|
|
|
# gather inputs which is partitioned or checkpointed before first forward
|
|
if PARTITION_ACTIVATIONS:
|
|
# with get_accelerator().stream(transport_stream):
|
|
inputs = gather_partitioned_activations(deepspeed_saved_tensors,
|
|
device=cuda_device if CPU_CHECKPOINT else None)
|
|
detached_inputs = detach_variable(inputs)
|
|
elif CPU_CHECKPOINT:
|
|
inputs = move_to_device(deepspeed_saved_tensors, cuda_device, is_activation_to_checkpoint)
|
|
detached_inputs = detach_variable(inputs)
|
|
else:
|
|
inputs = deepspeed_saved_tensors
|
|
detached_inputs = detach_variable(inputs)
|
|
|
|
# Add non tensor input args
|
|
detached_inputs = merge_tensors(tensor_objects=detached_inputs,
|
|
non_tensor_objects=non_tensor_args,
|
|
tensor_flags=tensor_flags)
|
|
|
|
# Store the current states.
|
|
bwd_cpu_rng_state = torch.get_rng_state()
|
|
bwd_cuda_rng_state = get_accelerator().get_rng_state()
|
|
bwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states()
|
|
|
|
# Set the states to what it used to be before the forward pass.
|
|
torch.set_rng_state(fwd_cpu_rng_state)
|
|
_set_cuda_rng_state(fwd_cuda_rng_state)
|
|
get_cuda_rng_tracker().set_states(fwd_cuda_rng_state_tracker)
|
|
|
|
see_memory_usage("In backward checkpointing code before forward", force=False)
|
|
with torch.enable_grad(), torch.autograd.graph.saved_tensors_hooks(replay_pack, replay_unpack):
|
|
_unused = function(*detached_inputs)
|
|
|
|
see_memory_usage("In backward checkpointing code after forward", force=False)
|
|
# Set the states back to what it was at the start of this function.
|
|
torch.set_rng_state(bwd_cpu_rng_state)
|
|
_set_cuda_rng_state(bwd_cuda_rng_state)
|
|
get_cuda_rng_tracker().set_states(bwd_cuda_rng_state_tracker)
|
|
|
|
deepspeed_saved_tensors = None
|
|
non_tensor_args = None
|
|
tensor_flags = None
|
|
|
|
if holder_from_backward not in storage:
|
|
raise RuntimeError("Attempt to retrieve a tensor saved by autograd multiple times without checkpoint"
|
|
" recomputation being triggered in between, this is not currently supported.")
|
|
|
|
return storage[holder_from_backward]
|
|
|
|
def after_backward_hook(_nonuse_grads):
|
|
"""the hook registered to all leaf tensors"""
|
|
nonlocal leaf_tensors, backward_visited_leaf_nodes
|
|
backward_visited_leaf_nodes += 1
|
|
|
|
if backward_visited_leaf_nodes == len(leaf_tensors):
|
|
see_memory_usage("After backward checkpointing code after backward", force=False)
|
|
|
|
if PROFILE_TIME:
|
|
timers('backward').stop()
|
|
timers.log(['backward'])
|
|
if SYNCHRONIZE:
|
|
get_accelerator().synchronize()
|
|
|
|
with torch.autograd.graph.saved_tensors_hooks(checkpoint_pack, checkpoint_unpack):
|
|
outputs = function(*inputs_cuda)
|
|
if PROFILE_TIME or SYNCHRONIZE:
|
|
for leaf_tensor in leaf_tensors:
|
|
leaf_tensor.register_hook(after_backward_hook)
|
|
|
|
see_memory_usage("After running forward on the layer", force=False)
|
|
|
|
if PROFILE_TIME:
|
|
timers(FORWARD_GLOBAL_TIMER).stop()
|
|
timers.log([FORWARD_GLOBAL_TIMER])
|
|
if SYNCHRONIZE:
|
|
get_accelerator().synchronize()
|
|
|
|
all_outputs = []
|
|
if torch.is_tensor(outputs):
|
|
all_outputs += [outputs]
|
|
else:
|
|
all_outputs += outputs
|
|
|
|
if len(all_outputs) == 1:
|
|
return all_outputs[0]
|
|
else:
|
|
return tuple(all_outputs)
|
|
|
|
|
|
@compiler.disable # WA from Pytorch repo for compile + zero 3 accuracy issue
|
|
def checkpoint(function, *args):
|
|
"""Checkpoint a model or part of the model.
|
|
This has been directly copied from torch.utils.checkpoint. """
|
|
|
|
all_outputs = []
|
|
CheckpointFunction.apply(function, all_outputs, *args)
|
|
if len(all_outputs) == 1:
|
|
return all_outputs[0]
|
|
else:
|
|
return tuple(all_outputs)
|
|
|
|
|
|
def partition_activations_in_checkpoint(partition_activation):
|
|
global PARTITION_ACTIVATIONS
|
|
PARTITION_ACTIVATIONS = partition_activation
|
|
if dist.get_rank() == 0:
|
|
logger.info(f"**************Partition Activations {PARTITION_ACTIVATIONS}************")
|
|
|
|
|
|
def set_num_layers(nlayers):
|
|
global num_layers
|
|
num_layers = nlayers
|
|
|
|
|
|
def reset():
|
|
"""Resets memory buffers related to contiguous memory optimizations.
|
|
Should be called during eval when multiple forward propagations are
|
|
computed without any backward propagation that usually clears these
|
|
buffers.
|
|
Arguments:
|
|
None
|
|
|
|
Return:
|
|
None
|
|
"""
|
|
if CONTIGUOUS_CHECKPOINTING:
|
|
global data_offsets, size_offsets
|
|
global contiguous_data_buffers, contiguous_size_buffers
|
|
|
|
for buffers in contiguous_data_buffers:
|
|
buffers = []
|
|
|
|
# frees up all the pointers to the checkpoints except for the ones
|
|
# stored by save for backward
|
|
contiguous_data_buffers = []
|
|
contiguous_size_buffers = []
|
|
data_offsets = []
|
|
size_offsets = []
|
|
|
|
|
|
def _configure_using_config_file(config, mpu=None):
|
|
global num_layers, PARTITION_ACTIVATIONS, CONTIGUOUS_CHECKPOINTING, \
|
|
CPU_CHECKPOINT, SYNCHRONIZE, PROFILE_TIME
|
|
|
|
config = DeepSpeedConfig(config, mpu=mpu).activation_checkpointing_config
|
|
if dist.get_rank() == 0:
|
|
logger.info(config.repr())
|
|
PARTITION_ACTIVATIONS = config.partition_activations
|
|
CONTIGUOUS_CHECKPOINTING = config.contiguous_memory_optimization
|
|
num_layers = config.number_checkpoints
|
|
CPU_CHECKPOINT = config.cpu_checkpointing
|
|
SYNCHRONIZE = config.synchronize_checkpoint_boundary
|
|
PROFILE_TIME = config.profile
|
|
|
|
|
|
def _configure_defaults():
|
|
|
|
global mpu, num_layers, deepspeed_checkpointing_enabled
|
|
|
|
global PARTITION_ACTIVATIONS, CONTIGUOUS_CHECKPOINTING, \
|
|
CPU_CHECKPOINT, SYNCHRONIZE, PROFILE_TIME
|
|
|
|
PARTITION_ACTIVATIONS = False
|
|
CONTIGUOUS_CHECKPOINTING = False
|
|
num_layers = None
|
|
CPU_CHECKPOINT = False
|
|
SYNCHRONIZE = False
|
|
PROFILE_TIME = False
|
|
deepspeed_checkpointing_enabled = True
|
|
|
|
|
|
def configure(
|
|
mpu_,
|
|
deepspeed_config=None,
|
|
partition_activations=None,
|
|
contiguous_checkpointing=None,
|
|
num_checkpoints=None,
|
|
checkpoint_in_cpu=None,
|
|
synchronize=None,
|
|
profile=None,
|
|
):
|
|
"""Configure DeepSpeed Activation Checkpointing.
|
|
|
|
Arguments:
|
|
mpu_: Optional: An object that implements the following methods
|
|
get_model_parallel_rank/group/world_size, and get_data_parallel_rank/group/world_size
|
|
|
|
deepspeed_config: Optional: DeepSpeed Config json file when provided will be used to
|
|
configure DeepSpeed Activation Checkpointing
|
|
|
|
partition_activations: Optional: Partitions activation checkpoint across model parallel
|
|
GPUs when enabled. By default False. Will overwrite deepspeed_config if provided
|
|
|
|
contiguous_checkpointing: Optional: Copies activation checkpoints to a contiguous memory
|
|
buffer. Works only with homogeneous checkpoints when partition_activations is enabled.
|
|
Must provide num_checkpoints. By default False. Will overwrite deepspeed_config if
|
|
provided
|
|
|
|
num_checkpoints: Optional: Number of activation checkpoints stored during the forward
|
|
propagation of the model. Used to calculate the buffer size for contiguous_checkpointing
|
|
Will overwrite deepspeed_config if provided
|
|
|
|
checkpoint_in_cpu: Optional: Moves the activation checkpoint to CPU. Only works with
|
|
partition_activation. Default is false. Will overwrite deepspeed_config if provided
|
|
|
|
synchronize: Optional: Performs get_accelerator().synchronize() at the beginning and end of
|
|
each call to deepspeed.checkpointing.checkpoint for both forward and backward pass.
|
|
By default false. Will overwrite deepspeed_config if provided
|
|
|
|
profile: Optional: Logs the forward and backward time for each
|
|
deepspeed.checkpointing.checkpoint invocation. Will overwrite deepspeed_config
|
|
if provided
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
global mpu, num_layers, deepspeed_checkpointing_enabled
|
|
|
|
global PARTITION_ACTIVATIONS, CONTIGUOUS_CHECKPOINTING, \
|
|
CPU_CHECKPOINT, SYNCHRONIZE, PROFILE_TIME
|
|
|
|
_configure_defaults()
|
|
|
|
if mpu_ is not None:
|
|
mpu = mpu_
|
|
|
|
if deepspeed_config is not None:
|
|
_configure_using_config_file(deepspeed_config, mpu=mpu)
|
|
|
|
if partition_activations is not None:
|
|
PARTITION_ACTIVATIONS = partition_activations
|
|
|
|
if contiguous_checkpointing is not None:
|
|
CONTIGUOUS_CHECKPOINTING = contiguous_checkpointing
|
|
|
|
if num_checkpoints is not None:
|
|
num_layers = num_checkpoints
|
|
|
|
if checkpoint_in_cpu is not None:
|
|
CPU_CHECKPOINT = checkpoint_in_cpu
|
|
|
|
if synchronize is not None:
|
|
SYNCHRONIZE = synchronize
|
|
|
|
if profile is not None:
|
|
PROFILE_TIME = profile
|
|
|
|
if CONTIGUOUS_CHECKPOINTING:
|
|
assert PARTITION_ACTIVATIONS, "Contiguous Checkpointing is only available with partitioned activations. Set partitioned activations to true in deepspeed config"
|
|
if CONTIGUOUS_CHECKPOINTING:
|
|
assert num_layers is not None, "Must specify the number of layers with contiguous memory checkpointing"
|
|
|
|
global mp_rank, mp_size, mp_group
|
|
|
|
if mpu is not None:
|
|
if hasattr(mpu, 'get_tensor_model_parallel_rank'):
|
|
mp_rank = mpu.get_tensor_model_parallel_rank()
|
|
mp_size = mpu.get_tensor_model_parallel_world_size()
|
|
mp_group = mpu.get_tensor_model_parallel_group()
|
|
else:
|
|
mp_rank = mpu.get_model_parallel_rank()
|
|
mp_size = mpu.get_model_parallel_world_size()
|
|
mp_group = mpu.get_model_parallel_group()
|
|
|
|
#print configuration only once
|
|
see_memory_usage("After configuration", force=False)
|
|
if dist.get_rank() == 0:
|
|
logger.info("Activation Checkpointing Information")
|
|
logger.info(f"----Partition Activations {PARTITION_ACTIVATIONS}, CPU CHECKPOINTING {CPU_CHECKPOINT}")
|
|
logger.info(f"----contiguous Memory Checkpointing {CONTIGUOUS_CHECKPOINTING} with {num_layers} total layers")
|
|
logger.info(f"----Synchronization {SYNCHRONIZE}")
|
|
logger.info(f"----Profiling time in checkpointing {PROFILE_TIME}")
|
|
|
|
|
|
def is_configured():
|
|
"""True if deepspeed activation checkpointing has been configured
|
|
by calling deepspeed.checkpointing.configure, else returns false
|
|
|
|
Arguments:
|
|
None
|
|
|
|
Return:
|
|
True of configured, else False
|
|
"""
|
|
return deepspeed_checkpointing_enabled
|