Internal change

PiperOrigin-RevId: 397842944

official/nlp/configs/pretraining_experiments.py

@@ -51,6 +51,7 @@ _TRAINER = cfg.TrainerConfig(
 def bert_pretraining() -> cfg.ExperimentConfig:
   """BERT pretraining experiment."""
   config = cfg.ExperimentConfig(
+      runtime=cfg.RuntimeConfig(enable_xla=True),
       task=masked_lm.MaskedLMConfig(
           train_data=pretrain_dataloader.BertPretrainDataConfig(),
           validation_data=pretrain_dataloader.BertPretrainDataConfig(
@@ -70,6 +71,7 @@ def bert_dynamic() -> cfg.ExperimentConfig:
   TPU needs to run with tf.data service with round-robin behavior.
   """
   config = cfg.ExperimentConfig(
+      runtime=cfg.RuntimeConfig(enable_xla=True),
       task=masked_lm.MaskedLMConfig(
           train_data=pretrain_dynamic_dataloader.BertPretrainDataConfig(),
           validation_data=pretrain_dataloader.BertPretrainDataConfig(

official/nlp/configs/wmt_transformer_experiments.py

@@ -43,6 +43,7 @@ def wmt_transformer_large() -> cfg.ExperimentConfig:
   encdecoder = translation.EncDecoder(
       num_attention_heads=16, intermediate_size=hidden_size * 4)
   config = cfg.ExperimentConfig(
+      runtime=cfg.RuntimeConfig(enable_xla=True),
       task=translation.TranslationConfig(
           model=translation.ModelConfig(
               encoder=encdecoder,

official/vision/beta/configs/image_classification.py

@@ -119,6 +119,7 @@ def image_classification_imagenet() -> cfg.ExperimentConfig:
   eval_batch_size = 4096
   steps_per_epoch = IMAGENET_TRAIN_EXAMPLES // train_batch_size
   config = cfg.ExperimentConfig(
+      runtime=cfg.RuntimeConfig(enable_xla=True),
       task=ImageClassificationTask(
           model=ImageClassificationModel(
               num_classes=1001,

official/vision/beta/configs/maskrcnn.py

@@ -292,7 +292,8 @@ def maskrcnn_resnetfpn_coco() -> cfg.ExperimentConfig:
   eval_batch_size = 8
 
   config = cfg.ExperimentConfig(
-      runtime=cfg.RuntimeConfig(mixed_precision_dtype='bfloat16'),
+      runtime=cfg.RuntimeConfig(
+          mixed_precision_dtype='bfloat16', enable_xla=True),
       task=MaskRCNNTask(
           init_checkpoint='gs://cloud-tpu-checkpoints/vision-2.0/resnet50_imagenet/ckpt-28080',
           init_checkpoint_modules='backbone',