bert_pretrain_benchmark.py 12.7 KB
Newer Older
Chen Chen's avatar
Chen Chen 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
27
28
29
30
# Lint as: python3
# Copyright 2020 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 benchmark testing for bert pretraining."""
# pylint: disable=line-too-long
from __future__ import print_function

import json
import os
import time
from typing import Optional

from absl import flags
from absl import logging
import tensorflow as tf  # pylint: disable=g-bad-import-order

from official.benchmark import benchmark_wrappers
from official.benchmark import bert_benchmark_utils
Jing Li's avatar
Jing Li committed
31
from official.benchmark import owner_utils
Chen Chen's avatar
Chen Chen committed
32
33
34
35
36
from official.nlp.bert import run_pretraining
from official.utils.flags import core as flags_core
from official.utils.misc import distribution_utils

# Pretrain masked lanauge modeling accuracy range:
Chen Chen's avatar
Chen Chen committed
37
38
MIN_MLM_ACCURACY = 0.635
MAX_MLM_ACCURACY = 0.645
Chen Chen's avatar
Chen Chen committed
39
40

# Pretrain next sentence prediction accuracy range:
Chen Chen's avatar
Chen Chen committed
41
42
MIN_NSP_ACCURACY = 0.94
MAX_NSP_ACCURACY = 0.96
Chen Chen's avatar
Chen Chen committed
43

Zongwei Zhou's avatar
Zongwei Zhou committed
44
45
46
47
48
49
50
51
52
53

# Pretrain masked lanauge modeling accuracy range:
MIN_MLM_ACCURACY_GPU = 0.378
MAX_MLM_ACCURACY_GPU = 0.388

# Pretrain next sentence prediction accuracy range:
MIN_NSP_ACCURACY_GPU = 0.82
MAX_NSP_ACCURACY_GPU = 0.84


Chen Chen's avatar
Chen Chen committed
54
55
56
57
58
59
60
61
62
63
64
BERT_PRETRAIN_FILES_SEQ128 = 'gs://mlcompass-data/bert/pretraining_data/seq_128/wikipedia.tfrecord*,gs://mlcompass-data/bert/pretraining_data/seq_128/books.tfrecord*'
BERT_BASE_CONFIG_FILE = 'gs://cloud-tpu-checkpoints/bert/keras_bert/uncased_L-12_H-768_A-12/bert_config.json'

FLAGS = flags.FLAGS


class BertPretrainAccuracyBenchmark(bert_benchmark_utils.BertBenchmarkBase):
  """Benchmark accuracy tests for BERT Pretraining."""

  def __init__(self,
               output_dir: Optional[str] = None,
Chen Chen's avatar
Chen Chen committed
65
66
               tpu: Optional[str] = None,
               **kwargs):
Chen Chen's avatar
Chen Chen committed
67
68
69
70
71
    """Inits BertPretrainAccuracyBenchmark class.

    Args:
      output_dir: Directory where to output e.g. log files
      tpu: TPU name to use in a TPU benchmark.
Chen Chen's avatar
Chen Chen committed
72
      **kwargs: Additional keyword arguments.
Chen Chen's avatar
Chen Chen committed
73
74
    """
    super(BertPretrainAccuracyBenchmark, self).__init__(
Chen Chen's avatar
Chen Chen committed
75
        output_dir=output_dir, tpu=tpu, **kwargs)
Chen Chen's avatar
Chen Chen committed
76
77

  @benchmark_wrappers.enable_runtime_flags
Zongwei Zhou's avatar
Zongwei Zhou committed
78
79
  def _run_and_report_benchmark(self, summary_path: str, report_accuracy: bool,
                                ds_type: str):
Chen Chen's avatar
Chen Chen committed
80
81
    """Runs and reports the benchmark given the provided configuration."""
    distribution = distribution_utils.get_distribution_strategy(
Zongwei Zhou's avatar
Zongwei Zhou committed
82
        distribution_strategy=ds_type, tpu_address=self.tpu)
Chen Chen's avatar
Chen Chen committed
83
84
85
86
87
88
89
90
    logging.info('Flags: %s', flags_core.get_nondefault_flags_as_str())
    start_time_sec = time.time()
    run_pretraining.run_bert_pretrain(
        strategy=distribution, custom_callbacks=self.timer_callback)
    wall_time_sec = time.time() - start_time_sec

    with tf.io.gfile.GFile(summary_path, 'rb') as reader:
      summary = json.loads(reader.read().decode('utf-8'))
Jing Li's avatar
Jing Li committed
91
    self._report_benchmark(summary, start_time_sec, wall_time_sec,
Zongwei Zhou's avatar
Zongwei Zhou committed
92
                           report_accuracy, ds_type)
Chen Chen's avatar
Chen Chen committed
93

Jing Li's avatar
Jing Li committed
94
  def _report_benchmark(self, summary, start_time_sec, wall_time_sec,
Zongwei Zhou's avatar
Zongwei Zhou committed
95
                        report_accuracy, ds_type):
Chen Chen's avatar
Chen Chen committed
96
97
98
99
100
    metrics = [{
        'name': 'train_loss',
        'value': summary['train_loss'],
    }, {
        'name':
Jing Li's avatar
Jing Li committed
101
            'exp_per_second',
Chen Chen's avatar
Chen Chen committed
102
103
104
105
106
107
108
        'value':
            self.timer_callback.get_examples_per_sec(FLAGS.train_batch_size *
                                                     FLAGS.steps_per_loop)
    }, {
        'name': 'startup_time',
        'value': self.timer_callback.get_startup_time(start_time_sec)
    }]
Jing Li's avatar
Jing Li committed
109
    if report_accuracy:
Zongwei Zhou's avatar
Zongwei Zhou committed
110
111
112
113
114
115
116
117
118
119
      if ds_type == 'tpu':
        min_mlm_acc = MIN_MLM_ACCURACY
        max_mlm_acc = MAX_MLM_ACCURACY
        min_nsp_acc = MIN_NSP_ACCURACY
        max_nsp_acc = MAX_NSP_ACCURACY
      else:
        min_mlm_acc = MIN_MLM_ACCURACY_GPU
        max_mlm_acc = MAX_MLM_ACCURACY_GPU
        min_nsp_acc = MIN_NSP_ACCURACY_GPU
        max_nsp_acc = MAX_NSP_ACCURACY_GPU
Jing Li's avatar
Jing Li committed
120
121
122
      metrics.extend([{
          'name': 'masked_lm_accuracy',
          'value': summary['masked_lm_accuracy'],
Zongwei Zhou's avatar
Zongwei Zhou committed
123
124
          'min_value': min_mlm_acc,
          'max_value': max_mlm_acc,
Jing Li's avatar
Jing Li committed
125
126
127
      }, {
          'name': 'next_sentence_accuracy',
          'value': summary['next_sentence_accuracy'],
Zongwei Zhou's avatar
Zongwei Zhou committed
128
129
          'min_value': min_nsp_acc,
          'max_value': max_nsp_acc,
Jing Li's avatar
Jing Li committed
130
      }])
Chen Chen's avatar
Chen Chen committed
131
132
133
134
135
136
137
138
139
140
    self.report_benchmark(
        iters=summary['total_training_steps'],
        wall_time=wall_time_sec,
        metrics=metrics,
        extras={'flags': flags_core.get_nondefault_flags_as_str()})

  def _specify_common_flags(self):
    FLAGS.bert_config_file = BERT_BASE_CONFIG_FILE
    FLAGS.learning_rate = 1e-4
    FLAGS.warmup_steps = 10000
Chen Chen's avatar
Chen Chen committed
141
    FLAGS.steps_per_loop = 10000
Chen Chen's avatar
Chen Chen committed
142
143
144
    FLAGS.input_files = BERT_PRETRAIN_FILES_SEQ128
    FLAGS.max_seq_length = 128
    FLAGS.max_predictions_per_seq = 20
Zongwei Zhou's avatar
Zongwei Zhou committed
145
146
147

  def _specify_tpu_common_flags(self):
    FLAGS.distribution_strategy = 'tpu'
Chen Chen's avatar
Chen Chen committed
148
149
    FLAGS.dtype = 'bf16'

Zongwei Zhou's avatar
Zongwei Zhou committed
150
151
152
153
154
  def _specify_gpu_common_flags(self):
    FLAGS.distribution_strategy = 'mirrored'
    FLAGS.dtype = 'fp16'
    FLAGS.loss_scale = 'dynamic'

Jing Li's avatar
Jing Li committed
155
  @owner_utils.Owner('tf-model-garden')
Chen Chen's avatar
Chen Chen committed
156
157
  def benchmark_accuracy_8x8_tpu_bf16_seq128_500k_steps(self):
    """Test bert pretraining with 8x8 TPU for 500k steps."""
Chen Chen's avatar
Chen Chen committed
158
159
160
    # This is used for accuracy test.
    self._setup()
    self._specify_common_flags()
Zongwei Zhou's avatar
Zongwei Zhou committed
161
162
    self._specify_tpu_common_flags()
    FLAGS.train_batch_size = 512
Chen Chen's avatar
Chen Chen committed
163
    FLAGS.num_steps_per_epoch = 500000
Chen Chen's avatar
Chen Chen committed
164
    FLAGS.num_train_epochs = 1
Chen Chen's avatar
Chen Chen committed
165
    FLAGS.model_dir = self._get_model_dir(
Chen Chen's avatar
Chen Chen committed
166
        'benchmark_accuracy_8x8_tpu_bf16_seq128_500k_steps')
Chen Chen's avatar
Chen Chen committed
167
168
    summary_path = os.path.join(FLAGS.model_dir,
                                'summaries/training_summary.txt')
Chen Chen's avatar
Chen Chen committed
169
170
171
172
    # Set train_summary_interval to -1 to disable training summary, because
    # writing summary to gcs may fail and summaries are not needed for this
    # accuracy benchmark test.
    FLAGS.train_summary_interval = -1
Hongkun Yu's avatar
Hongkun Yu committed
173
    self._run_and_report_benchmark(
Zongwei Zhou's avatar
Zongwei Zhou committed
174
175
176
        summary_path=summary_path,
        report_accuracy=True,
        ds_type=FLAGS.distribution_strategy)
Chen Chen's avatar
Chen Chen committed
177

Allen Wang's avatar
Allen Wang committed
178
179
180
181
182
  @owner_utils.Owner('tf-model-garden')
  def benchmark_perf_2x2_tpu_bf16_seq128_10k_steps(self):
    """Test bert pretraining with 2x2 TPU for 10000 steps."""
    self._setup()
    self._specify_common_flags()
Zongwei Zhou's avatar
Zongwei Zhou committed
183
    self._specify_tpu_common_flags()
Allen Wang's avatar
Allen Wang committed
184
185
186
187
188
189
190
191
192
    FLAGS.num_steps_per_epoch = 5000
    FLAGS.num_train_epochs = 2
    FLAGS.train_batch_size = 128
    FLAGS.model_dir = self._get_model_dir(
        'benchmark_perf_2x2_tpu_bf16_seq128_10k_steps')
    summary_path = os.path.join(FLAGS.model_dir,
                                'summaries/training_summary.txt')
    # Disable accuracy check.
    self._run_and_report_benchmark(
Zongwei Zhou's avatar
Zongwei Zhou committed
193
194
195
        summary_path=summary_path,
        report_accuracy=False,
        ds_type=FLAGS.distribution_strategy)
Allen Wang's avatar
Allen Wang committed
196
197
198
199
200
201

  @owner_utils.Owner('tf-model-garden')
  def benchmark_perf_2x2_tpu_bf16_seq128_10k_steps_mlir(self):
    """Test bert pretraining with 2x2 TPU with MLIR for 10000 steps."""
    self._setup()
    self._specify_common_flags()
Zongwei Zhou's avatar
Zongwei Zhou committed
202
    self._specify_tpu_common_flags()
Allen Wang's avatar
Allen Wang committed
203
204
205
206
207
208
209
210
211
212
    FLAGS.num_steps_per_epoch = 5000
    FLAGS.num_train_epochs = 2
    FLAGS.train_batch_size = 128
    FLAGS.model_dir = self._get_model_dir(
        'benchmark_perf_2x2_tpu_bf16_seq128_10k_steps_mlir')
    summary_path = os.path.join(FLAGS.model_dir,
                                'summaries/training_summary.txt')
    tf.config.experimental.enable_mlir_bridge()
    # Disable accuracy check.
    self._run_and_report_benchmark(
Zongwei Zhou's avatar
Zongwei Zhou committed
213
214
215
        summary_path=summary_path,
        report_accuracy=False,
        ds_type=FLAGS.distribution_strategy)
Allen Wang's avatar
Allen Wang committed
216

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
217
218
219
220
221
  @owner_utils.Owner('tf-model-garden')
  def benchmark_perf_4x4_tpu_bf16_seq128_10k_steps(self):
    """Test bert pretraining with 4x4 TPU for 10000 steps."""
    self._setup()
    self._specify_common_flags()
Zongwei Zhou's avatar
Zongwei Zhou committed
222
223
    self._specify_tpu_common_flags()
    FLAGS.train_batch_size = 512
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
224
225
226
227
228
229
230
231
    FLAGS.num_steps_per_epoch = 5000
    FLAGS.num_train_epochs = 2
    FLAGS.model_dir = self._get_model_dir(
        'benchmark_perf_4x4_tpu_bf16_seq128_10k_steps')
    summary_path = os.path.join(FLAGS.model_dir,
                                'summaries/training_summary.txt')
    # Disable accuracy check.
    self._run_and_report_benchmark(
Zongwei Zhou's avatar
Zongwei Zhou committed
232
233
234
        summary_path=summary_path,
        report_accuracy=False,
        ds_type=FLAGS.distribution_strategy)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
235

Allen Wang's avatar
Allen Wang committed
236
237
238
239
240
  @owner_utils.Owner('tf-model-garden')
  def benchmark_perf_4x4_tpu_bf16_seq128_10k_steps_mlir(self):
    """Test bert pretraining with 4x4 TPU with MLIR for 10000 steps."""
    self._setup()
    self._specify_common_flags()
Zongwei Zhou's avatar
Zongwei Zhou committed
241
242
    self._specify_tpu_common_flags()
    FLAGS.train_batch_size = 512
Allen Wang's avatar
Allen Wang committed
243
244
245
246
247
248
249
250
251
    FLAGS.num_steps_per_epoch = 5000
    FLAGS.num_train_epochs = 2
    FLAGS.model_dir = self._get_model_dir(
        'benchmark_perf_4x4_tpu_bf16_seq128_10k_steps_mlir')
    summary_path = os.path.join(FLAGS.model_dir,
                                'summaries/training_summary.txt')
    tf.config.experimental.enable_mlir_bridge()
    # Disable accuracy check.
    self._run_and_report_benchmark(
Zongwei Zhou's avatar
Zongwei Zhou committed
252
253
254
        summary_path=summary_path,
        report_accuracy=False,
        ds_type=FLAGS.distribution_strategy)
Allen Wang's avatar
Allen Wang committed
255

Jing Li's avatar
Jing Li committed
256
257
258
  @owner_utils.Owner('tf-model-garden')
  def benchmark_perf_8x8_tpu_bf16_seq128_10k_steps(self):
    """Test bert pretraining with 8x8 TPU for 10000 steps."""
Chen Chen's avatar
Chen Chen committed
259
260
    self._setup()
    self._specify_common_flags()
Zongwei Zhou's avatar
Zongwei Zhou committed
261
262
    self._specify_tpu_common_flags()
    FLAGS.train_batch_size = 512
Jing Li's avatar
Jing Li committed
263
264
    FLAGS.num_steps_per_epoch = 5000
    FLAGS.num_train_epochs = 2
Chen Chen's avatar
Chen Chen committed
265
    FLAGS.model_dir = self._get_model_dir(
Jing Li's avatar
Jing Li committed
266
        'benchmark_perf_8x8_tpu_bf16_seq128_10k_steps')
Chen Chen's avatar
Chen Chen committed
267
268
    summary_path = os.path.join(FLAGS.model_dir,
                                'summaries/training_summary.txt')
Jing Li's avatar
Jing Li committed
269
    # Disable accuracy check.
Hongkun Yu's avatar
Hongkun Yu committed
270
    self._run_and_report_benchmark(
Zongwei Zhou's avatar
Zongwei Zhou committed
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
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
332
333
334
335
336
337
338
339
        summary_path=summary_path,
        report_accuracy=False,
        ds_type=FLAGS.distribution_strategy)

  @owner_utils.Owner('tf-dist-strat')
  def benchmark_accuracy_1x8_gpu_fp16_seq128_15k_steps(self):
    """Test bert pretraining with 8 GPU for 15k steps."""
    # This is used for accuracy test.
    self._setup()
    self._specify_common_flags()
    self._specify_gpu_common_flags()
    FLAGS.train_batch_size = 96
    FLAGS.num_steps_per_epoch = 5000
    FLAGS.num_train_epochs = 3
    FLAGS.steps_per_loop = 5000
    FLAGS.model_dir = self._get_model_dir(
        'benchmark_accuracy_1x8_gpu_fp16_seq128_15k_steps')
    summary_path = os.path.join(FLAGS.model_dir,
                                'summaries/training_summary.txt')
    # Set train_summary_interval to -1 to disable training summary, because
    # writing summary to gcs may fail and summaries are not needed for this
    # accuracy benchmark test.
    FLAGS.train_summary_interval = -1
    self._run_and_report_benchmark(
        summary_path=summary_path,
        report_accuracy=True,
        ds_type=FLAGS.distribution_strategy)

  @owner_utils.Owner('tf-dist-strat')
  def benchmark_perf_1x1_gpu_fp16_seq128_200_steps(self):
    """Test bert pretraining with 1 GPU for 200 steps."""
    self._setup()
    self._specify_common_flags()
    self._specify_gpu_common_flags()
    FLAGS.num_steps_per_epoch = 200
    FLAGS.num_train_epochs = 1
    FLAGS.num_gpus = 1
    FLAGS.train_batch_size = 12
    FLAGS.steps_per_loop = 100
    FLAGS.model_dir = self._get_model_dir(
        'benchmark_perf_1x1_gpu_fp16_seq128_200_steps')
    summary_path = os.path.join(FLAGS.model_dir,
                                'summaries/training_summary.txt')
    # Disable accuracy check.
    self._run_and_report_benchmark(
        summary_path=summary_path,
        report_accuracy=False,
        ds_type=FLAGS.distribution_strategy)

  @owner_utils.Owner('tf-dist-strat')
  def benchmark_perf_1x8_gpu_fp16_seq128_200_steps(self):
    """Test bert pretraining with 8 GPU for 200 steps."""
    self._setup()
    self._specify_common_flags()
    self._specify_gpu_common_flags()
    FLAGS.num_steps_per_epoch = 200
    FLAGS.num_train_epochs = 1
    FLAGS.num_gpus = 8
    FLAGS.train_batch_size = 96
    FLAGS.steps_per_loop = 100
    FLAGS.model_dir = self._get_model_dir(
        'benchmark_perf_1x8_gpu_fp16_seq128_200_steps')
    summary_path = os.path.join(FLAGS.model_dir,
                                'summaries/training_summary.txt')
    # Disable accuracy check.
    self._run_and_report_benchmark(
        summary_path=summary_path,
        report_accuracy=False,
        ds_type=FLAGS.distribution_strategy)
Chen Chen's avatar
Chen Chen committed
340
341
342
343


if __name__ == '__main__':
  tf.test.main()