Skip to content

skip torch.zeros and tensor.copy_ when model parallel is not used #2479

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

Merged
tjruwase merged 5 commits into from
Dec 16, 2022
Merged

skip torch.zeros and tensor.copy_ when model parallel is not used #2479

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
82831af
skip torch.zeros and tensor.copy_ when model parallel is not used
guoyejun Nov 7, 2022
e6c2781
Merge branch 'master' into checkpointing2
tjruwase Nov 28, 2022
19bb900
Merge branch 'master' into checkpointing2
tjruwase Nov 30, 2022
7aaad6e
Merge branch 'master' into checkpointing2
tjruwase Dec 1, 2022
1db711d
Merge branch 'master' into checkpointing2
tjruwase Dec 16, 2022
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
7 changes: 3 additions & 4 deletions deepspeed/runtime/activation_checkpointing/checkpointing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -270,8 +270,8 @@ def gather_partitioned_activations(tensors, device=None):
inputs.append(item)
continue

# don't need to do all_gather if model parallel size is 1
if mp_size == 1:
# 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()))
inputs.append(item)
continue
Expand All @@ -290,8 +290,7 @@ def gather_partitioned_activations(tensors, device=None):
if i == mp_rank:
part_i.copy_(item)
partitions.append(part_i)
if mp_group is not None:
dist.all_gather(partitions, partitions[mp_rank], group=mp_group)
dist.all_gather(partitions, partitions[mp_rank], group=mp_group)
input_tensor = flat_tensor.view(list(size.numpy()))
item.data = input_tensor.data

Expand Down