Disable grouping by dtype and device if compiling (pytorch#102771)

Disable grouping if we are compiling, this happens during lowering
Pull Request resolved: pytorch#102771
Approved by: https://github.com/janeyx99

torch/optim/adadelta.py

@@ -3,7 +3,6 @@
 
 from .optimizer import (Optimizer, _use_grad_for_differentiable, _default_to_fused_or_foreach,
                         _differentiable_doc, _foreach_doc, _maximize_doc)
-from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype
 from typing import List, Optional
 
 __all__ = ["Adadelta", "adadelta"]
@@ -276,7 +275,7 @@ def _multi_tensor_adadelta(
     if len(params) == 0:
         return
 
-    grouped_tensors = _group_tensors_by_device_and_dtype([params, grads, square_avgs, acc_deltas])
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype([params, grads, square_avgs, acc_deltas])
     for device_params, device_grads, device_square_avgs, device_acc_deltas in grouped_tensors.values():
         if maximize:
             device_grads = torch._foreach_neg(device_grads)

torch/optim/adagrad.py

@@ -3,7 +3,6 @@
 
 from .optimizer import (Optimizer, _use_grad_for_differentiable, _get_value,
                         _default_to_fused_or_foreach, _differentiable_doc, _foreach_doc, _maximize_doc)
-from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype
 from typing import List, Optional
 
 __all__ = ["Adagrad", "adagrad"]
@@ -321,7 +320,7 @@ def _multi_tensor_adagrad(
     if len(params) == 0:
         return
 
-    grouped_tensorlists = _group_tensors_by_device_and_dtype([params, grads, state_sums, state_steps])
+    grouped_tensorlists = Optimizer._group_tensors_by_device_and_dtype([params, grads, state_sums, state_steps])
     for device_params, device_grads, device_state_sums, device_state_steps in grouped_tensorlists.values():
 
         if maximize:

torch/optim/adam.py

@@ -5,7 +5,6 @@
 from .optimizer import (Optimizer, _use_grad_for_differentiable, _get_value, _stack_if_compiling,
                         _dispatch_sqrt, _default_to_fused_or_foreach, _capturable_doc,
                         _differentiable_doc, _foreach_doc, _fused_doc, _maximize_doc)
-from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype
 
 __all__ = ['Adam', 'adam']
 
@@ -424,7 +423,8 @@ def _multi_tensor_adam(params: List[Tensor],
 
     assert not differentiable, "_foreach ops don't support autograd"
 
-    grouped_tensors = _group_tensors_by_device_and_dtype([params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps])
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(
+        [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps])
     for (device_params, device_grads, device_exp_avgs, device_exp_avg_sqs,
          device_max_exp_avg_sqs, device_state_steps) in grouped_tensors.values():
 
@@ -532,7 +532,8 @@ def _fused_adam(
 ) -> None:
     grad_scale_dict = {grad_scale.device: grad_scale} if grad_scale is not None else None
     found_inf_dict = {found_inf.device: found_inf} if found_inf is not None else None
-    grouped_tensors = _group_tensors_by_device_and_dtype([params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps])
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(
+        [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps])
     for (device, dtype) in grouped_tensors:
         (
             device_params,

torch/optim/adamax.py

@@ -4,7 +4,6 @@
 from .optimizer import (Optimizer, _use_grad_for_differentiable, _get_value, _stack_if_compiling,
                         _default_to_fused_or_foreach, _differentiable_doc, _maximize_doc, _foreach_doc)
 from typing import List, Optional
-from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype
 
 __all__ = ["Adamax", "adamax"]
 
@@ -305,7 +304,7 @@ def _multi_tensor_adamax(
     if len(params) == 0:
         return
 
-    grouped_tensors = _group_tensors_by_device_and_dtype([params, grads, exp_avgs, exp_infs, state_steps])
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype([params, grads, exp_avgs, exp_infs, state_steps])
     for grouped_params, grouped_grads, grouped_exp_avgs, grouped_exp_infs, grouped_state_steps in grouped_tensors.values():
         if maximize:
             grouped_grads = torch._foreach_neg(grouped_grads)

torch/optim/adamw.py

@@ -4,7 +4,6 @@
                         _stack_if_compiling, _capturable_doc, _differentiable_doc, _foreach_doc,
                         _fused_doc, _maximize_doc, _default_to_fused_or_foreach)
 from typing import List, Optional
-from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype
 
 __all__ = ["AdamW", "adamw"]
 
@@ -476,7 +475,7 @@ def _multi_tensor_adamw(
 
     assert grad_scale is None and found_inf is None
 
-    grouped_tensors = _group_tensors_by_device_and_dtype([
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype([
         params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps])
     for (device_params, device_grads, device_exp_avgs, device_exp_avg_sqs,
          device_max_exp_avg_sqs, device_state_steps) in grouped_tensors.values():
@@ -593,7 +592,8 @@ def _fused_adamw(
         raise RuntimeError("_fused_adamw is not differentiable")
     grad_scale_dict = {grad_scale.device: grad_scale} if grad_scale is not None else None
     found_inf_dict = {found_inf.device: found_inf} if found_inf is not None else None
-    grouped_tensors = _group_tensors_by_device_and_dtype([params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps])
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(
+        [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps])
     for (device, dtype) in grouped_tensors:
         (
             device_params,

torch/optim/asgd.py

@@ -4,7 +4,6 @@
 from .optimizer import (Optimizer, _use_grad_for_differentiable, _get_value, _default_to_fused_or_foreach,
                         _differentiable_doc, _foreach_doc, _maximize_doc)
 from torch._utils import is_compiling
-from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype
 from typing import List, Optional
 
 __all__ = ["ASGD", "asgd"]
@@ -294,7 +293,7 @@ def _multi_tensor_asgd(
 
     assert not differentiable, "_foreach ops don't support autograd"
 
-    grouped_tensors = _group_tensors_by_device_and_dtype([params, grads, axs, mus, etas, state_steps])
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype([params, grads, axs, mus, etas, state_steps])
     for (grouped_params, grouped_grads, grouped_axs, grouped_mus,
          grouped_etas, grouped_state_steps) in grouped_tensors.values():
         if maximize:

torch/optim/nadam.py

@@ -3,7 +3,6 @@
 from .optimizer import (Optimizer, _use_grad_for_differentiable, _get_value, _dispatch_sqrt, _stack_if_compiling,
                         _differentiable_doc, _foreach_doc, _default_to_fused_or_foreach)
 from typing import List, Optional
-from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype
 
 __all__ = ['NAdam', 'nadam']
 
@@ -291,8 +290,7 @@ def _multi_tensor_nadam(params: List[Tensor],
 
     assert not differentiable, "_foreach ops don't support autograd"
 
-    grouped_tensors = _group_tensors_by_device_and_dtype([params, grads, exp_avgs, exp_avg_sqs,
-                                                          mu_products, state_steps])
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype([params, grads, exp_avgs, exp_avg_sqs, mu_products, state_steps])
     for (grouped_params, grouped_grads, grouped_exp_avgs,
          grouped_exp_avg_sqs, grouped_mu_products, grouped_state_steps) in grouped_tensors.values():
 

torch/optim/optimizer.py

@@ -11,6 +11,7 @@
 import torch.utils.hooks as hooks
 from torch.utils.hooks import RemovableHandle
 from torch._utils import is_compiling
+from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype
 
 __all__ = ['Optimizer', 'register_optimizer_step_pre_hook', 'register_optimizer_step_post_hook']
 _global_optimizer_pre_hooks: Dict[int, Callable] = OrderedDict()
@@ -288,6 +289,18 @@ def wrapper(*args, **kwargs):
 
         return wrapper
 
+    @staticmethod
+    def _group_tensors_by_device_and_dtype(tensorlistlist, with_indices=False):
+        """Groups a list of lists of tensors by device and dtype.
+        Skips this step if we are compiling since this will occur during inductor lowering."""
+        if is_compiling():
+            if with_indices:
+                indices = list(range(len(tensorlistlist[0])))
+                tensorlistlist.append(indices)
+            return {(None, None): tensorlistlist}
+        else:
+            return _group_tensors_by_device_and_dtype(tensorlistlist, with_indices)
+
     def _patch_step_function(self):
         self._zero_grad_profile_name = "Optimizer.zero_grad#{}.zero_grad".format(self.__class__.__name__)
         hooked = getattr(self.__class__.step, "hooked", None)

torch/optim/radam.py

@@ -5,7 +5,6 @@
 from .optimizer import (Optimizer, _use_grad_for_differentiable, _get_value, _dispatch_sqrt, _stack_if_compiling,
                         _default_to_fused_or_foreach, _differentiable_doc, _foreach_doc)
 from typing import List, Optional
-from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype
 
 __all__ = ["RAdam", "radam"]
 
@@ -315,7 +314,7 @@ def _multi_tensor_radam(
 
     assert not differentiable, "_foreach ops don't support autograd"
 
-    grouped_tensors = _group_tensors_by_device_and_dtype([params, grads, exp_avgs, exp_avg_sqs, state_steps])
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype([params, grads, exp_avgs, exp_avg_sqs, state_steps])
     for grouped_params, grouped_grads, grouped_exp_avgs, grouped_exp_avg_sqs, grouped_state_steps in grouped_tensors.values():
         # Update steps
         torch._foreach_add_(grouped_state_steps, 1)

torch/optim/rmsprop.py

@@ -3,7 +3,6 @@
 from .optimizer import (Optimizer, _default_to_fused_or_foreach, _use_grad_for_differentiable,
                         _differentiable_doc, _foreach_doc, _maximize_doc)
 from typing import List, Optional
-from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype
 
 __all__ = ["RMSprop", "rmsprop"]
 
@@ -326,7 +325,7 @@ def _multi_tensor_rmsprop(
 
     assert not differentiable, "_foreach ops don't support autograd"
 
-    grouped_tensors = _group_tensors_by_device_and_dtype([params, grads, square_avgs, grad_avgs, momentum_buffer_list])
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype([params, grads, square_avgs, grad_avgs, momentum_buffer_list])
     for (grouped_params, grouped_grads, grouped_square_avgs, grouped_grad_avgs,
          grouped_momentum_buffer_list) in grouped_tensors.values():
         if maximize:

torch/optim/rprop.py

@@ -3,7 +3,6 @@
 from .optimizer import (Optimizer, _use_grad_for_differentiable, _default_to_fused_or_foreach,
                         _differentiable_doc, _foreach_doc, _maximize_doc)
 from typing import List, Optional
-from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype
 
 __all__ = ["Rprop", "rprop"]
 
@@ -281,7 +280,7 @@ def _multi_tensor_rprop(
 
     assert not differentiable, "_foreach ops don't support autograd"
 
-    grouped_tensors = _group_tensors_by_device_and_dtype([params, grads, prevs, step_sizes])
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype([params, grads, prevs, step_sizes])
     for grouped_params, grouped_grads, grouped_prevs, grouped_step_sizes in grouped_tensors.values():
         # Handle complex params
         def _view_complex_as_real(tensor_list):

torch/optim/sgd.py

@@ -3,7 +3,6 @@
 from .optimizer import (Optimizer, required, _use_grad_for_differentiable, _default_to_fused_or_foreach,
                         _differentiable_doc, _foreach_doc, _maximize_doc)
 from typing import List, Optional
-from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype
 
 __all__ = ['SGD', 'sgd']
 
@@ -280,7 +279,7 @@ def _multi_tensor_sgd(params: List[Tensor],
     if len(params) == 0:
         return
 
-    grouped_tensors = _group_tensors_by_device_and_dtype([params, grads, momentum_buffer_list], with_indices=True)
+    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)