Static grad scaler

python/oneflow/amp/__init__.py

@@ -14,3 +14,4 @@
 limitations under the License.
 """
 from .grad_scaler import GradScaler
+from .grad_scaler import StaticGradScaler

python/oneflow/amp/grad_scaler.py

@@ -43,3 +43,14 @@ def generate_conf_for_graph(self, train_conf):
         train_conf.mutable_dynamic_loss_scale_policy().set_multiplier(
             self._growth_factor
         )
+
+
+class StaticGradScaler(object):
+    def __init__(self, scale_factor):
+        if scale_factor <= 0.0:
+            raise ValueError("scale_factor must > 0.0")
+
+        self.scale_factor = scale_factor
+
+    def generate_conf_for_graph(self, train_conf):
+        train_conf.set_loss_scale_factor(self.scale_factor)

python/oneflow/nn/graph/graph.py

@@ -29,7 +29,7 @@
 from oneflow.nn.graph.block import Block, BlockType
 from oneflow.nn.graph.config import GraphConfig
 from oneflow.nn.graph.optimizer import OptDict, VariableConfig
-from oneflow.amp import GradScaler
+from oneflow.amp import GradScaler, StaticGradScaler
 from oneflow.nn.graph.util import add_indent, sys_exc_error_msg, list_to_func_return
 from oneflow.nn.module import Module
 from oneflow.nn.optimizer.optimizer import Optimizer
@@ -220,7 +220,7 @@ def add_optimizer(
     def set_grad_scaler(self, grad_scaler: GradScaler = None):
         r"""Set the GradScaler for gradient and loss scaling.
         """
-        assert isinstance(grad_scaler, GradScaler)
+        assert isinstance(grad_scaler, (GradScaler, StaticGradScaler))
         self._grad_scaler = grad_scaler
 
     def __call__(self, *args):