diff --git a/deepspeed/pt/deepspeed_checkpointing.py b/deepspeed/pt/deepspeed_checkpointing.py index 68d3b6bff4f7..a38db1f498db 100755 --- a/deepspeed/pt/deepspeed_checkpointing.py +++ b/deepspeed/pt/deepspeed_checkpointing.py @@ -14,6 +14,7 @@ # Parts of the code here are adapted from PyTorch # repo: https://github.com/pytorch/pytorch import contextlib +import copy import torch.distributed as dist import torch from torch import _C @@ -108,7 +109,7 @@ def detach_variable(inputs, device=None): def _set_cuda_rng_state(new_state, device=-1): """Sets the random number generator state of the current GPU. - Argumentss: + Arguments: new_state (torch.ByteTensor): The desired state This function is adapted from PyTorch repo (torch.cuda.set_rng_state) with a single change: the input state is not cloned. Cloning caused @@ -160,10 +161,7 @@ def reset(self): 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.""" - states = {} - for name in self.states_: - states[name] = self.states_[name] - return states + return copy.copy(self.states_) def set_states(self, states): """Set the rng states. For efficiency purposes, we do not check @@ -720,5 +718,4 @@ def is_configured(): Return: True of configured, else False """ - global deepspeed_checkpointing_enabled return deepspeed_checkpointing_enabled diff --git a/deepspeed/pt/deepspeed_config_utils.py b/deepspeed/pt/deepspeed_config_utils.py index d99c5ab197db..38fdb647f61d 100755 --- a/deepspeed/pt/deepspeed_config_utils.py +++ b/deepspeed/pt/deepspeed_config_utils.py @@ -6,20 +6,18 @@ Collection of DeepSpeed configuration utilities """ +from collections import Counter + def get_scalar_param(param_dict, param_name, param_default_value): - if param_name in param_dict.keys(): - return param_dict[param_name] - else: - return param_default_value + return param_dict.get(param_name, param_default_value) def dict_raise_error_on_duplicate_keys(ordered_pairs): """Reject duplicate keys.""" - d = {} - for k, v in ordered_pairs: - if k in d: - raise ValueError("Duplicate key in DeepSpeed config: %r" % (k, )) - else: - d[k] = v + d = dict((k, v) for k, v in ordered_pairs) + if len(d) != len(ordered_pairs): + counter = Counter([pair[0] for pair in ordered_pairs]) + keys = [key for key, value in counter.items() if value > 1] + raise ValueError("Duplicate keys in DeepSpeed config: {}".format(keys)) return d