BC for grad scaler type (#753)

Summary: Pull Request resolved: #753

Reviewed By: diego-urgell

Differential Revision: D55132971

fbshipit-source-id: c1008baf411ad89922d51184bc5ac2951d31704b

Author: galrotem

torchtnt/framework/auto_unit.py

@@ -38,6 +38,7 @@
 from torchtnt.utils.precision import (
     convert_precision_str_to_dtype,
     get_grad_scaler_from_precision,
+    GradScaler,
 )
 from torchtnt.utils.prepare_module import (
     _is_fsdp_module,
@@ -505,7 +506,7 @@ def __init__(
             enable_compiled_autograd=enable_compiled_autograd,
         )
 
-        self.grad_scaler: Optional[torch.amp.GradScaler] = None
+        self.grad_scaler: Optional[GradScaler] = None
         if self.precision:
             self.grad_scaler = get_grad_scaler_from_precision(
                 self.precision,

torchtnt/utils/precision.py

@@ -10,7 +10,12 @@
 from typing import Mapping, Optional
 
 import torch
-from torch.cuda.amp.grad_scaler import GradScaler
+from torch.cuda.amp.grad_scaler import GradScaler as CudaGradScaler
+
+try:
+    from torch.amp.grad_scaler import GradScaler
+except Exception:
+    GradScaler = CudaGradScaler
 
 _DTYPE_STRING_TO_DTYPE_MAPPING: Mapping[str, Optional[torch.dtype]] = {
     "fp16": torch.float16,
@@ -39,7 +44,7 @@ def convert_precision_str_to_dtype(precision: str) -> Optional[torch.dtype]:
 
 def get_grad_scaler_from_precision(
     precision: torch.dtype, *, is_fsdp_module: Optional[bool] = False
-) -> Optional[torch.amp.GradScaler]:
+) -> Optional[GradScaler]:
     """
     Returns the correct grad scaler to use based on the precision and whether
     or not the model is FSDP.
@@ -58,5 +63,5 @@ def get_grad_scaler_from_precision(
 
             return ShardedGradScaler()
         else:
-            return GradScaler()
+            return CudaGradScaler()
     return None