Skip to content

Fix ZeRO3 save_checkpoint #857

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
jeffra merged 12 commits into from
Mar 16, 2021
Merged

Fix ZeRO3 save_checkpoint #857

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
3aea965
Fix ZeRO3 save_checkpoint
tjruwase Mar 12, 2021
c256197
Merge branch 'master' into olruwase/zero3_save_checkpoint
jeffra Mar 12, 2021
69b9aaa
Merge branch 'master' into olruwase/zero3_save_checkpoint
tjruwase Mar 15, 2021
9e4b306
Merge branch 'master' of github.com:microsoft/DeepSpeed into olruwase…
tjruwase Mar 15, 2021
74ce362
Merge branch 'master' into olruwase/zero3_save_checkpoint
jeffra Mar 15, 2021
bf47165
turn checkpoint test back on
jeffra Mar 15, 2021
1809b9a
Merge branch 'master' into olruwase/zero3_save_checkpoint
tjruwase Mar 15, 2021
0371abb
formatting
jeffra Mar 16, 2021
9016e20
debug prints
tjruwase Mar 16, 2021
45f446d
Remove debug prints; formatting
tjruwase Mar 16, 2021
2770cb9
Merge branch 'master' into olruwase/zero3_save_checkpoint
tjruwase Mar 16, 2021
9d7e2ae
Merge branch 'master' into olruwase/zero3_save_checkpoint
jeffra Mar 16, 2021
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
12 changes: 5 additions & 7 deletions deepspeed/runtime/zero/stage3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2269,7 +2269,7 @@ def _prepare_fp32_grad_for_sub_group(self, sub_group_id):

assert single_grad_partition.numel() == self.fp32_partitioned_groups_flat[sub_group_id].numel(), \
"averaged gradients have different number of elements that partition size {} {} {} {}".format(
single_grad_partition.numel(), self.partition_size[sub_group_id], sub_group_id, partition_id)
single_grad_partition.numel(), self.fp32_partitioned_groups_flat[sub_group_id].numel(), sub_group_id, partition_id)

self.fp32_partitioned_groups_flat[sub_group_id].grad = single_grad_partition

Expand Down Expand Up @@ -2638,14 +2638,12 @@ def get_groups_without_padding(self, groups_with_padding):
def _set_fp32_optimizer_param_groups(self):
for sub_group_id, _ in enumerate(self.fp16_groups):
param_group_id = self.sub_group_to_group_id[sub_group_id]
self.optimizer.param_groups[param_group_id]['params'] = [
self.fp32_partitioned_groups_flat[sub_group_id]
]
self.optimizer.param_groups[param_group_id]['params'].append(
self.fp32_partitioned_groups_flat[sub_group_id])

def _clear_fp32_optimizer_param_groups(self):
for sub_group_id, _ in enumerate(self.fp16_groups):
param_group_id = self.sub_group_to_group_id[sub_group_id]
self.optimizer.param_groups[param_group_id]['params'] = []
for param_group in self.optimizer.param_groups:
param_group['params'] = []

def _rigid_state_dict(self):
state_dict = {}
Expand Down
44 changes: 24 additions & 20 deletions tests/unit/test_checkpointing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ def compare_model_states(saved_model, loaded_model, compare_optimizer=True):
if FP16_DeepSpeedZeroOptimizer_Stage3 is not None and isinstance(
saved_model.optimizer,
FP16_DeepSpeedZeroOptimizer_Stage3):
for p0, p1 in zip(saved_model.optimizer.fp32_groups_flat, loaded_model.optimizer.fp32_groups_flat):
for p0, p1 in zip(saved_model.optimizer.fp32_partitioned_groups_flat, loaded_model.optimizer.fp32_partitioned_groups_flat):
assert torch.allclose(p0, p1, atol=1e-07), f"Fp32 model states {p0} is not equal to {p1}"

elif isinstance(saved_model.optimizer, FP16_DeepSpeedZeroOptimizer):
Expand Down Expand Up @@ -303,12 +303,13 @@ def _test_checkpoint_fused_optimizer(args,
'deepspeed_adam'),
(3,
False,
'Adam')])
'Adam'),
(3,
True,
'deepspeed_adam')])
def test_checkpoint_zero_optimizer(tmpdir, zero_stage, use_cpu_offload, adam_optimizer):
if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
pytest.skip("cpu-adam is not compatible")
if zero_stage == 3:
pytest.skip('Skip checkpointing tests for ZeRO3')

config_dict = {
"train_batch_size": 2,
Expand All @@ -324,8 +325,10 @@ def test_checkpoint_zero_optimizer(tmpdir, zero_stage, use_cpu_offload, adam_opt
}
},
"fp16": {
"enabled": True
"enabled": True,
"initial_scale_power": 8
},
"wall_clock_breakdown": True,
"zero_optimization": {
"stage": zero_stage,
"cpu_offload": use_cpu_offload
Expand All @@ -340,9 +343,7 @@ def _test_checkpoint_zero_optimizer(args,
hidden_dim,
load_optimizer_states):
if zero_stage == 3:
global FP16_DeepSpeedZeroOptimizer_Stage3
from deepspeed.runtime.zero.stage3 import FP16_DeepSpeedZeroOptimizer_Stage3
with deepspeed.ScatteredParameters(zero_modules=True):
with deepspeed.zero.Init():
models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
else:
models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
Expand Down Expand Up @@ -371,15 +372,16 @@ def _test_checkpoint_zero_optimizer(args,
'deepspeed_adam'),
(3,
False,
'Adam')])
'Adam'),
(3,
True,
'deepspeed_adam')])
def test_checkpoint_zero_no_optimizer(tmpdir,
zero_stage,
use_cpu_offload,
adam_optimizer):
if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
pytest.skip("cpu-adam is not compatible")
if zero_stage == 3:
pytest.skip('Skip checkpointing tests for ZeRO3')

config_dict = {
"train_batch_size": 2,
Expand Down Expand Up @@ -413,7 +415,7 @@ def _test_checkpoint_zero_no_optimizer(args,
if zero_stage == 3:
global FP16_DeepSpeedZeroOptimizer_Stage3
from deepspeed.runtime.zero.stage3 import FP16_DeepSpeedZeroOptimizer_Stage3
with deepspeed.ScatteredParameters(zero_modules=True):
with deepspeed.zero.Init():
models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
else:
models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
Expand Down Expand Up @@ -445,12 +447,13 @@ def _test_checkpoint_zero_no_optimizer(args,
'deepspeed_adam'),
(3,
False,
'Adam')])
'Adam'),
(3,
True,
'deepspeed_adam')])
def test_checkpoint_lr_scheduler(tmpdir, zero_stage, use_cpu_offload, adam_optimizer):
if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
pytest.skip("cpu-adam is not compatible")
if zero_stage == 3:
pytest.skip('Skip checkpointing tests for ZeRO3')

config_dict = {
"train_batch_size": 2,
Expand Down Expand Up @@ -493,7 +496,7 @@ def _test_checkpoint_lr_scheduler(args,
if zero_stage == 3:
global FP16_DeepSpeedZeroOptimizer_Stage3
from deepspeed.runtime.zero.stage3 import FP16_DeepSpeedZeroOptimizer_Stage3
with deepspeed.ScatteredParameters(zero_modules=True):
with deepspeed.zero.Init():
models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
else:
models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
Expand Down Expand Up @@ -526,14 +529,15 @@ def _test_checkpoint_lr_scheduler(args,
(2,
True,
'deepspeed_adam'),
(3,
False,
'Adam'),
(3,
True,
'Adam')])
'deepspeed_adam')])
def test_checkpoint_no_lr_scheduler(tmpdir, zero_stage, use_cpu_offload, adam_optimizer):
if use_cpu_offload and not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
pytest.skip("cpu-adam is not compatible")
if zero_stage == 3:
pytest.skip('Skip checkpointing tests for ZeRO3')

config_dict = {
"train_batch_size": 2,
Expand Down Expand Up @@ -570,7 +574,7 @@ def _test_checkpoint_no_lr_scheduler(args,
load_optimizer_states,
load_lr_scheduler_states):
if zero_stage == 3:
with deepspeed.ScatteredParameters(zero_modules=True):
with deepspeed.zero.Init():
models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
else:
models = [SimpleModel(hidden_dim, empty_grad=False) for _ in range(2)]
Expand Down