Add distribution_strategy and all_reduce_alg flags to TensorFlow BERT pretraining

[rapsealk]

Hello mlcommons team!

I've noticed that some utilities, such as additional flags, are missing in BERT pretraining, unlike ResNet50 image classification. This pull request is expected to be helpful to run BERT training on distributed environment.

References are below:

training/image_classification/tensorflow2/tf2_common/utils/flags/_base.py

Lines 140 to 150 in f0a7d0c

	if distribution_strategy:
	flags.DEFINE_string(
	name="distribution_strategy", short_name="ds", default="mirrored",
	help=help_wrap("The Distribution Strategy to use for training. "
	"Accepted values are 'off', 'one_device', "
	"'mirrored', 'parameter_server', 'collective', "
	"case insensitive. 'off' means not to use "
	"Distribution Strategy; 'default' means to choose "
	"from `MirroredStrategy` or `OneDeviceStrategy` "
	"according to the number of GPUs.")
	)

training/image_classification/tensorflow2/tf2_common/utils/flags/_performance.py

Lines 224 to 234 in f0a7d0c

	if all_reduce_alg:
	flags.DEFINE_string(
	name="all_reduce_alg", short_name="ara", default=None,
	help=help_wrap("Defines the algorithm to use for performing all-reduce."
	"When specified with MirroredStrategy for single "
	"worker, this controls "
	"tf.contrib.distribute.AllReduceCrossTowerOps. When "
	"specified with MultiWorkerMirroredStrategy, this "
	"controls "
	"tf.distribute.experimental.CollectiveCommunication; "
	"valid options are `ring` and `nccl`."))

Refs

• [BERT] Multi-gpu, eval and step counting #384

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅