Internal change

PiperOrigin-RevId: 339901099

official/nlp/tasks/masked_lm.py

@@ -36,6 +36,9 @@ class MaskedLMConfig(cfg.TaskConfig):
       bert.ClsHeadConfig(
           inner_dim=768, num_classes=2, dropout_rate=0.1, name='next_sentence')
   ])
+  # TODO(b/154564893): Mathematically, scale_loss should be True.
+  # However, it works better with scale_loss being False.
+  scale_loss: bool = False
   train_data: cfg.DataConfig = cfg.DataConfig()
   validation_data: cfg.DataConfig = cfg.DataConfig()
 
@@ -161,11 +164,14 @@ class MaskedLMTask(base_task.Task):
           model_outputs=outputs,
           metrics=metrics,
           aux_losses=model.losses)
+      if self.task_config.scale_loss:
         # Scales loss as the default gradients allreduce performs sum inside the
         # optimizer.
-      # TODO(b/154564893): enable loss scaling.
-      # scaled_loss = loss / tf.distribute.get_strategy().num_replicas_in_sync
+        scaled_loss = loss / tf.distribute.get_strategy().num_replicas_in_sync
     tvars = model.trainable_variables
+    if self.task_config.scale_loss:
+      grads = tape.gradient(scaled_loss, tvars)
+    else:
       grads = tape.gradient(loss, tvars)
     optimizer.apply_gradients(list(zip(grads, tvars)))
     self.process_metrics(metrics, inputs, outputs)

official/nlp/tasks/masked_lm_test.py

@@ -28,6 +28,7 @@ class MLMTaskTest(tf.test.TestCase):
   def test_task(self):
     config = masked_lm.MaskedLMConfig(
         init_checkpoint=self.get_temp_dir(),
+        scale_loss=True,
         model=bert.PretrainerConfig(
             encoder=encoders.EncoderConfig(
                 bert=encoders.BertEncoderConfig(vocab_size=30522,