Use 'not in'

Signed-off-by: cyy <[email protected]>

Author: cyyever

deepspeed/checkpoint/deepspeed_checkpoint.py

@@ -137,7 +137,7 @@ def get_final_norm_layer_id(self):
         return self.layer_keys[self.final_layer_norm_idx]
 
     def get_iteration(self):
-        if not ITERATION_KEY in self.global_state:
+        if ITERATION_KEY not in self.global_state:
             sd = torch.load(self.mp_rank_files[0], map_location=torch.device('cpu'), weights_only=False)
             self.global_state[ITERATION_KEY] = sd.get(ITERATION_KEY, 0)
 
@@ -157,7 +157,7 @@ def get_embedding_files(self, tp_index: int) -> list:
         return self.tp_to_embedding_map[tp_index]
 
     def _get_checkpoint_value(self, key):
-        if not key in self.global_state:
+        if key not in self.global_state:
             sd = torch.load(self.mp_rank_files[0], map_location=torch.device('cpu'), weights_only=False)
             self.global_state[key] = sd.get(key, None)
 
@@ -254,7 +254,7 @@ def _build_transformer_file_map(self):
             layer_file_partitions = partition_data(layer_files, self.tp_degree)
             for tp_index in range(self.tp_degree):
                 map_key = (tp_index, pp_index)
-                if not map_key in file_map.keys():
+                if map_key not in file_map.keys():
                     file_map[map_key] = []
                 file_map[map_key].append(layer_file_partitions[tp_index])
 
@@ -286,7 +286,7 @@ def _get_layer_keys(self):
     def _merge_state_dicts(self, sd_list):
         merged_sd = {}
         for key in sd_list[0].keys():
-            if not key in SEQUENTIAL_LAYERS:
+            if key not in SEQUENTIAL_LAYERS:
                 cat_dim = LAYER_CONCAT_DIM.get(key, 0)
                 merged_sd[key] = torch.cat([sd[key] for sd in sd_list], dim=cat_dim)
             else:

deepspeed/checkpoint/reshape_meg_2d.py

@@ -24,7 +24,7 @@ def add_data(self, pp_index, tp_index, data):
         assert type(data) is list
 
         key = self._make_key(pp_index, tp_index)
-        if not key in self.map.keys():
+        if key not in self.map.keys():
             self.map[key] = []
         self.map[key] += data
 

deepspeed/compile/fx.py

@@ -115,7 +115,7 @@ def add_free_activations(graph_id: int, graph: Graph, activation_node_names: Lis
     def _should_free(node: Node) -> bool:
         if not hasattr(node, "meta"):
             return False
-        if not "tensor_meta" in node.meta:
+        if "tensor_meta" not in node.meta:
             return False
         return True
 

deepspeed/compile/passes/offload_activation.py

@@ -34,7 +34,7 @@ def _gen():
 def _should_offload(node: Node) -> bool:
     if not hasattr(node, "meta"):
         return False
-    if not "tensor_meta" in node.meta:
+    if "tensor_meta" not in node.meta:
         return False
 
     return True

deepspeed/nvme/parse_nvme_stats.py

@@ -101,7 +101,7 @@ def get_metric(file, metric):
 
 
 def validate_args(args):
-    if not args.metric in PERF_METRICS:
+    if args.metric not in PERF_METRICS:
         print(f'{args.metric} is not a valid performance metrics')
         return False
 

deepspeed/runtime/engine.py

@@ -3198,7 +3198,7 @@ def get_sparse_tensor_module_names(original_set, loaded_set, original_parameters
         if load_optimizer_states:
             deepspeed_states.append('optimizer')
 
-        client_state = {key: value for key, value in checkpoint.items() if not key in deepspeed_states}
+        client_state = {key: value for key, value in checkpoint.items() if key not in deepspeed_states}
 
         if optim_checkpoint is not None:
             client_state['optimizer'] = optim_checkpoint['optimizer']

deepspeed/runtime/lr_schedules.py

@@ -209,7 +209,7 @@ def get_config_from_args(args):
     if not hasattr(args, LR_SCHEDULE) or args.lr_schedule is None:
         return None, '--{} not specified on command line'.format(LR_SCHEDULE)
 
-    if not args.lr_schedule in VALID_LR_SCHEDULES:
+    if args.lr_schedule not in VALID_LR_SCHEDULES:
         return None, '{} is not supported LR schedule'.format(args.lr_schedule)
 
     config = {}
@@ -227,16 +227,16 @@ def get_config_from_args(args):
 
 
 def get_lr_from_config(config):
-    if not 'type' in config:
+    if 'type' not in config:
         return None, 'LR schedule type not defined in config'
 
-    if not 'params' in config:
+    if 'params' not in config:
         return None, 'LR schedule params not defined in config'
 
     lr_schedule = config['type']
     lr_params = config['params']
 
-    if not lr_schedule in VALID_LR_SCHEDULES:
+    if lr_schedule not in VALID_LR_SCHEDULES:
         return None, '{} is not a valid LR schedule'.format(lr_schedule)
 
     if lr_schedule == LR_RANGE_TEST:

deepspeed/runtime/swap_tensor/optimizer_utils.py

@@ -102,7 +102,7 @@ def get_swap_buffers_and_paths(self, pinned):
     def get_or_create_gradient_paths(self, offsets, lengths):
         gradient_paths = []
         for offset, length in zip(offsets, lengths):
-            if not offset in self.swapped_gradients.keys():
+            if offset not in self.swapped_gradients.keys():
                 path = os.path.join(self.swap_folder, f'{self.param_id}_gradient_{offset}_{length}.tensor.swp')
                 self.swapped_gradients[offset] = FlattenedTensorSwapInfo(path, length, offset)
 
@@ -233,7 +233,7 @@ def _flush_gradient_swapper(self, gradient_swapper):
             self.timer_names.update(gradient_swapper.get_timer_names())
 
     def _swap_out_gradients(self, parameter, gradient_offsets, gradient_tensors, gradient_swapper):
-        if not OptimizerSwapper.parameter_id(parameter) in self.swap_params_info.keys():
+        if OptimizerSwapper.parameter_id(parameter) not in self.swap_params_info.keys():
             return
 
         swap_info = self.swap_params_info[OptimizerSwapper.parameter_id(parameter)]
@@ -471,7 +471,7 @@ def _retrieve_unswapped_grad_partitions(self, swap_info, dest_buffer):
             )
 
     def _get_state_tensors(self, parameter):
-        if not parameter in self.optimizer.state:
+        if parameter not in self.optimizer.state:
             return []
 
         tensor_list = []
@@ -490,7 +490,7 @@ def _update_param_state_info(self, swap_info, parameter):
 
     def _create_param_swap_info(self, parameter, numel):
         param_id = OptimizerSwapper.parameter_id(parameter)
-        assert not param_id in self.swap_params_info
+        assert param_id not in self.swap_params_info
 
         self.swap_params_info[param_id] = OptimizerStateSwapInfo(parameter=parameter,
                                                                  numel=numel,

deepspeed/runtime/swap_tensor/utils.py

@@ -30,7 +30,7 @@ def swap_out_tensors(swap_handle, tensor_buffers, swap_paths):
 def print_object(obj, name, exclude_list=[]):
     logger.info('{}:'.format(name))
     for arg in sorted(vars(obj)):
-        if not arg in exclude_list:
+        if arg not in exclude_list:
             dots = '.' * (29 - len(arg))
             logger.info('  {} {} {}'.format(arg, dots, getattr(obj, arg)))
 
@@ -55,7 +55,7 @@ def insert_tensor(self, tensor, swap_path, aligned_numel):
 
     def allocate_tensor(self, swap_path, numel, aligned_numel):
         assert self.has_space(aligned_numel)
-        assert not self.offset in self.swap_tensors
+        assert self.offset not in self.swap_tensors
 
         allocate_offset = self.offset
         swap_tensor = self.buffer.narrow(0, allocate_offset, aligned_numel)

deepspeed/runtime/zero/contiguous_memory_allocator.py

@@ -85,7 +85,7 @@ def assign_to_param(self, tensor, param, numel, shape):
 
         assert tensor_id in self.tensor_map.keys(), "No such tensor allocated by the allocator."
         assert tensor.numel() >= numel, "Assert tensor buffer does is not large enough"
-        assert not tensor_id in self.id_to_params.keys(), "This tensor has already been assigned to a param"
+        assert tensor_id not in self.id_to_params.keys(), "This tensor has already been assigned to a param"
 
         self.id_to_params[tensor_id] = [param]