Skip to content

GitLab

Commit 3d6d5b66 authored by Zongwei Zhou's avatar Zongwei Zhou Committed by A. Unique TensorFlower
Browse files

Internal change 

PiperOrigin-RevId: 328614383
parent f91b59c6
Hide whitespace changes
Inline Side-by-side
Showing with 5 additions and 0 deletions
official/nlp/bert/run_pretraining.py
View file @ 3d6d5b66
...@@ -184,9 +184,14 @@ def main(_): ...@@ -184,9 +184,14 @@ def main(_):
gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_param) gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_param)
if not FLAGS.model_dir: if not FLAGS.model_dir:
FLAGS.model_dir = '/tmp/bert20/' FLAGS.model_dir = '/tmp/bert20/'
# Configures cluster spec for multi-worker distribution strategy.
if FLAGS.num_gpus > 0:
_ = distribution_utils.configure_cluster(FLAGS.worker_hosts,
FLAGS.task_index)
strategy = distribution_utils.get_distribution_strategy( strategy = distribution_utils.get_distribution_strategy(
distribution_strategy=FLAGS.distribution_strategy, distribution_strategy=FLAGS.distribution_strategy,
num_gpus=FLAGS.num_gpus, num_gpus=FLAGS.num_gpus,
all_reduce_alg=FLAGS.all_reduce_alg,
tpu_address=FLAGS.tpu) tpu_address=FLAGS.tpu)
if strategy: if strategy:
print('***** Number of cores used : ', strategy.num_replicas_in_sync) print('***** Number of cores used : ', strategy.num_replicas_in_sync)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment