bert_squad_benchmark.py 21.6 KB
Newer Older
davidmochen's avatar
davidmochen committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Executes BERT SQuAD benchmarks and accuracy tests."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import json
import os
import time

# pylint: disable=g-bad-import-order
from absl import flags
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
27
from absl import logging
davidmochen's avatar
davidmochen committed
28
from absl.testing import flagsaver
29
import tensorflow as tf
davidmochen's avatar
davidmochen committed
30
31
# pylint: enable=g-bad-import-order

32
from official.benchmark import bert_benchmark_utils as benchmark_utils
Jing Li's avatar
Jing Li committed
33
from official.benchmark import owner_utils
34
from official.nlp.bert import run_squad
davidmochen's avatar
davidmochen committed
35
from official.utils.misc import distribution_utils
36
from official.utils.misc import keras_utils
37
from official.benchmark import benchmark_wrappers
38

davidmochen's avatar
davidmochen committed
39
40

# pylint: disable=line-too-long
David Chen's avatar
David Chen committed
41
PRETRAINED_CHECKPOINT_PATH = 'gs://cloud-tpu-checkpoints/bert/keras_bert/uncased_L-24_H-1024_A-16/bert_model.ckpt'
davidmochen's avatar
davidmochen committed
42
43
SQUAD_TRAIN_DATA_PATH = 'gs://tf-perfzero-data/bert/squad/squad_train.tf_record'
SQUAD_PREDICT_FILE = 'gs://tf-perfzero-data/bert/squad/dev-v1.1.json'
David Chen's avatar
David Chen committed
44
SQUAD_VOCAB_FILE = 'gs://tf-perfzero-data/bert/squad/vocab.txt'
David Chen's avatar
David Chen committed
45
SQUAD_MEDIUM_INPUT_META_DATA_PATH = 'gs://tf-perfzero-data/bert/squad/squad_medium_meta_data'
Zongwei Zhou's avatar
Zongwei Zhou committed
46
SQUAD_LONG_INPUT_META_DATA_PATH = 'gs://tf-perfzero-data/bert/squad/squad_long_meta_data'
47
SQUAD_FULL_INPUT_META_DATA_PATH = 'gs://tf-perfzero-data/bert/squad/squad_full_meta_data'
David Chen's avatar
David Chen committed
48
MODEL_CONFIG_FILE_PATH = 'gs://cloud-tpu-checkpoints/bert/keras_bert/uncased_L-24_H-1024_A-16/bert_config.json'
davidmochen's avatar
davidmochen committed
49
50
# pylint: enable=line-too-long

David Chen's avatar
David Chen committed
51
TMP_DIR = os.getenv('TMPDIR')
davidmochen's avatar
davidmochen committed
52
53
54
55
56
57
FLAGS = flags.FLAGS


class BertSquadBenchmarkBase(benchmark_utils.BertBenchmarkBase):
  """Base class to hold methods common to test classes in the module."""

David Chen's avatar
David Chen committed
58
  def __init__(self, output_dir=None, tpu=None):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
59
    super(BertSquadBenchmarkBase, self).__init__(output_dir=output_dir, tpu=tpu)
David Chen's avatar
David Chen committed
60

61
62
  def _read_training_summary_from_file(self):
    """Reads the training summary from a file."""
63
64
    summary_path = os.path.join(FLAGS.model_dir,
                                'summaries/training_summary.txt')
65
66
    with tf.io.gfile.GFile(summary_path, 'rb') as reader:
      return json.loads(reader.read().decode('utf-8'))
67

68
69
70
71
  def _read_input_meta_data_from_file(self):
    """Reads the input metadata from a file."""
    with tf.io.gfile.GFile(FLAGS.input_meta_data_path, 'rb') as reader:
      return json.loads(reader.read().decode('utf-8'))
72

73
74
75
76
77
78
79
80
81
82
  def _get_distribution_strategy(self, ds_type='mirrored'):
    """Gets the distribution strategy.

    Args:
      ds_type: String, the distribution strategy type to be used. Can be
      'mirrored', 'multi_worker_mirrored', 'tpu' and 'off'.

    Returns:
      A `tf.distribute.DistibutionStrategy` object.
    """
Jing Li's avatar
Jing Li committed
83
    if self.default_flags['tpu'] or ds_type == 'tpu':
David Chen's avatar
David Chen committed
84
      return distribution_utils.get_distribution_strategy(
Jing Li's avatar
Jing Li committed
85
          distribution_strategy='tpu', tpu_address=self.default_flags['tpu'])
86
87
88
89
90
91
92
93
    elif ds_type == 'multi_worker_mirrored':
      # Configures cluster spec for multi-worker distribution strategy.
      _ = distribution_utils.configure_cluster(FLAGS.worker_hosts,
                                               FLAGS.task_index)
    return distribution_utils.get_distribution_strategy(
        distribution_strategy=ds_type,
        num_gpus=self.num_gpus,
        all_reduce_alg=FLAGS.all_reduce_alg)
94

95
96
97
98
99
100
101
102
103
  def _init_gpu_and_data_threads(self):
    """Set env variables before any TF calls."""
    if FLAGS.tf_gpu_thread_mode:
      keras_utils.set_gpu_thread_mode_and_count(
          per_gpu_thread_count=FLAGS.per_gpu_thread_count,
          gpu_thread_mode=FLAGS.tf_gpu_thread_mode,
          num_gpus=self.num_gpus,
          datasets_num_private_threads=FLAGS.datasets_num_private_threads)

davidmochen's avatar
davidmochen committed
104
  @flagsaver.flagsaver
105
106
  def _train_squad(self, run_eagerly=False, ds_type='mirrored'):
    """Runs BERT SQuAD training. Uses mirrored strategy by default."""
107
    self._init_gpu_and_data_threads()
108
    input_meta_data = self._read_input_meta_data_from_file()
109
    strategy = self._get_distribution_strategy(ds_type)
davidmochen's avatar
davidmochen committed
110
111
112
113

    run_squad.train_squad(
        strategy=strategy,
        input_meta_data=input_meta_data,
114
        run_eagerly=run_eagerly,
davidmochen's avatar
davidmochen committed
115
        custom_callbacks=[self.timer_callback])
116
117

  @flagsaver.flagsaver
118
119
  def _evaluate_squad(self, ds_type='mirrored'):
    """Runs BERT SQuAD evaluation. Uses mirrored strategy by default."""
120
    self._init_gpu_and_data_threads()
121
    input_meta_data = self._read_input_meta_data_from_file()
122
    strategy = self._get_distribution_strategy(ds_type)
123

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
124
125
126
127
    if input_meta_data.get('version_2_with_negative', False):
      logging.error('In memory evaluation result for SQuAD v2 is not accurate')
    eval_metrics = run_squad.eval_squad(strategy=strategy,
                                        input_meta_data=input_meta_data)
128
    # Use F1 score as reported evaluation metric.
Hongkun Yu's avatar
Hongkun Yu committed
129
    self.eval_metrics = eval_metrics['final_f1']
davidmochen's avatar
davidmochen committed
130
131


132
class BertSquadBenchmarkReal(BertSquadBenchmarkBase):
davidmochen's avatar
davidmochen committed
133
134
135
136
  """Short benchmark performance tests for BERT SQuAD model.

  Tests BERT SQuAD performance in different GPU configurations.
  The naming convention of below test cases follow
David Chen's avatar
David Chen committed
137
138
  `benchmark_(number of gpus)_gpu` format for GPUs and
  `benchmark_(topology)_tpu` format for TPUs.
davidmochen's avatar
davidmochen committed
139
140
  """

David Chen's avatar
David Chen committed
141
142
  def __init__(self, output_dir=TMP_DIR, tpu=None, **kwargs):
    super(BertSquadBenchmarkReal, self).__init__(output_dir=output_dir, tpu=tpu)
davidmochen's avatar
davidmochen committed
143
144

  def _setup(self):
145
146
    """Sets up the benchmark and SQuAD flags."""
    super(BertSquadBenchmarkReal, self)._setup()
davidmochen's avatar
davidmochen committed
147
148
149
150
151
    FLAGS.train_data_path = SQUAD_TRAIN_DATA_PATH
    FLAGS.predict_file = SQUAD_PREDICT_FILE
    FLAGS.vocab_file = SQUAD_VOCAB_FILE
    FLAGS.bert_config_file = MODEL_CONFIG_FILE_PATH
    FLAGS.num_train_epochs = 1
152
    FLAGS.steps_per_loop = 100
davidmochen's avatar
davidmochen committed
153

154
  @benchmark_wrappers.enable_runtime_flags
155
  def _run_and_report_benchmark(self,
156
157
                                run_eagerly=False,
                                ds_type='mirrored'):
158
    """Runs the benchmark and reports various metrics."""
159
    if FLAGS.train_batch_size <= 4 or run_eagerly:
160
161
162
      FLAGS.input_meta_data_path = SQUAD_MEDIUM_INPUT_META_DATA_PATH
    else:
      FLAGS.input_meta_data_path = SQUAD_LONG_INPUT_META_DATA_PATH
163
    start_time_sec = time.time()
164
    self._train_squad(run_eagerly=run_eagerly, ds_type=ds_type)
165
166
167
    wall_time_sec = time.time() - start_time_sec

    summary = self._read_training_summary_from_file()
David Chen's avatar
David Chen committed
168
    summary['start_time_sec'] = start_time_sec
169
170
171
172
173
174

    super(BertSquadBenchmarkReal, self)._report_benchmark(
        stats=summary,
        wall_time_sec=wall_time_sec,
        min_accuracy=0,
        max_accuracy=1)
davidmochen's avatar
davidmochen committed
175
176

  def benchmark_1_gpu(self):
177
    """Tests BERT SQuAD model performance with 1 GPU."""
davidmochen's avatar
davidmochen committed
178
179
180
181

    self._setup()
    self.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu_squad')
182
    FLAGS.train_batch_size = 4
davidmochen's avatar
davidmochen committed
183

184
    self._run_and_report_benchmark()
davidmochen's avatar
davidmochen committed
185

186
187
188
189
190
191
192
193
  def benchmark_1_gpu_eager(self):
    """Tests BERT SQuAD model performance with 1 GPU."""

    self._setup()
    self.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu_squad_eager')
    FLAGS.train_batch_size = 2

Zongwei Zhou's avatar
Zongwei Zhou committed
194
    self._run_and_report_benchmark(run_eagerly=True)
195

196
197
198
199
200
201
  def benchmark_1_gpu_xla(self):
    """Tests BERT SQuAD model performance with 1 GPU with XLA."""

    self._setup()
    self.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu_xla_squad')
202
203
    # XLA runs out of memory when running with batch size 4.
    FLAGS.train_batch_size = 3
204
    FLAGS.enable_xla = True
205

206
    self._run_and_report_benchmark()
207
208
209
210
211
212
213

  def benchmark_1_gpu_no_dist_strat(self):
    """Tests BERT SQuAD model performance with 1 GPU without DS."""

    self._setup()
    self.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu_no_dist_strat_squad')
214
    FLAGS.train_batch_size = 4
215

216
    self._run_and_report_benchmark(ds_type='off')
217
218
219
220
221
222
223
224

  def benchmark_1_gpu_eager_no_dist_strat(self):
    """Tests BERT SQuAD model performance with 1 GPU with eager execution."""

    self._setup()
    self.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir(
        'benchmark_1_gpu_eager_no_dist_strat_squad')
225
    FLAGS.train_batch_size = 4
226

227
    self._run_and_report_benchmark(ds_type='off', run_eagerly=True)
228

Jing Li's avatar
Jing Li committed
229
  @owner_utils.Owner('tf-model-garden')
davidmochen's avatar
davidmochen committed
230
  def benchmark_8_gpu(self):
231
232
233
234
235
    """Tests BERT SQuAD model performance with 8 GPUs."""

    self._setup()
    self.num_gpus = 8
    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_squad')
Sai Ganesh Bandiatmakuri's avatar
Sai Ganesh Bandiatmakuri committed
236
    FLAGS.train_batch_size = 24
237
    FLAGS.tf_gpu_thread_mode = 'gpu_private'
238

239
    self._run_and_report_benchmark()
240

241
242
243
244
245
246
247
248
249
250
  def benchmark_1_gpu_fp16_eager(self):
    """Tests BERT SQuAD model performance with 1 GPU and FP16."""

    self._setup()
    self.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu_squad_fp16_eager')
    FLAGS.train_batch_size = 4
    FLAGS.dtype = 'fp16'
    FLAGS.loss_scale = 'dynamic'

Zongwei Zhou's avatar
Zongwei Zhou committed
251
    self._run_and_report_benchmark(run_eagerly=True)
252

253
254
255
256
257
258
259
260
261
262
263
264
  def benchmark_1_gpu_fp16(self):
    """Tests BERT SQuAD model performance with 1 GPU and FP16."""

    self._setup()
    self.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu_squad_fp16')
    FLAGS.train_batch_size = 4
    FLAGS.dtype = 'fp16'
    FLAGS.loss_scale = 'dynamic'

    self._run_and_report_benchmark()

265
266
267
268
269
270
271
272
273
274
275
276
277
  def benchmark_1_gpu_xla_fp16(self):
    """Tests BERT SQuAD model performance with 1 GPU with XLA and FP16."""

    self._setup()
    self.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu_xla_squad_fp16')
    FLAGS.train_batch_size = 4
    FLAGS.enable_xla = True
    FLAGS.dtype = 'fp16'
    FLAGS.loss_scale = 'dynamic'

    self._run_and_report_benchmark()

278
279
280
281
282
283
284
285
286
  def benchmark_8_gpu_fp16(self):
    """Tests BERT SQuAD model performance with 8 GPUs."""

    self._setup()
    self.num_gpus = 8
    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_squad_fp16')
    FLAGS.train_batch_size = 32
    FLAGS.dtype = 'fp16'
    FLAGS.loss_scale = 'dynamic'
287
    FLAGS.tf_gpu_thread_mode = 'gpu_private'
288
289
290

    self._run_and_report_benchmark()

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
291
292
293
294
295
296
297
298
299
300
301
302
303
  def benchmark_8_gpu_xla_fp16(self):
    """Tests BERT SQuAD model performance with 8 GPUs with XLA."""

    self._setup()
    self.num_gpus = 8
    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_squad_fp16')
    FLAGS.train_batch_size = 32
    FLAGS.enable_xla = True
    FLAGS.dtype = 'fp16'
    FLAGS.loss_scale = 'dynamic'

    self._run_and_report_benchmark()

304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
  def benchmark_1_gpu_amp(self):
    """Tests BERT SQuAD model performance with 1 GPU with automatic mixed precision."""

    self._setup()
    self.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu_amp_squad')
    FLAGS.train_batch_size = 4
    FLAGS.dtype = 'fp16'
    FLAGS.fp16_implementation = 'graph_rewrite'

    self._run_and_report_benchmark()

  def benchmark_8_gpu_amp(self):
    """Tests BERT SQuAD model performance with 1 GPU with automatic mixed precision."""

    self._setup()
    self.num_gpus = 8
    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_amp_squad')
    FLAGS.train_batch_size = 32
    FLAGS.dtype = 'fp16'
    FLAGS.fp16_implementation = 'graph_rewrite'
325
    FLAGS.tf_gpu_thread_mode = 'gpu_private'
326
327

    self._run_and_report_benchmark()
328

Jing Li's avatar
Jing Li committed
329
  @owner_utils.Owner('tf-model-garden')
David Chen's avatar
David Chen committed
330
331
332
333
334
335
  def benchmark_2x2_tpu(self):
    """Tests BERT SQuAD model performance with 2x2 TPU."""

    self._setup()
    FLAGS.model_dir = self._get_model_dir('benchmark_2x2_tpu')
    FLAGS.train_batch_size = 48
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
336
337
338
339
340
341
342
    FLAGS.predict_batch_size = 48
    FLAGS.mode = 'train'
    FLAGS.learning_rate = 8e-5
    FLAGS.num_train_epochs = 1
    FLAGS.steps_per_loop = 100
    FLAGS.do_lower_case = True
    FLAGS.init_checkpoint = PRETRAINED_CHECKPOINT_PATH
David Chen's avatar
David Chen committed
343
344
    self._run_and_report_benchmark()

345
346
347
348
349

class BertSquadAccuracy(BertSquadBenchmarkBase):
  """Short accuracy test for BERT SQuAD model.

  Tests BERT SQuAD accuracy. The naming convention of below test cases follow
David Chen's avatar
David Chen committed
350
351
  `benchmark_(number of gpus)_gpu` format for GPUs and
  `benchmark_(topology)_tpu` format for TPUs.
352
353
  """

David Chen's avatar
David Chen committed
354
355
  def __init__(self, output_dir=None, tpu=None, **kwargs):
    super(BertSquadAccuracy, self).__init__(output_dir=output_dir, tpu=tpu)
356
357
358
359
360
361
362
363
364
365
366

  def _setup(self):
    """Sets up the benchmark and SQuAD flags."""
    super(BertSquadAccuracy, self)._setup()
    FLAGS.train_data_path = SQUAD_TRAIN_DATA_PATH
    FLAGS.predict_file = SQUAD_PREDICT_FILE
    FLAGS.vocab_file = SQUAD_VOCAB_FILE
    FLAGS.input_meta_data_path = SQUAD_FULL_INPUT_META_DATA_PATH
    FLAGS.bert_config_file = MODEL_CONFIG_FILE_PATH
    FLAGS.init_checkpoint = PRETRAINED_CHECKPOINT_PATH
    FLAGS.num_train_epochs = 2
367
    FLAGS.steps_per_loop = 100
368

369
  @benchmark_wrappers.enable_runtime_flags
370
  def _run_and_report_benchmark(self,
371
372
                                run_eagerly=False,
                                ds_type='mirrored'):
373
    """Runs the benchmark and reports various metrics."""
374
    start_time_sec = time.time()
375
376
    self._train_squad(run_eagerly=run_eagerly, ds_type=ds_type)
    self._evaluate_squad(ds_type=ds_type)
377
378
379
380
    wall_time_sec = time.time() - start_time_sec

    summary = self._read_training_summary_from_file()
    summary['eval_metrics'] = self.eval_metrics
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
381
    summary['start_time_sec'] = start_time_sec
382
383
384
385

    super(BertSquadAccuracy, self)._report_benchmark(
        stats=summary,
        wall_time_sec=wall_time_sec,
386
        min_accuracy=0.900,
387
        max_accuracy=0.920)
388

389
390
391
392
393
394
395
396
  def benchmark_1_gpu_eager(self):
    """Tests BERT SQuAD model accuracy with 1 GPU with eager execution."""

    self._setup()
    self.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu_squad_eager')
    FLAGS.train_batch_size = 4

397
    self._run_and_report_benchmark(ds_type='off', run_eagerly=True)
398

Jing Li's avatar
Jing Li committed
399
  @owner_utils.Owner('tf-model-garden')
400
401
  def benchmark_8_gpu(self):
    """Tests BERT SQuAD model accuracy with 8 GPUs."""
davidmochen's avatar
davidmochen committed
402
403
404
405

    self._setup()
    self.num_gpus = 8
    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_squad')
406
    FLAGS.train_batch_size = 24
407
    FLAGS.tf_gpu_thread_mode = 'gpu_private'
davidmochen's avatar
davidmochen committed
408

409
    self._run_and_report_benchmark()
davidmochen's avatar
davidmochen committed
410

411
412
413
414
415
416
417
418
419
  def benchmark_8_gpu_fp16(self):
    """Tests BERT SQuAD model accuracy with 8 GPUs and FP16."""

    self._setup()
    self.num_gpus = 8
    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_squad_fp16')
    FLAGS.train_batch_size = 32
    FLAGS.dtype = 'fp16'
    FLAGS.loss_scale = 'dynamic'
420
    FLAGS.tf_gpu_thread_mode = 'gpu_private'
421
422
423

    self._run_and_report_benchmark()

424
425
426
427
428
429
430
  def benchmark_8_gpu_xla(self):
    """Tests BERT SQuAD model accuracy with 8 GPUs."""

    self._setup()
    self.num_gpus = 8
    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_squad_xla')
    FLAGS.train_batch_size = 32
431
    FLAGS.enable_xla = True
432
    FLAGS.tf_gpu_thread_mode = 'gpu_private'
433

434
    self._run_and_report_benchmark()
435

Jing Li's avatar
Jing Li committed
436
  @owner_utils.Owner('tf-model-garden')
David Chen's avatar
David Chen committed
437
438
439
440
441
442
443
444
445
  def benchmark_2x2_tpu(self):
    """Tests BERT SQuAD model accuracy with 2x2 TPU."""

    self._setup()
    FLAGS.model_dir = self._get_model_dir('benchmark_2x2_tpu')
    FLAGS.train_batch_size = 48

    self._run_and_report_benchmark()

davidmochen's avatar
davidmochen committed
446

447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
class BertSquadMultiWorkerAccuracy(BertSquadBenchmarkBase):
  """BERT SQuAD distributed accuracy tests with multiple workers."""

  def __init__(self, output_dir=None, tpu=None, **kwargs):
    super(BertSquadMultiWorkerAccuracy, self).__init__(
        output_dir=output_dir, tpu=tpu)

  def _setup(self):
    """Sets up the benchmark and SQuAD flags."""
    super(BertSquadMultiWorkerAccuracy, self)._setup()
    FLAGS.train_data_path = SQUAD_TRAIN_DATA_PATH
    FLAGS.predict_file = SQUAD_PREDICT_FILE
    FLAGS.vocab_file = SQUAD_VOCAB_FILE
    FLAGS.input_meta_data_path = SQUAD_FULL_INPUT_META_DATA_PATH
    FLAGS.bert_config_file = MODEL_CONFIG_FILE_PATH
    FLAGS.init_checkpoint = PRETRAINED_CHECKPOINT_PATH
    FLAGS.num_train_epochs = 2
464
    FLAGS.steps_per_loop = 100
465
466
467
468
469
470
471

  @benchmark_wrappers.enable_runtime_flags
  def _run_and_report_benchmark(self,
                                use_ds=True,
                                run_eagerly=False):
    """Runs the benchmark and reports various metrics."""
    start_time_sec = time.time()
472
473
474
    self._train_squad(run_eagerly=run_eagerly,
                      ds_type='multi_worker_mirrored')
    self._evaluate_squad(ds_type='multi_worker_mirrored')
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
    wall_time_sec = time.time() - start_time_sec

    summary = self._read_training_summary_from_file()
    summary['eval_metrics'] = self.eval_metrics

    super(BertSquadMultiWorkerAccuracy, self)._report_benchmark(
        stats=summary,
        wall_time_sec=wall_time_sec,
        min_accuracy=0.900,
        max_accuracy=0.920)

  def _benchmark_common(self, num_workers, all_reduce_alg):
    """Common to all benchmarks in this class."""
    self._setup()

    num_gpus = 8
    FLAGS.num_gpus = num_gpus
    FLAGS.dtype = 'fp16'
    FLAGS.enable_xla = False
    FLAGS.distribution_strategy = 'multi_worker_mirrored'
    FLAGS.tf_gpu_thread_mode = 'gpu_private'
    FLAGS.datasets_num_private_threads = 32
    FLAGS.model_dir = self._get_model_dir(
        'benchmark_8_gpu_{}_worker_fp16_{}_tweaked'.format(
            num_workers, all_reduce_alg))
    FLAGS.train_batch_size = 4 * num_gpus * num_workers
    FLAGS.all_reduce_alg = all_reduce_alg

    self._run_and_report_benchmark()

Yanhui Liang's avatar
Yanhui Liang committed
505
506
507
508
509
510
511
512
  def benchmark_eager_8_gpu_2_workers_fp16_ring_tweaked(self):
    """8 GPUs per worker, 2 workers, fp16, ring all-reduce."""
    self._benchmark_common(num_workers=2, all_reduce_alg='ring')

  def benchmark_eager_8_gpu_2_workers_fp16_nccl_tweaked(self):
    """8 GPUs per worker, 2 workers, fp16, nccl all-reduce."""
    self._benchmark_common(num_workers=2, all_reduce_alg='nccl')

513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
  def benchmark_8_gpu_8_workers_fp16_ring_tweaked(self):
    """8 GPUs per worker, 8 workers, fp16, ring all-reduce."""
    self._benchmark_common(num_workers=8, all_reduce_alg='ring')

  def benchmark_8_gpu_8_workers_fp16_nccl_tweaked(self):
    """8 GPUs per worker, 8 workers, fp16, nccl all-reduce."""
    self._benchmark_common(num_workers=8, all_reduce_alg='nccl')


class BertSquadMultiWorkerBenchmark(BertSquadBenchmarkBase):
  """BERT SQuAD distributed benchmark tests with multiple workers."""

  def __init__(self, output_dir=TMP_DIR, tpu=None, **kwargs):
    super(BertSquadMultiWorkerBenchmark, self).__init__(
        output_dir=output_dir, tpu=tpu)

  def _setup(self):
    """Sets up the benchmark and SQuAD flags."""
    super(BertSquadMultiWorkerBenchmark, self)._setup()
    FLAGS.train_data_path = SQUAD_TRAIN_DATA_PATH
    FLAGS.predict_file = SQUAD_PREDICT_FILE
    FLAGS.vocab_file = SQUAD_VOCAB_FILE
535
    FLAGS.input_meta_data_path = SQUAD_FULL_INPUT_META_DATA_PATH
536
537
    FLAGS.bert_config_file = MODEL_CONFIG_FILE_PATH
    FLAGS.num_train_epochs = 1
538
    FLAGS.steps_per_loop = 100
539
540
541
542
543
544

  @benchmark_wrappers.enable_runtime_flags
  def _run_and_report_benchmark(self,
                                use_ds=True,
                                run_eagerly=False):
    """Runs the benchmark and reports various metrics."""
545
546
547
548
    if FLAGS.train_batch_size <= 4 * 8:
      FLAGS.input_meta_data_path = SQUAD_LONG_INPUT_META_DATA_PATH
    else:
      FLAGS.input_meta_data_path = SQUAD_FULL_INPUT_META_DATA_PATH
549
    start_time_sec = time.time()
550
551
    self._train_squad(run_eagerly=run_eagerly,
                      ds_type='multi_worker_mirrored')
552
553
554
555
556
557
558
559
560
561
562
    wall_time_sec = time.time() - start_time_sec

    summary = self._read_training_summary_from_file()
    summary['start_time_sec'] = start_time_sec

    super(BertSquadMultiWorkerBenchmark, self)._report_benchmark(
        stats=summary,
        wall_time_sec=wall_time_sec,
        min_accuracy=0,
        max_accuracy=1)

Hongkun Yu's avatar
Hongkun Yu committed
563
  def _benchmark_common(self, num_workers, all_reduce_alg):
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
    """Common to all benchmarks in this class."""
    self._setup()

    num_gpus = 8
    FLAGS.num_gpus = num_gpus
    FLAGS.dtype = 'fp16'
    FLAGS.enable_xla = False
    FLAGS.distribution_strategy = 'multi_worker_mirrored'
    FLAGS.tf_gpu_thread_mode = 'gpu_private'
    FLAGS.datasets_num_private_threads = 32
    FLAGS.model_dir = self._get_model_dir(
        'benchmark_8_gpu_{}_worker_fp16_{}_tweaked'.format(
            num_workers, all_reduce_alg))
    FLAGS.train_batch_size = 4 * num_gpus * num_workers
    FLAGS.all_reduce_alg = all_reduce_alg

    self._run_and_report_benchmark()

  def benchmark_8_gpu_1_worker_fp16_ring_tweaked(self):
    """8 GPUs per worker, 1 worker, fp16, ring all-reduce."""
    self._benchmark_common(num_workers=1, all_reduce_alg='ring')

  def benchmark_8_gpu_1_worker_fp16_nccl_tweaked(self):
    """8 GPUs per worker, 1 worker, fp16, nccl all-reduce."""
    self._benchmark_common(num_workers=1, all_reduce_alg='nccl')

  def benchmark_8_gpu_2_workers_fp16_ring_tweaked(self):
    """8 GPUs per worker, 2 workers, fp16, ring all-reduce."""
Hongkun Yu's avatar
Hongkun Yu committed
592
    self._benchmark_common(num_workers=2, all_reduce_alg='ring')
593
594
595

  def benchmark_8_gpu_2_workers_fp16_nccl_tweaked(self):
    """8 GPUs per worker, 2 workers, fp16, nccl all-reduce."""
Hongkun Yu's avatar
Hongkun Yu committed
596
    self._benchmark_common(num_workers=2, all_reduce_alg='nccl')
597
598
599

  def benchmark_8_gpu_8_workers_fp16_ring_tweaked(self):
    """8 GPUs per worker, 8 workers, fp16, ring all-reduce."""
Hongkun Yu's avatar
Hongkun Yu committed
600
    self._benchmark_common(num_workers=8, all_reduce_alg='ring')
601
602
603

  def benchmark_8_gpu_8_workers_fp16_nccl_tweaked(self):
    """8 GPUs per worker, 8 workers, fp16, nccl all-reduce."""
Hongkun Yu's avatar
Hongkun Yu committed
604
    self._benchmark_common(num_workers=8, all_reduce_alg='nccl')
605
606


davidmochen's avatar
davidmochen committed
607
608
if __name__ == '__main__':
  tf.test.main()