bert_squad_benchmark.py 22.7 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
33
from official.nlp.bert import run_squad
davidmochen's avatar
davidmochen committed
34
from official.utils.misc import distribution_utils
35
from official.utils.misc import keras_utils
36
37
from official.utils.testing import benchmark_wrappers

davidmochen's avatar
davidmochen committed
38
39

# pylint: disable=line-too-long
David Chen's avatar
David Chen committed
40
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
41
42
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
43
SQUAD_VOCAB_FILE = 'gs://tf-perfzero-data/bert/squad/vocab.txt'
David Chen's avatar
David Chen committed
44
SQUAD_MEDIUM_INPUT_META_DATA_PATH = 'gs://tf-perfzero-data/bert/squad/squad_medium_meta_data'
Zongwei Zhou's avatar
Zongwei Zhou committed
45
SQUAD_LONG_INPUT_META_DATA_PATH = 'gs://tf-perfzero-data/bert/squad/squad_long_meta_data'
46
SQUAD_FULL_INPUT_META_DATA_PATH = 'gs://tf-perfzero-data/bert/squad/squad_full_meta_data'
David Chen's avatar
David Chen committed
47
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
48
49
# pylint: enable=line-too-long

David Chen's avatar
David Chen committed
50
TMP_DIR = os.getenv('TMPDIR')
davidmochen's avatar
davidmochen committed
51
52
53
54
55
56
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
57
58
59
60
  def __init__(self, output_dir=None, tpu=None):
    super(BertSquadBenchmarkBase, self).__init__(output_dir=output_dir)
    self.tpu = tpu

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
83
  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.
    """
    if self.tpu or ds_type == 'tpu':
David Chen's avatar
David Chen committed
84
85
      return distribution_utils.get_distribution_strategy(
          distribution_strategy='tpu', tpu_address=self.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."""
David Chen's avatar
David Chen committed
107
    assert tf.version.VERSION.startswith('2.')
108
    self._init_gpu_and_data_threads()
109
    input_meta_data = self._read_input_meta_data_from_file()
110
    strategy = self._get_distribution_strategy(ds_type)
davidmochen's avatar
davidmochen committed
111
112
113
114

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

  @flagsaver.flagsaver
119
120
  def _evaluate_squad(self, ds_type='mirrored'):
    """Runs BERT SQuAD evaluation. Uses mirrored strategy by default."""
David Chen's avatar
David Chen committed
121
    assert tf.version.VERSION.startswith('2.')
122
    self._init_gpu_and_data_threads()
123
    input_meta_data = self._read_input_meta_data_from_file()
124
    strategy = self._get_distribution_strategy(ds_type)
125

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
126
127
128
129
    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)
130
131
    # Use F1 score as reported evaluation metric.
    self.eval_metrics = eval_metrics['f1']
davidmochen's avatar
davidmochen committed
132
133


134
class BertSquadBenchmarkReal(BertSquadBenchmarkBase):
davidmochen's avatar
davidmochen committed
135
136
137
138
  """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
139
140
  `benchmark_(number of gpus)_gpu` format for GPUs and
  `benchmark_(topology)_tpu` format for TPUs.
davidmochen's avatar
davidmochen committed
141
142
  """

David Chen's avatar
David Chen committed
143
144
  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
145
146

  def _setup(self):
147
148
    """Sets up the benchmark and SQuAD flags."""
    super(BertSquadBenchmarkReal, self)._setup()
davidmochen's avatar
davidmochen committed
149
150
151
152
153
    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
154
    FLAGS.steps_per_loop = 100
davidmochen's avatar
davidmochen committed
155

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

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

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

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

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

186
    self._run_and_report_benchmark()
davidmochen's avatar
davidmochen committed
187

188
189
190
191
192
193
194
195
  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
196
    self._run_and_report_benchmark(run_eagerly=True)
197

198
199
200
201
202
203
  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')
204
205
    # XLA runs out of memory when running with batch size 4.
    FLAGS.train_batch_size = 3
206
    FLAGS.enable_xla = True
207

208
    self._run_and_report_benchmark()
209
210
211
212
213
214
215

  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')
216
    FLAGS.train_batch_size = 4
217

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

  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')
227
    FLAGS.train_batch_size = 4
228

229
    self._run_and_report_benchmark(ds_type='off', run_eagerly=True)
230

davidmochen's avatar
davidmochen committed
231
  def benchmark_2_gpu(self):
232
    """Tests BERT SQuAD model performance with 2 GPUs."""
davidmochen's avatar
davidmochen committed
233
234
235
236

    self._setup()
    self.num_gpus = 2
    FLAGS.model_dir = self._get_model_dir('benchmark_2_gpu_squad')
237
    FLAGS.train_batch_size = 8
davidmochen's avatar
davidmochen committed
238

239
    self._run_and_report_benchmark()
davidmochen's avatar
davidmochen committed
240
241

  def benchmark_4_gpu(self):
242
    """Tests BERT SQuAD model performance with 4 GPUs."""
davidmochen's avatar
davidmochen committed
243
244
245
246

    self._setup()
    self.num_gpus = 4
    FLAGS.model_dir = self._get_model_dir('benchmark_4_gpu_squad')
247
    FLAGS.train_batch_size = 16
davidmochen's avatar
davidmochen committed
248

249
    self._run_and_report_benchmark()
davidmochen's avatar
davidmochen committed
250
251

  def benchmark_8_gpu(self):
252
253
254
255
256
    """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')
257
    FLAGS.train_batch_size = 32
258
    FLAGS.tf_gpu_thread_mode = 'gpu_private'
259

260
    self._run_and_report_benchmark()
261

262
263
264
265
266
267
268
269
270
271
  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
272
    self._run_and_report_benchmark(run_eagerly=True)
273

274
275
276
277
278
279
280
281
282
283
284
285
  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()

286
287
288
289
290
291
292
293
294
295
296
297
298
  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()

299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
  def benchmark_2_gpu_fp16(self):
    """Tests BERT SQuAD model performance with 2 GPUs and FP16."""

    self._setup()
    self.num_gpus = 2
    FLAGS.model_dir = self._get_model_dir('benchmark_2_gpu_squad_fp16')
    FLAGS.train_batch_size = 8
    FLAGS.dtype = 'fp16'
    FLAGS.loss_scale = 'dynamic'

    self._run_and_report_benchmark()

  def benchmark_4_gpu_fp16(self):
    """Tests BERT SQuAD model performance with 4 GPUs and FP16."""

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

    self._run_and_report_benchmark()

  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'
332
    FLAGS.tf_gpu_thread_mode = 'gpu_private'
333
334
335

    self._run_and_report_benchmark()

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
336
337
338
339
340
341
342
343
344
345
346
347
348
  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()

349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
  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_4_gpu_amp(self):
    """Tests BERT SQuAD model performance with 1 GPU with automatic mixed precision."""

    self._setup()
    self.num_gpus = 4
    FLAGS.model_dir = self._get_model_dir('benchmark_4_gpu_amp_squad')
    FLAGS.train_batch_size = 16
    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'
382
    FLAGS.tf_gpu_thread_mode = 'gpu_private'
383
384

    self._run_and_report_benchmark()
385

David Chen's avatar
David Chen committed
386
387
388
389
390
391
392
393
394
  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

    self._run_and_report_benchmark()

395
396
397
398
399

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
400
401
  `benchmark_(number of gpus)_gpu` format for GPUs and
  `benchmark_(topology)_tpu` format for TPUs.
402
403
  """

David Chen's avatar
David Chen committed
404
405
  def __init__(self, output_dir=None, tpu=None, **kwargs):
    super(BertSquadAccuracy, self).__init__(output_dir=output_dir, tpu=tpu)
406
407
408
409
410
411
412
413
414
415
416

  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
417
    FLAGS.steps_per_loop = 100
418

419
  @benchmark_wrappers.enable_runtime_flags
420
  def _run_and_report_benchmark(self,
421
422
                                run_eagerly=False,
                                ds_type='mirrored'):
423
    """Runs the benchmark and reports various metrics."""
424
    start_time_sec = time.time()
425
426
    self._train_squad(run_eagerly=run_eagerly, ds_type=ds_type)
    self._evaluate_squad(ds_type=ds_type)
427
428
429
430
    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
431
    summary['start_time_sec'] = start_time_sec
432
433
434
435

    super(BertSquadAccuracy, self)._report_benchmark(
        stats=summary,
        wall_time_sec=wall_time_sec,
436
        min_accuracy=0.900,
437
        max_accuracy=0.920)
438

439
440
441
442
443
444
445
446
  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

447
    self._run_and_report_benchmark(ds_type='off', run_eagerly=True)
448

449
450
  def benchmark_8_gpu(self):
    """Tests BERT SQuAD model accuracy with 8 GPUs."""
davidmochen's avatar
davidmochen committed
451
452
453
454

    self._setup()
    self.num_gpus = 8
    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_squad')
455
    FLAGS.train_batch_size = 24
456
    FLAGS.tf_gpu_thread_mode = 'gpu_private'
davidmochen's avatar
davidmochen committed
457

458
    self._run_and_report_benchmark()
davidmochen's avatar
davidmochen committed
459

460
461
462
463
464
465
466
467
468
  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'
469
    FLAGS.tf_gpu_thread_mode = 'gpu_private'
470
471
472

    self._run_and_report_benchmark()

473
474
475
476
477
478
479
  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
480
    FLAGS.enable_xla = True
481
    FLAGS.tf_gpu_thread_mode = 'gpu_private'
482

483
    self._run_and_report_benchmark()
484

David Chen's avatar
David Chen committed
485
486
487
488
489
490
491
492
493
  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
494

495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
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
512
    FLAGS.steps_per_loop = 100
513
514
515
516
517
518
519

  @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()
520
521
522
    self._train_squad(run_eagerly=run_eagerly,
                      ds_type='multi_worker_mirrored')
    self._evaluate_squad(ds_type='multi_worker_mirrored')
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
    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
553
554
555
556
557
558
559
560
  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')

561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
  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
583
    FLAGS.input_meta_data_path = SQUAD_FULL_INPUT_META_DATA_PATH
584
585
    FLAGS.bert_config_file = MODEL_CONFIG_FILE_PATH
    FLAGS.num_train_epochs = 1
586
    FLAGS.steps_per_loop = 100
587
588
589
590
591
592

  @benchmark_wrappers.enable_runtime_flags
  def _run_and_report_benchmark(self,
                                use_ds=True,
                                run_eagerly=False):
    """Runs the benchmark and reports various metrics."""
593
594
595
596
    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
597
    start_time_sec = time.time()
598
599
    self._train_squad(run_eagerly=run_eagerly,
                      ds_type='multi_worker_mirrored')
600
601
602
603
604
605
606
607
608
609
610
    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
611
  def _benchmark_common(self, num_workers, all_reduce_alg):
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
    """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
640
    self._benchmark_common(num_workers=2, all_reduce_alg='ring')
641
642
643

  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
644
    self._benchmark_common(num_workers=2, all_reduce_alg='nccl')
645
646
647

  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
648
    self._benchmark_common(num_workers=8, all_reduce_alg='ring')
649
650
651

  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
652
    self._benchmark_common(num_workers=8, all_reduce_alg='nccl')
653
654


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