From 812f6441c7e3ee28df571fabdffd5459e6007c0f Mon Sep 17 00:00:00 2001 From: Vinh Nguyen Date: Mon, 22 Jul 2019 16:10:29 +1000 Subject: [PATCH 1/6] adding GPU auto mixed precision training --- tensor2tensor/utils/optimize.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py index ae962c32f..e7b14f3af 100644 --- a/tensor2tensor/utils/optimize.py +++ b/tensor2tensor/utils/optimize.py @@ -18,6 +18,7 @@ from __future__ import division from __future__ import print_function import numpy as np +import os from tensor2tensor.layers import common_layers from tensor2tensor.utils import adafactor as adafactor_lib @@ -40,7 +41,7 @@ def _mixed_precision_is_enabled(hparams): return activation_dtype == tf.float16 and weight_dtype == tf.float32 -def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None): +def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None, gpu_auto_mixed_precision=False): """Minimize loss.""" loss = weight_decay_and_noise(loss, hparams, learning_rate) loss = tf.identity(loss, name="total_loss") @@ -65,6 +66,12 @@ def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None): opt = ConditionalOptimizer(hparams.optimizer, learning_rate, hparams, use_tpu) if use_tpu: opt = tf.contrib.tpu.CrossShardOptimizer(opt) + if os.environ.get('GPU_AUTO_MIXED_PRECISION', default=False) or gpu_auto_mixed_precision: + if use_tpu: + raise(RuntimeError("GPU auto mixed precision cannot be used with TPU")) + else: + opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt) + opt_summaries = [] if common_layers.should_generate_summaries(): tf.summary.scalar("learning_rate", learning_rate) From e31d47609b5089fe84f3bcadcf07edc838561ac4 Mon Sep 17 00:00:00 2001 From: Vinh Nguyen Date: Tue, 23 Jul 2019 10:33:44 +1000 Subject: [PATCH 2/6] add warning: manual vs. amp --- tensor2tensor/utils/optimize.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py index e7b14f3af..355a7db7d 100644 --- a/tensor2tensor/utils/optimize.py +++ b/tensor2tensor/utils/optimize.py @@ -69,6 +69,8 @@ def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None, gpu_au if os.environ.get('GPU_AUTO_MIXED_PRECISION', default=False) or gpu_auto_mixed_precision: if use_tpu: raise(RuntimeError("GPU auto mixed precision cannot be used with TPU")) + elif _mixed_precision_is_enabled(hparams): + raise(RuntimeError("GPU auto mixed precision cannot be used with manual mixed precision")) else: opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt) From d028d0c0fa5210b46d5be838cf400bf46bf358f1 Mon Sep 17 00:00:00 2001 From: Vinh Nguyen Date: Tue, 23 Jul 2019 11:00:43 +1000 Subject: [PATCH 3/6] fix os environment check --- tensor2tensor/utils/optimize.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py index 355a7db7d..bbfb6f4b0 100644 --- a/tensor2tensor/utils/optimize.py +++ b/tensor2tensor/utils/optimize.py @@ -66,7 +66,7 @@ def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None, gpu_au opt = ConditionalOptimizer(hparams.optimizer, learning_rate, hparams, use_tpu) if use_tpu: opt = tf.contrib.tpu.CrossShardOptimizer(opt) - if os.environ.get('GPU_AUTO_MIXED_PRECISION', default=False) or gpu_auto_mixed_precision: + if os.environ.get('GPU_AUTO_MIXED_PRECISION', default='0') == '1' or gpu_auto_mixed_precision: if use_tpu: raise(RuntimeError("GPU auto mixed precision cannot be used with TPU")) elif _mixed_precision_is_enabled(hparams): From 8687315f65c4f3d6c67c1b609eac40a6b2e4c99f Mon Sep 17 00:00:00 2001 From: Vinh Nguyen Date: Tue, 23 Jul 2019 13:05:41 +1000 Subject: [PATCH 4/6] add _use_locking attribute to optimizer --- tensor2tensor/utils/optimize.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py index bbfb6f4b0..43d6ad4bc 100644 --- a/tensor2tensor/utils/optimize.py +++ b/tensor2tensor/utils/optimize.py @@ -72,6 +72,7 @@ def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None, gpu_au elif _mixed_precision_is_enabled(hparams): raise(RuntimeError("GPU auto mixed precision cannot be used with manual mixed precision")) else: + setattr(opt, '_use_locking', 'True') opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt) opt_summaries = [] From 752220e4fbc8cd3668576dc8e7d28bc7ca4bb36e Mon Sep 17 00:00:00 2001 From: Vinh Nguyen Date: Tue, 23 Jul 2019 13:39:25 +1000 Subject: [PATCH 5/6] add _name attribute to optimizer --- tensor2tensor/utils/optimize.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py index 43d6ad4bc..b1c1431af 100644 --- a/tensor2tensor/utils/optimize.py +++ b/tensor2tensor/utils/optimize.py @@ -73,6 +73,7 @@ def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None, gpu_au raise(RuntimeError("GPU auto mixed precision cannot be used with manual mixed precision")) else: setattr(opt, '_use_locking', 'True') + setattr(opt, '_name', 'ConditionalOptimizer') opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt) opt_summaries = [] From 2a5d1d2661844dab9a4ba219c8b41c6dcc5f66e5 Mon Sep 17 00:00:00 2001 From: Vinh Nguyen Date: Tue, 23 Jul 2019 14:26:59 +1000 Subject: [PATCH 6/6] change OS flag to TF_ENABLE_AUTO_MIXED_PRECISION --- tensor2tensor/utils/optimize.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensor2tensor/utils/optimize.py b/tensor2tensor/utils/optimize.py index b1c1431af..fa620da3b 100644 --- a/tensor2tensor/utils/optimize.py +++ b/tensor2tensor/utils/optimize.py @@ -66,7 +66,7 @@ def optimize(loss, learning_rate, hparams, use_tpu=False, variables=None, gpu_au opt = ConditionalOptimizer(hparams.optimizer, learning_rate, hparams, use_tpu) if use_tpu: opt = tf.contrib.tpu.CrossShardOptimizer(opt) - if os.environ.get('GPU_AUTO_MIXED_PRECISION', default='0') == '1' or gpu_auto_mixed_precision: + if os.environ.get('TF_ENABLE_AUTO_MIXED_PRECISION', default='0') == '1' or gpu_auto_mixed_precision: if use_tpu: raise(RuntimeError("GPU auto mixed precision cannot be used with TPU")) elif _mixed_precision_is_enabled(hparams):