perf(optim/dynamo): shortcut is_sparse iteration in SGD multi_tensor (#110648)

jon-chuang · pytorchmergebot · commit d776dd04ac9b · 2023-10-06T08:56:18.000Z
Originated: #110353 (comment) Speeds up significantly in non-sparse path (majority use-case). Benchmarks: #110506 (comment) CC: @janeyx99 Pull Request resolved: #110648 Approved by: https://github.com/janeyx99
diff --git a/torch/optim/sgd.py b/torch/optim/sgd.py
@@ -280,7 +280,7 @@ def _multi_tensor_sgd(params: List[Tensor],
 
     grouped_tensors = Optimizer._group_tensors_by_device_and_dtype([params, grads, momentum_buffer_list], with_indices=True)
     for ((device_params, device_grads, device_momentum_buffer_list), indices) in grouped_tensors.values():
-        device_has_sparse_grad = any(grad.is_sparse for grad in device_grads)
+        device_has_sparse_grad = has_sparse_grad and any(grad.is_sparse for grad in device_grads)
 
         if maximize:
             device_grads = torch._foreach_neg(device_grads)
