Add more robust arg checks

deepspeed/comm/comm.py

@@ -87,7 +87,7 @@ class ReduceOp(Enum):
 # - Global profiling (profile all comms)
 # - Op-type profiling (e.g. profile all all_reduce comms)
 # - Op profiling (e.g. profile a specific all_reduce op)
-prof_all = False
+prof_all = True
 prof_op = None
 
 
@@ -121,17 +121,24 @@ def log_wrapper(*args, **kwargs):
         # Need func args and their defaults
         func_args = get_default_args(func)
         func_args.update(kwargs)
-        tensor_pos = get_tensor_position(func)
+        if len(args) > 0:
+            tensor_pos = get_tensor_position(func)
+            tensor_arg = args[tensor_pos]
+        else:
+            if len(kwargs) > 0:
+                tensor_arg = func_args['tensor']
         # Get size of tensor to be communicated
         # set msg_size = 0 for barrier
-        msg_size = 0
-        # Sum of tensor sizes for list colls
-        if type(args[tensor_pos]) is list:
-            msg_size = sum(x.element_size() * x.nelement()
-                           for x in func_args['tensor_list'])
-        # msg_size = tensor size for most colls
+        if len(kwargs) == 0:
+            msg_size = 0
         else:
-            msg_size = args[tensor_pos].element_size() * args[tensor_pos].nelement()
+            # Sum of tensor sizes for list colls
+            if type(tensor_arg) is list:
+                msg_size = sum(x.element_size() * x.nelement()
+                               for x in func_args['tensor_list'])
+            # msg_size = tensor size for most colls
+            else:
+                msg_size = tensor_arg.element_size() * tensor_arg.nelement()
         # Start the timer if arg is set or it's a default
         if func_args['prof'] or prof_all or func_args['log_name'] == prof_op:
             timers(func_args['log_name']).start()

deepspeed/comm/utils.py

@@ -103,4 +103,7 @@ def get_tensor_position(func):
     # all_to_all and torch multiGPU colls
     elif 'input_tensor_list' in sig_params:
         arg = 'input_tensor_list'
-    return list(sig_params).index(arg)
+    if arg is None:
+        return -1
+    else:
+        return list(sig_params).index(arg)

deepspeed/runtime/zero/stage3.py

@@ -102,11 +102,11 @@ def _apply_to_tensors_only(module, functional, backward_function, outputs):
     elif type(outputs) is torch.Tensor:
         return functional.apply(module, backward_function, outputs)
     else:
-        if not is_builtin_type(outputs):
-            logger.warning(
-                f"A module has unknown inputs or outputs type ({type(outputs)}) and the tensors embedded in it cannot be detected. "
-                "The ZeRO-3 hooks designed to trigger before or after backward pass of the module relies on knowing the input and "
-                "output tensors and therefore may not get triggered properly.")
+        #if not is_builtin_type(outputs):
+        #    logger.warning(
+        #        f"A module has unknown inputs or outputs type ({type(outputs)}) and the tensors embedded in it cannot be detected. "
+        #        "The ZeRO-3 hooks designed to trigger before or after backward pass of the module relies on knowing the input and "
+        #        "output tensors and therefore may not get triggered properly.")
         return outputs
 
 

tests/comm/all_reduce.py

@@ -14,7 +14,7 @@
 def timed_allreduce(mat):
     torch.cuda.synchronize()
     pre = time.perf_counter()
-    dist.all_reduce(mat)
+    dist.all_reduce(tensor=mat)
     #print('ignore me', mat[0][0])  # required due to lazy evaluation
     torch.cuda.synchronize()
     duration = time.perf_counter() - pre
@@ -62,7 +62,7 @@ def init_processes(fn, backend='nccl', use_deepspeed=False):
     print(f'local rank = {dist.get_local_rank()}')
     torch.cuda.set_device(dist.get_local_rank())
     dist.start_profiling_comms()
-    dist.set_comms_log_verbose(True)
+    #dist.set_comms_log_verbose(True)
     fn(dist.get_local_rank())