Repository: OpenNMT/OpenNMT-tf

Issue:
Multi-GPU synchronous training
Distributed training is currently asynchronous. It should also support a local multi-GPU synchronous mode.

Discussion:

Patch:
diff --git a/README.md b/README.md
index a37f315e..7a6dfd11 100644
--- a/README.md
+++ b/README.md
@@ -25,7 +25,7 @@ and all of the above can be used simultaneously to train novel and complex archi
 
 OpenNMT-tf is also compatible with some of the best TensorFlow features:
 
-* asynchronous distributed training
+* replicated and distributed training
 * monitoring with [TensorBoard](https://www.tensorflow.org/get_started/summaries_and_tensorboard)
 * inference with [TensorFlow Serving](https://www.tensorflow.org/serving/) and the TensorFlow C++ API.
 
diff --git a/bin/average_checkpoints.py b/bin/average_checkpoints.py
index 26faed37..97d3b473 100644
--- a/bin/average_checkpoints.py
+++ b/bin/average_checkpoints.py
@@ -34,7 +34,7 @@ def main():
                       help="The model directory containing the checkpoints.")
   parser.add_argument("--output_dir", required=True,
                       help="The output directory where the averaged checkpoint will be saved.")
-  parser.add_argument("--max_count", default=8,
+  parser.add_argument("--max_count", type=int, default=8,
                       help="The maximal number of checkpoints to average.")
   args = parser.parse_args()
 
diff --git a/bin/main.py b/bin/main.py
index 9c128bb0..52792cd1 100644
--- a/bin/main.py
+++ b/bin/main.py
@@ -74,13 +74,14 @@ def load_model(model_dir, model_file=None):
 
   return model
 
-def train(estimator, model, config):
+def train(estimator, model, config, num_devices=1):
   """Runs training.
 
   Args:
     estimator: A `tf.estimator.Estimator`.
     model: A `opennmt.models.Model`.
     config: The configuration.
+    num_devices: The number of devices used for training.
   """
   if "eval" not in config:
     config["eval"] = {}
@@ -132,6 +133,7 @@ def train(estimator, model, config):
           config["data"]["train_features_file"],
           labels_file=config["data"]["train_labels_file"],
           batch_type=train_batch_type,
+          batch_multiplier=num_devices,
           bucket_width=config["train"].get("bucket_width", 5),
           sample_buffer_size=config["train"].get(
               "sample_buffer_size", default_sample_buffer_size),
@@ -241,6 +243,8 @@ def main():
   parser.add_argument("--checkpoint_path", default=None,
                       help=("Checkpoint or directory to use for inference or export "
                             "(when a directory is set, the latest checkpoint is used)."))
+  parser.add_argument("--num_gpus", type=int, default=1,
+                      help="Number of GPUs to use for in-graph replication.")
   parser.add_argument("--chief_host", default="",
                       help="hostname:port of the chief worker (for distributed training).")
   parser.add_argument("--worker_hosts", default="",
@@ -286,8 +290,11 @@ def main():
     tf.logging.info("Creating model directory %s", config["model_dir"])
     os.makedirs(config["model_dir"])
 
-  session_config = tf.ConfigProto()
-  session_config.gpu_options.allow_growth = args.gpu_allow_growth
+  session_config = tf.ConfigProto(
+      allow_soft_placement=True,
+      log_device_placement=False,
+      gpu_options=tf.GPUOptions(
+          allow_growth=args.gpu_allow_growth))
 
   run_config = tf.estimator.RunConfig(
       model_dir=config["model_dir"],
@@ -309,7 +316,7 @@ def main():
   model = load_model(config["model_dir"], model_file=args.model)
 
   estimator = tf.estimator.Estimator(
-      model,
+      model.model_fn(num_devices=args.num_gpus),
       config=run_config,
       params=config["params"])
 
@@ -320,7 +327,7 @@ def main():
   if args.run == "train":
     if args.data_dir:
       config["data"] = _prefix_paths(args.data_dir, config["data"])
-    train(estimator, model, config)
+    train(estimator, model, config, num_devices=args.num_gpus)
   elif args.run == "infer":
     if not args.features_file:
       parser.error("--features_file is required for inference.")
diff --git a/docs/package/opennmt.utils.parallel.rst b/docs/package/opennmt.utils.parallel.rst
new file mode 100644
index 00000000..c85542c7
--- /dev/null
+++ b/docs/package/opennmt.utils.parallel.rst
@@ -0,0 +1,7 @@
+opennmt\.utils\.parallel module
+===============================
+
+.. automodule:: opennmt.utils.parallel
+    :members:
+    :undoc-members:
+    :show-inheritance:
diff --git a/docs/package/opennmt.utils.rst b/docs/package/opennmt.utils.rst
index f82c1cf3..123c73cf 100644
--- a/docs/package/opennmt.utils.rst
+++ b/docs/package/opennmt.utils.rst
@@ -20,6 +20,7 @@ Submodules
    opennmt.utils.losses
    opennmt.utils.misc
    opennmt.utils.optim
+   opennmt.utils.parallel
    opennmt.utils.position
    opennmt.utils.reducer
    opennmt.utils.transformer
diff --git a/docs/training.md b/docs/training.md
index 9b21e822..3558687c 100644
--- a/docs/training.md
+++ b/docs/training.md
@@ -18,9 +18,23 @@ then open the URL displayed in the shell to monitor and visualize several data,
 * word embeddings
 * decoder sampling probability
 
-## Distributed
+## Replicated training
 
-OpenNMT-tf supports asynchronous distributed training. The user should set on the command line:
+OpenNMT-tf training can make use of multiple GPUs with *in-graph replication*. In this mode, the main section of the graph is replicated over multiple devices and batches are processed in parallel. The resulting graph is equivalent to train with batches `N` times larger, where `N` is the number of used GPUs.
+
+For example, if your machine has 4 GPUs, simply add the `--num_gpus` option:
+
+```bash
+python -m bin.main train [...] --num_gpus 4
+```
+
+Note that evaluation and inference will run on a single device.
+
+## Distributed training
+
+OpenNMT-tf also supports asynchronous distributed training with *between-graph replication*. In this mode, each graph replica processes a batch independently, compute the gradients, and asynchronously update a shared set of parameters.
+
+To enable distributed training, the user should set on the command line:
 
 * a **chief worker** host that runs a training loop and manages checkpoints, summaries, etc.
 * a list of **worker** hosts that run a training loop
@@ -32,6 +46,6 @@ Then a training instance should be started on each host with a selected task, e.
 CUDA_VISIBLE_DEVICES=0 python -m bin.main train [...] --ps_hosts localhost:2222 --chief_host localhost:2223 --worker_hosts localhost:2224,localhost:2225 --task_type worker --task_index 1
 ```
 
-will start the worker 1 on the current machine and first GPU.
+will start the worker 1 on the current machine and first GPU. By setting `CUDA_VISIBLE_DEVICES` correctly, asynchronous distributed training can be run on a single multi-GPU machine.
 
-For more details, see the documentation of [`tf.estimator.train_and_evaluate`](https://www.tensorflow.org/versions/r1.4/api_docs/python/tf/estimator/train_and_evaluate). Also see [tensorflow/ecosystem](https://github.com/tensorflow/ecosystem) to integrate distributed training with open-source frameworks like Docker or Kubernetes.
+For more details, see the documentation of [`tf.estimator.train_and_evaluate`](https://www.tensorflow.org/api_docs/python/tf/estimator/train_and_evaluate). Also see [tensorflow/ecosystem](https://github.com/tensorflow/ecosystem) to integrate distributed training with open-source frameworks like Docker or Kubernetes.
diff --git a/opennmt/decoders/self_attention_decoder.py b/opennmt/decoders/self_attention_decoder.py
index c45ecbba..b5d37865 100644
--- a/opennmt/decoders/self_attention_decoder.py
+++ b/opennmt/decoders/self_attention_decoder.py
@@ -46,26 +46,42 @@ class SelfAttentionDecoder(Decoder):
     self.relu_dropout = relu_dropout
     self.position_encoder = position_encoder
 
-  def _init_cache(self, memory, memory_sequence_length):
+  def _build_memory_mask(self, memory, memory_sequence_length=None):
+    if memory_sequence_length is None:
+      return None
+    else:
+      return transformer.build_sequence_mask(
+          memory_sequence_length,
+          num_heads=self.num_heads,
+          maximum_length=tf.shape(memory)[1],
+          dtype=memory.dtype)
+
+  def _init_cache(self, memory, memory_sequence_length=None):
     cache = {
         "memory": memory,
-        "memory_sequence_length": memory_sequence_length
+        "memory_mask": self._build_memory_mask(
+            memory, memory_sequence_length=memory_sequence_length)
     }
 
     batch_size = tf.shape(memory)[0]
     depth = memory.get_shape().as_list()[-1]
 
-    for l in range(self.num_layers):
-      keys = tf.zeros([batch_size, 0, depth])
-      values = tf.zeros([batch_size, 0, depth])
-
-      # Ensure shape invariance for tf.while_loop.
-      keys._shape = tf.TensorShape([None, None, depth])  # pylint: disable=protected-access
-      values._shape = tf.TensorShape([None, None, depth])  # pylint: disable=protected-access
+    def _create_placeholder(shape, loop_shape=None):
+      placeholder = tf.zeros(shape)
+      if loop_shape is not None:
+        placeholder._shape = tf.TensorShape(loop_shape)  # pylint: disable=protected-access
+      return placeholder
 
+    for l in range(self.num_layers):
       cache["layer_{}".format(l)] = {
-          "keys": keys,
-          "values": values
+          "self_keys": _create_placeholder(
+              [batch_size, 0, depth], loop_shape=[None, None, depth]),
+          "self_values": _create_placeholder(
+              [batch_size, 0, depth], loop_shape=[None, None, depth]),
+          "memory_keys": _create_placeholder(
+              [batch_size, 0, depth], loop_shape=[None, None, depth]),
+          "memory_values": _create_placeholder(
+              [batch_size, 0, depth], loop_shape=[None, None, depth]),
       }
 
     return cache
@@ -75,13 +91,14 @@ class SelfAttentionDecoder(Decoder):
 
     def _impl(ids, step, cache):
       inputs = embedding_fn(ids[:, -1:])
+      inputs *= self.num_units**0.5
       inputs = self.position_encoder.apply_one(inputs, step + 1)
       outputs = self._self_attention_stack(
           inputs,
           mode=mode,
           cache=cache,
           memory=cache["memory"],
-          memory_sequence_length=cache["memory_sequence_length"])
+          memory_sequence_length=None)
       outputs = outputs[:, -1:, :]
       logits = tf.layers.dense(outputs, vocab_size)
       return logits, cache
@@ -109,23 +126,22 @@ class SelfAttentionDecoder(Decoder):
           num_heads=self.num_heads,
           maximum_length=tf.shape(inputs)[1],
           dtype=inputs.dtype)
-    if memory_sequence_length is not None:
-      memory_mask = transformer.build_sequence_mask(
-          memory_sequence_length,
-          num_heads=self.num_heads,
-          maximum_length=tf.shape(memory)[1],
-          dtype=memory.dtype)
+    if memory is not None:
+      if cache is not None:
+        memory_mask = cache["memory_mask"]
+      elif memory_sequence_length is not None:
+        memory_mask = self._build_memory_mask(
+            memory, memory_sequence_length=memory_sequence_length)
 
     for l in range(self.num_layers):
       layer_name = "layer_{}".format(l)
       layer_cache = cache[layer_name] if cache is not None else None
       with tf.variable_scope(layer_name):
         with tf.variable_scope("masked_multi_head"):
-          inputs_norm = transformer.norm(inputs)
           encoded = transformer.multi_head_attention(
               self.num_heads,
-              inputs_norm,
-              inputs_norm,
+              transformer.norm(inputs),
+              None,
               mode,
               num_units=self.num_units,
               mask=decoder_mask,
@@ -145,6 +161,7 @@ class SelfAttentionDecoder(Decoder):
                 memory,
                 mode,
                 mask=memory_mask,
+                cache=layer_cache,
                 dropout=self.attention_dropout)
             context = transformer.drop_and_add(
                 encoded,
@@ -182,6 +199,7 @@ class SelfAttentionDecoder(Decoder):
     if sampling_probability is not None:
       raise ValueError("Scheduled sampling is not supported with SelfAttentionDecoder")
 
+    inputs *= self.num_units**0.5
     if self.position_encoder is not None:
       inputs = self.position_encoder(inputs, sequence_length=sequence_length)
 
@@ -213,7 +231,7 @@ class SelfAttentionDecoder(Decoder):
     inputs = tf.expand_dims(start_tokens, 1)
     lengths = tf.zeros([batch_size], dtype=tf.int32)
     log_probs = tf.zeros([batch_size])
-    cache = self._init_cache(memory, memory_sequence_length)
+    cache = self._init_cache(memory, memory_sequence_length=memory_sequence_length)
 
     symbols_to_logits_fn = self._symbols_to_logits_fn(embedding, vocab_size, mode)
 
@@ -282,7 +300,7 @@ class SelfAttentionDecoder(Decoder):
                                 mode=tf.estimator.ModeKeys.PREDICT,
                                 memory=None,
                                 memory_sequence_length=None):
-    cache = self._init_cache(memory, memory_sequence_length)
+    cache = self._init_cache(memory, memory_sequence_length=memory_sequence_length)
     symbols_to_logits_fn = self._symbols_to_logits_fn(embedding, vocab_size, mode)
 
     outputs, log_probs = beam_search(
diff --git a/opennmt/encoders/self_attention_encoder.py b/opennmt/encoders/self_attention_encoder.py
index a556f366..3f3677fb 100644
--- a/opennmt/encoders/self_attention_encoder.py
+++ b/opennmt/encoders/self_attention_encoder.py
@@ -46,6 +46,7 @@ class SelfAttentionEncoder(Encoder):
     self.position_encoder = position_encoder
 
   def encode(self, inputs, sequence_length=None, mode=tf.estimator.ModeKeys.TRAIN):
+    inputs *= self.num_units**0.5
     if self.position_encoder is not None:
       inputs = self.position_encoder(inputs, sequence_length=sequence_length)
 
@@ -64,11 +65,10 @@ class SelfAttentionEncoder(Encoder):
     for l in range(self.num_layers):
       with tf.variable_scope("layer_{}".format(l)):
         with tf.variable_scope("multi_head"):
-          inputs_norm = transformer.norm(inputs)
           context = transformer.multi_head_attention(
               self.num_heads,
-              inputs_norm,
-              inputs_norm,
+              transformer.norm(inputs),
+              None,
               mode,
               num_units=self.num_units,
               mask=mask,
diff --git a/opennmt/models/model.py b/opennmt/models/model.py
index 39456578..7d36e9af 100644
--- a/opennmt/models/model.py
+++ b/opennmt/models/model.py
@@ -10,6 +10,29 @@ import tensorflow as tf
 
 from opennmt.utils.optim import optimize
 from opennmt.utils.misc import add_dict_to_collection, item_or_tuple
+from opennmt.utils.parallel import GraphDispatcher
+
+
+def filter_irregular_batches(multiple):
+  """Transformation that filters out batches based on their size.
+
+  Args:
+    multiple: The divisor of the batch size.
+
+  Returns:
+    A ``tf.data.Dataset`` transformation.
+  """
+  def _apply_fn(dataset):
+    """Transformation function."""
+
+    def _predicate(*x):
+      flat = tf.contrib.framework.nest.flatten(x)
+      batch_size = tf.shape(flat[0])[0]
+      return tf.equal(tf.mod(batch_size, multiple), 0)
+
+    return dataset.filter(_predicate)
+
+  return _apply_fn
 
 
 @six.add_metaclass(abc.ABCMeta)
@@ -19,48 +42,94 @@ class Model(object):
   def __init__(self, name):
     self.name = name
 
-  def __call__(self, features, labels, params, mode, config):
-    """Creates the model.
+  def model_fn(self, num_devices=1):
+    """Returns the model function.
+
+    Args:
+      num_devices: The number of devices used for training.
 
     See Also:
       ``tf.estimator.Estimator`` 's ``model_fn`` argument for more details about
       arguments and the returned value.
     """
-    if mode == tf.estimator.ModeKeys.TRAIN:
-      self._register_word_counters(features, labels)
-
-    with tf.variable_scope(self.name, initializer=self._initializer(params)) as model_scope:
-      outputs, predictions = self._build(features, labels, params, mode, config)
+    dispatcher = GraphDispatcher(num_devices)
+
+    def _loss_op(features, labels, params, mode, config):
+      """Single callable to compute the loss."""
+      logits, _ = self._build(features, labels, params, mode, config)
+      return self._compute_loss(features, labels, logits, params, mode)
+
+    def _normalize_loss(num, den=None):
+      """Normalizes the loss."""
+      if isinstance(num, list):  # Sharded mode.
+        if den is not None:
+          assert isinstance(den, list)
+          return tf.add_n(num) / tf.add_n(den)
+        else:
+          return tf.reduce_mean(num)
+      elif den is not None:
+        return num / den
+      else:
+        return num
 
-    if predictions is not None:
-      # Register predictions in a collection so that hooks can easily fetch them.
-      add_dict_to_collection("predictions", predictions)
+    def _extract_loss(loss):
+      """Extracts and summarizes the loss."""
+      if not isinstance(loss, tuple):
+        actual_loss = _normalize_loss(loss)
+        tboard_loss = actual_loss
+      else:
+        actual_loss = _normalize_loss(loss[0], den=loss[1])
+        tboard_loss = _normalize_loss(loss[0], den=loss[2]) if len(loss) > 2 else actual_loss
+      tf.summary.scalar("loss", tboard_loss)
+      return actual_loss
 
-    if mode != tf.estimator.ModeKeys.PREDICT:
-      with tf.variable_scope(model_scope):
-        loss = self._compute_loss(features, labels, outputs, params, mode)
+    def _model_fn(features, labels, params, mode, config):
+      """model_fn implementation."""
+      if mode == tf.estimator.ModeKeys.TRAIN:
+        self._register_word_counters(features, labels)
 
-      if isinstance(loss, tuple):
-        loss, display_loss = loss
-      else:
-        display_loss = loss
+        features_shards = dispatcher.shard(features)
+        labels_shards = dispatcher.shard(labels)
 
-      tf.summary.scalar("loss", display_loss)
+        with tf.variable_scope(self.name, initializer=self._initializer(params)):
+          losses_shards = dispatcher(
+              _loss_op, features_shards, labels_shards, params, mode, config)
 
-      if mode == tf.estimator.ModeKeys.TRAIN:
+        loss = _extract_loss(losses_shards)
         train_op = optimize(loss, params)
         return tf.estimator.EstimatorSpec(
-            mode, loss=loss, train_op=train_op)
-      else:
+            mode,
+            loss=loss,
+            train_op=train_op)
+      elif mode == tf.estimator.ModeKeys.EVAL:
+        with tf.variable_scope(self.name):
+          logits, predictions = self._build(features, labels, params, mode, config)
+          loss = self._compute_loss(features, labels, logits, params, mode)
+
+        loss = _extract_loss(loss)
         eval_metric_ops = self._compute_metrics(features, labels, predictions)
+        if predictions is not None:
+          # Register predictions in a collection so that hooks can easily fetch them.
+          add_dict_to_collection("predictions", predictions)
+
         return tf.estimator.EstimatorSpec(
-            mode, loss=loss, eval_metric_ops=eval_metric_ops)
-    else:
-      export_outputs = {}
-      export_outputs[tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = \
-          tf.estimator.export.PredictOutput(predictions)
-      return tf.estimator.EstimatorSpec(
-          mode, predictions=predictions, export_outputs=export_outputs)
+            mode,
+            loss=loss,
+            eval_metric_ops=eval_metric_ops)
+      elif mode == tf.estimator.ModeKeys.PREDICT:
+        with tf.variable_scope(self.name):
+          _, predictions = self._build(features, labels, params, mode, config)
+
+        export_outputs = {}
+        export_outputs[tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = (
+            tf.estimator.export.PredictOutput(predictions))
+
+        return tf.estimator.EstimatorSpec(
+            mode,
+            predictions=predictions,
+            export_outputs=export_outputs)
+
+    return _model_fn
 
   def _initializer(self, params):
     """Returns the global initializer for this model.
@@ -268,6 +337,7 @@ class Model(object):
                      features_file,
                      labels_file=None,
                      batch_type="examples",
+                     batch_multiplier=1,
                      bucket_width=None,
                      sample_buffer_size=None,
                      maximum_features_length=None,
@@ -306,6 +376,7 @@ class Model(object):
           labels,
           maximum_features_length=maximum_features_length,
           maximum_labels_length=maximum_labels_length))
+      batch_size = batch_size * batch_multiplier
 
     if mode == tf.estimator.ModeKeys.TRAIN and bucket_width is not None:
       # Form batches with sequences of similar lengths to improve efficiency.
@@ -333,7 +404,10 @@ class Model(object):
         if bucket_width > 1:
           key += 1  # For bucket_width == 1, key 0 is unassigned.
         size = batch_size // (key * bucket_width)
-        return tf.to_int64(tf.maximum(size, 1))
+        if batch_multiplier > 1:
+          # Make the window size a multiple of batch_multiplier.
+          size = size + batch_multiplier - size % batch_multiplier
+        return tf.to_int64(tf.maximum(size, batch_multiplier))
 
       if batch_type == "examples":
         batchify_fn = tf.contrib.data.group_by_window(
@@ -352,6 +426,8 @@ class Model(object):
           padded_shapes=padded_shapes)
 
     if mode == tf.estimator.ModeKeys.TRAIN:
+      if batch_multiplier > 1:
+        dataset = dataset.apply(filter_irregular_batches(batch_multiplier))
       dataset = dataset.repeat()
 
     iterator = dataset.make_initializable_iterator()
@@ -370,6 +446,7 @@ class Model(object):
                features_file,
                labels_file=None,
                batch_type="examples",
+               batch_multiplier=1,
                bucket_width=None,
                sample_buffer_size=None,
                maximum_features_length=None,
@@ -387,6 +464,8 @@ class Model(object):
       labels_file: The file containing output labels.
       batch_type: The training batching stragety to use: can be "examples" or
         "tokens".
+      batch_multiplier: The batch size multiplier to prepare splitting accross
+         replicated graph parts.
       bucket_width: The width of the length buckets to select batch candidates
         from. ``None`` to not constrain batch formation.
       sample_buffer_size: The number of elements from which to sample.
@@ -417,6 +496,7 @@ class Model(object):
         features_file,
         labels_file=labels_file,
         batch_type=batch_type,
+        batch_multiplier=batch_multiplier,
         bucket_width=bucket_width,
         sample_buffer_size=sample_buffer_size,
         maximum_features_length=maximum_features_length,
diff --git a/opennmt/models/sequence_tagger.py b/opennmt/models/sequence_tagger.py
index 4a3e7a75..20155bd5 100644
--- a/opennmt/models/sequence_tagger.py
+++ b/opennmt/models/sequence_tagger.py
@@ -134,7 +134,9 @@ class SequenceTagger(Model):
             outputs,
             tf.cast(labels["tags_id"], tf.int32),
             length)
-      return tf.reduce_mean(-log_likelihood)
+      loss = tf.reduce_sum(-log_likelihood)
+      loss_normalizer = tf.shape(log_likelihood)[0]
+      return loss, loss_normalizer
     else:
       return cross_entropy_sequence_loss(
           outputs,
diff --git a/opennmt/utils/losses.py b/opennmt/utils/losses.py
index a44c02df..5d7efbdd 100644
--- a/opennmt/utils/losses.py
+++ b/opennmt/utils/losses.py
@@ -26,7 +26,7 @@ def cross_entropy_sequence_loss(logits,
                                 label_smoothing=0.0,
                                 average_in_time=False,
                                 mode=tf.estimator.ModeKeys.TRAIN):
-  """Computes the reduced cross entropy loss of sequences.
+  """Computes the cross entropy loss of sequences.
 
   Args:
     logits: The unscaled probabilities.
@@ -37,7 +37,7 @@ def cross_entropy_sequence_loss(logits,
     mode: A ``tf.estimator.ModeKeys`` mode.
 
   Returns:
-    A tuple with the loss and the token level loss.
+    A tuple (cumulated loss, loss normalizer, token-level normalizer).
   """
   batch_size = tf.shape(logits)[0]
   max_time = tf.shape(logits)[1]
@@ -46,18 +46,20 @@ def cross_entropy_sequence_loss(logits,
   weights = tf.sequence_mask(
       sequence_length, maxlen=max_time, dtype=cross_entropy.dtype)
   loss = tf.reduce_sum(cross_entropy * weights)
-  loss_per_token = loss / tf.reduce_sum(weights)
+  loss_token_normalizer = tf.reduce_sum(weights)
 
   if average_in_time or mode != tf.estimator.ModeKeys.TRAIN:
-    return loss_per_token, loss_per_token
+    loss_normalizer = loss_token_normalizer
   else:
-    return loss / tf.cast(batch_size, loss.dtype), loss_per_token
+    loss_normalizer = tf.cast(batch_size, loss.dtype)
+
+  return loss, loss_normalizer, loss_token_normalizer
 
 def cross_entropy_loss(logits,
                        labels,
                        label_smoothing=0.0,
                        mode=tf.estimator.ModeKeys.TRAIN):
-  """Computes the reduced cross entropy loss.
+  """Computes the cross entropy loss.
 
   Args:
     logits: The unscaled probabilities.
@@ -66,8 +68,9 @@ def cross_entropy_loss(logits,
     mode: A ``tf.estimator.ModeKeys`` mode.
 
   Returns:
-    The loss.
+    The cumulated loss and the loss normalizer.
   """
   cross_entropy = _softmax_cross_entropy(logits, labels, label_smoothing, mode)
-  loss = tf.reduce_mean(cross_entropy)
-  return loss
+  loss = tf.reduce_sum(cross_entropy)
+  loss_normalizer = tf.shape(cross_entropy)[0]
+  return loss, loss_normalizer
diff --git a/opennmt/utils/optim.py b/opennmt/utils/optim.py
index 75fe515d..d41f9632 100644
--- a/opennmt/utils/optim.py
+++ b/opennmt/utils/optim.py
@@ -117,4 +117,5 @@ def optimize(loss, params):
       summaries=[
           "learning_rate",
           "global_gradient_norm",
-      ])
+      ],
+      colocate_gradients_with_ops=True)
diff --git a/opennmt/utils/parallel.py b/opennmt/utils/parallel.py
new file mode 100644
index 00000000..97db3cd5
--- /dev/null
+++ b/opennmt/utils/parallel.py
@@ -0,0 +1,175 @@
+"""Utilities to run execution in parallel."""
+
+import six
+
+import tensorflow as tf
+
+from tensorflow.python.client import device_lib
+
+
+class GraphDispatcher(object):
+  """Helper class to replicate graph parts on multiple devices and dispatch
+  sharded batches.
+  """
+
+  def __init__(self, num_devices):
+    """Initializes the dispatcher.
+
+    Args:
+      num_devices: The number of devices to dispatch on.
+
+    Raises:
+      ValueError: if the number of visible devices is lower than
+        :obj:`num_devices`.
+    """
+    devices = [x.name for x in device_lib.list_local_devices() if x.device_type == "GPU"]
+
+    if not devices:
+      self._n = 1
+      self._devices = [None]
+    elif len(devices) < num_devices:
+      raise ValueError("Only %d devices are visible but %d were requested"
+                       % (len(devices), num_replicas))
+    else:
+      self._n = num_devices
+      self._devices = devices[:self._n]
+
+  def shard(self, data):
+    """Shards a structure of ``tf.Tensor`` for dispatching.
+
+    Args:
+      data: A ``tf.Tensor`` of dictionary of ``tf.Tensor``.
+
+    Returns:
+      A list of the same ``tf.Tensor`` structure.
+    """
+    return split_batch(data, self._n)
+
+  def repeat(self, data):
+    """Ensures that the object is dispatchable list.
+
+    Args:
+      data: The object to convert.
+
+    Returns:
+      :obj:`data` if it is valid list or a list where :obj:`data` is replicated.
+
+    Raises:
+      ValueError: if :obj:`data` is a non dispatchable list.
+    """
+    if isinstance(data, list):
+      if len(data) != self._n:
+        raise ValueError("List arguments must contain %d elements, saw %d instead"
+                         % (self._n, len(data)))
+      return data
+    else:
+      return [data] * self._n
+
+  def _parallel_args(self, *args, **kwargs):
+    """Makes each argument dispatchable."""
+    if args:
+      parallel_args = [self.repeat(arg) for arg in args]
+      parallel_args = [list(arg) for arg in zip(*parallel_args)]
+    else:
+      parallel_args = [[] for _ in range(self._n)]
+    parallel_kwargs = [{} for _ in range(self._n)]
+    for k, v in six.iteritems(kwargs):
+      values = self.repeat(v)
+      for i in range(self._n):
+        parallel_kwargs[i][k] = values[i]
+    return parallel_args, parallel_kwargs
+
+  def __call__(self, fun, *args, **kwargs):
+    """Dispatches :obj:`fun` calls accross devices.
+
+    * Each argument must either not be a list or a list with length the
+      number of devices used for dispatching.
+    * Variables are copied in a daisy chain fashion between devices
+      (credits to Tensor2Tensor).
+
+    Args:
+      fun: A callable.
+      *args: The callable arguments.
+      **kwargs: The callable keyword arguments.
+
+    Returns:
+      The sharded outputs of :obj:`fun`.
+    """
+    funs = self.repeat(fun)
+    args, kwargs = self._parallel_args(*args, **kwargs)
+
+    outputs = []
+    cache = {}
+    tensor_to_var = {}
+
+    for i, device in enumerate(self._devices):
+
+      # pylint: disable=cell-var-from-loop
+      def _daisy_chain_getter(getter, name, *args, **kwargs):
+        """Get a variable and cache in a daisy chain."""
+        # Copyright 2017 The Tensor2Tensor Authors.
+        # Licensed under the Apache License, Version 2.0
+        device_var_key = (device, name)
+        if device_var_key in cache:
+          # if we have the variable on the correct device, return it.
+          return cache[device_var_key]
+        if name in cache:
+          # if we have it on a different device, copy it from the last device
+          last_device_v = cache[name]
+          var = tensor_to_var[last_device_v]
+          v = tf.identity(last_device_v)
+        else:
+          var = getter(name, *args, **kwargs)
+          v = tf.identity(var._ref())  # pylint: disable=protected-access
+
+        # keep track of the original variable
+        tensor_to_var[v] = var
+        v.read_value = lambda: tf.identity(v)
+        v.assign_sub = var.assign_sub
+        # update the cache
+        cache[name] = v
+        cache[device_var_key] = v
+        return v
+
+      with tf.name_scope("parallel_{}".format(i)):
+        with tf.variable_scope(
+            tf.get_variable_scope(),
+            reuse=True if i > 0 else None,
+            custom_getter=_daisy_chain_getter):
+          if device is None:
+            outputs.append(funs[i](*args[i], **kwargs[i]))
+          else:
+            with tf.device(device):
+              outputs.append(funs[i](*args[i], **kwargs[i]))
+
+    # If the function returned a tuple, also return a tuple of sharded results.
+    if isinstance(outputs[0], tuple):
+      outputs = tuple(list(output) for output in zip(*outputs))
+
+    return outputs
+
+
+def split_batch(data, num_shards):
+  """Split data into shards."""
+
+  def _split_dictionary(dictionary):
+    """Split a dictionary into shards."""
+    shards = [{} for _ in range(num_shards)]
+    for name, tensor in six.iteritems(dictionary):
+      if isinstance(tensor, tf.SparseTensor):
+        for i, shard in enumerate(tf.sparse_split(sp_input=tensor, num_split=num_shards, axis=0)):
+          shards[i][name] = shard
+      else:
+        for i, shard in enumerate(tf.split(tensor, num_shards)):
+          shards[i][name] = shard
+    return shards
+
+  with tf.name_scope("split_inputs"):
+    if data is None:
+      data_shards = None
+    elif isinstance(data, dict):
+      data_shards = _split_dictionary(data)
+    else:
+      data_shards = tf.split(data, num_shards)
+
+  return data_shards
diff --git a/opennmt/utils/transformer.py b/opennmt/utils/transformer.py
index 3fe296e9..961ec3ee 100644
--- a/opennmt/utils/transformer.py
+++ b/opennmt/utils/transformer.py
@@ -76,6 +76,20 @@ def build_future_mask(sequence_length,
     mask = tf.reshape(mask, [-1, num_heads, tf.shape(mask)[1], tf.shape(mask)[2]])
   return mask
 
+def fused_projection(inputs, num_units, num_outputs=1):
+  """Projects the same input into multiple output spaces.
+
+  Args:
+    inputs: The inputs to project.
+    num_units: The number of output units of each space.
+    num_outputs: The number of output spaces.
+
+  Returns:
+    :obj:`num_outputs` ``tf.Tensor`` of depth :obj:`num_units`.
+  """
+  return tf.split(
+      tf.layers.conv1d(inputs, num_units * num_outputs, 1), num_outputs, axis=2)
+
 def split_heads(inputs, num_heads):
   """Splits a tensor in depth.
 
@@ -103,14 +117,13 @@ def combine_heads(inputs):
   inputs = tf.reshape(inputs, [tf.shape(inputs)[0], tf.shape(inputs)[1], -1])
   return inputs
 
-def scaled_dot_attention(queries,
-                         keys,
-                         values,
-                         mode,
-                         mask=None,
-                         dropout=0.0):
-  """Computes the scaled dot-product attention as described
-  in https://arxiv.org/abs/1706.03762.
+def dot_product_attention(queries,
+                          keys,
+                          values,
+                          mode,
+                          mask=None,
+                          dropout=0.0):
+  """Computes the dot product attention.
 
   Args:
     queries: The sequence of queries. A tensor of shape :math:`[B, T_1, ...]`.
@@ -124,9 +137,8 @@ def scaled_dot_attention(queries,
   Returns:
     A tuple ``(context vector, attention vector)``.
   """
-  # Scaled dot-product between queries and keys.
+  # Dot product between queries and keys.
   dot = tf.matmul(queries, keys, transpose_b=True)
-  dot = tf.div(dot, tf.sqrt(tf.cast(tf.shape(keys)[-1], dot.dtype)))
 
   if mask is not None:
     dot = dot * mask + ((1.0 - mask) * dot.dtype.min)
@@ -159,6 +171,7 @@ def multi_head_attention(num_heads,
     num_heads: The number of attention heads.
     queries: The sequence of queries. A tensor of shape :math:`[B, T_1, ...]`.
     memory: The sequence to attend. A tensor of shape :math:`[B, T_2, ...]`.
+      If ``None``, computes self-attention.
     mode: A ``tf.estimator.ModeKeys`` mode.
     num_units: The number of hidden units. If not set, it is set to the input
       dimension.
@@ -175,21 +188,34 @@ def multi_head_attention(num_heads,
     raise ValueError("Multi head attention requires that num_units is a"
                      " multiple of {}".format(num_heads))
 
-  queries = tf.layers.conv1d(queries, num_units, 1)
-  memory = tf.layers.conv1d(memory, num_units * 2, 1)
-  keys, values = tf.split(memory, [num_units, num_units], axis=2)
-
-  if cache is not None:
-    keys = tf.concat([cache["keys"], keys], axis=1)
-    values = tf.concat([cache["values"], values], axis=1)
-    cache["keys"] = keys
-    cache["values"] = values
+  if memory is None:
+    queries, keys, values = fused_projection(queries, num_units, num_outputs=3)
+
+    if cache is not None:
+      keys = tf.concat([cache["self_keys"], keys], axis=1)
+      values = tf.concat([cache["self_values"], values], axis=1)
+      cache["self_keys"] = keys
+      cache["self_values"] = values
+  else:
+    queries = tf.layers.conv1d(queries, num_units, 1)
+
+    if cache is not None:
+      keys, values = tf.cond(
+          tf.equal(tf.shape(cache["memory_keys"])[1], 0),
+          true_fn=lambda: fused_projection(memory, num_units, num_outputs=2),
+          false_fn=lambda: [cache["memory_keys"], cache["memory_values"]])
+      cache["memory_keys"] = keys
+      cache["memory_values"] = values
+    else:
+      keys, values = fused_projection(memory, num_units, num_outputs=2)
 
   queries = split_heads(queries, num_heads)
   keys = split_heads(keys, num_heads)
   values = split_heads(values, num_heads)
 
-  heads, _ = scaled_dot_attention(
+  queries *= (num_units // num_heads)**-0.5
+
+  heads, _ = dot_product_attention(
       queries,
       keys,
       values,


Test patch:
diff --git a/opennmt/tests/transformer_test.py b/opennmt/tests/transformer_test.py
index 079b35df..4f8182a4 100644
--- a/opennmt/tests/transformer_test.py
+++ b/opennmt/tests/transformer_test.py
@@ -121,7 +121,7 @@ class TransformerTest(tf.test.TestCase):
     keys = values
 
     mask = transformer.build_sequence_mask(values_length, num_heads=num_heads)
-    context, attn = transformer.scaled_dot_attention(
+    context, attn = transformer.dot_product_attention(
         queries,
         keys,
         values,
@@ -152,7 +152,7 @@ class TransformerTest(tf.test.TestCase):
         shape=(None, num_heads, None, depth))
 
     mask = transformer.build_future_mask(queries_length, num_heads=num_heads)
-    context, attn = transformer.scaled_dot_attention(
+    context, attn = transformer.dot_product_attention(
         queries,
         queries,
         queries,