[#168 fix] add context manager to fake ScalingTensor/ScalingParameter's __class__ as torch.Tensor

examples/mnist.py

@@ -70,7 +70,7 @@ def train(args, model, device, train_loader, optimizer, epoch):
             output = model(data)
         loss = F.nll_loss(output, target)
         scaler.scale(loss).backward()
-        scaler.step(optimizer)
+        with msamp.common.tensor.tensor.pretend_scaling_is_torch(): scaler.step(optimizer)
         scaler.update()
         if batch_idx % args.log_interval == 0:
             print(

examples/mnist_ddp.py

@@ -74,7 +74,7 @@ def train(args, model, device, train_loader, optimizer, epoch):
             output = model(data)
         loss = F.nll_loss(output, target)
         scaler.scale(loss).backward()
-        scaler.step(optimizer)
+        with msamp.common.tensor.tensor.pretend_scaling_is_torch(): scaler.step(optimizer)
         scaler.update()
         if dist.get_rank() == 0:
             if batch_idx % args.log_interval == 0:

msamp/common/tensor/tensor.py

@@ -3,6 +3,7 @@
 
 """MS-AMP tensor module."""
 
+from contextlib import contextmanager
 import torch
 import torch.nn.functional as F
 from msamp.common.tensor import ScalingMeta
@@ -12,7 +13,16 @@
 from msamp.common.utils import TransformerEngineWrapper
 
 
+should_pretend_to_be_tt = False
+@contextmanager
+def pretend_scaling_is_torch():
+    global should_pretend_to_be_tt
+    should_pretend_to_be_tt = True
+    yield
+    should_pretend_to_be_tt = False
 class ScalingTensor:
+    @property
+    def __class__(self): return torch.Tensor if should_pretend_to_be_tt else ScalingTensor
     """Customized tensor with scaling."""
     class UniqueDtypeDecorator:
         """A decorator class to check whether dtype is supported and parameters are uniqie."""

msamp/nn/parameter.py

@@ -3,11 +3,15 @@
 
 """MS-AMP parameter module."""
 
+import torch
 from msamp.common.tensor import ScalingTensor
+import msamp.common.tensor.tensor as tensor_py
 
 
 class ScalingParameter(ScalingTensor):
     """Parameter class for ScalingTensor."""
+    @property
+    def __class__(self): return torch.Tensor if tensor_py.should_pretend_to_be_tt else ScalingParameter
     def __init__(self, tensor, requires_grad=True):
         """Constructor.