Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

move is_checkpointable call reducing torch.compile Graph breaks #5759

Merged
merged 5 commits into from
Jul 23, 2024
Merged

move is_checkpointable call reducing torch.compile Graph breaks #5759

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
897cf8e
move is_checkpointable call reducing t.comp breaks
NirSonnenschein Jul 8, 2024
bcf1966
Merge branch 'master' into torch_compile_is_checkpointable_fix_1
HeyangQin Jul 12, 2024
7d0eac6
Merge branch 'master' into torch_compile_is_checkpointable_fix_1
tohtana Jul 22, 2024
7636fe8
Merge branch 'master' into torch_compile_is_checkpointable_fix_1
loadams Jul 23, 2024
7acbf0f
Merge branch 'master' into torch_compile_is_checkpointable_fix_1
loadams Jul 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions deepspeed/runtime/pipe/engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,8 @@ def __init__(self, has_bool_tensors=False, *super_args, **super_kwargs):
self.module.activation_checkpoint_func = ds_checkpointing.non_reentrant_checkpoint
if self.grid.get_global_rank() == 0:
logger.info(f'CONFIG: activation_checkpoint_func=non_reentrant_checkpoint')
if self.module.activation_checkpoint_interval > 0:
self.module._precompute_checkpointable_values()

self.module.checkpoint_parallel_write_pipeline = self._config.checkpoint_parallel_write_pipeline

Expand Down
29 changes: 23 additions & 6 deletions deepspeed/runtime/pipe/module.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -196,17 +196,32 @@ def __init__(self,
#newseed = get_accelerator().initial_seed() + self._grid.get_stage_id()
#ds_utils.set_random_seed(newseed)

self.activation_checkpoint_interval = activation_checkpoint_interval

self.activation_checkpoint_func = activation_checkpoint_func

#storage for precomputed checkpointeble results
self.is_checkpointable_results = []
self.is_checkpointable_results_interval = None

# if configuration use_reentrant = False, self.activation_checkpoint_func will be set to ``checkpointing.non_reentrant_checkpoint``

#with torch.random.fork_rng(devices=[get_accelerator().current_device_name()]):
self._build()
self.to(get_accelerator().device_name(self.local_rank))

self.tied_comms = self._index_tied_modules()
self._synchronize_tied_weights()

self.activation_checkpoint_interval = activation_checkpoint_interval

self.activation_checkpoint_func = activation_checkpoint_func
# if configuration use_reentrant = False, self.activation_checkpoint_func will be set to ``checkpointing.non_reentrant_checkpoint``
def _precompute_checkpointable_values(self):
if self.activation_checkpoint_interval > 0 and self.is_checkpointable_results_interval != self.activation_checkpoint_interval:
num_layers = len(self.forward_funcs)
self.interval_was_zero = False
for start_idx in range(0, num_layers, self.activation_checkpoint_interval):
end_idx = min(start_idx + self.activation_checkpoint_interval, num_layers)
funcs = self.forward_funcs[start_idx:end_idx]
self.is_checkpointable_results.append(self._is_checkpointable(funcs))
self.is_checkpointable_results_interval = self.activation_checkpoint_interval

def _build(self):
specs = self._layer_specs
Expand Down Expand Up @@ -352,7 +367,9 @@ def exec_func(*inputs):
else:
num_layers = len(self.forward_funcs)
x = forward_input
for start_idx in range(0, num_layers, self.activation_checkpoint_interval):
for start_idx, is_checkpointable_result in \
zip(range(0, num_layers, self.activation_checkpoint_interval), self.is_checkpointable_results):

end_idx = min(start_idx + self.activation_checkpoint_interval, num_layers)

funcs = self.forward_funcs[start_idx:end_idx]
Expand All @@ -361,7 +378,7 @@ def exec_func(*inputs):
if not isinstance(x, tuple):
x = (x, )

if self._is_checkpointable(funcs):
if is_checkpointable_result:
x = self.activation_checkpoint_func(exec_range_func(start_idx, end_idx), *x)
else:
x = exec_range_func(start_idx, end_idx)(*x)
Expand Down
Loading