REALM pre-training in multi-GPU setup

Hi, thankyou for releasing the code and data.
I am trying to run the original REALM code. I see the default setting is disabled for TPUs. I have a machine with 8 GPUs of 16GB each. I am not able to understand why your scaled down code (with 16 as batch size, which I tried reducing down to 4) is also producing an OOM error on my machine. Do I need to explicitly need to define a train and eval distribute strategy here? https://github.com/google-research/language/blob/d9bdad8d007c53c4cb6c0be3edabcce0393d226c/language/common/utils/experiment_utils.py#L142

Please find the log:
```
2021-06-24 09:37:39.818921: W tensorflow/core/common_runtime/bfc_allocator.cc:419] Allocator (GPU_0_bfc) ran out of memory trying to allocate 64.0KiB (rounded to 65536). Current allocation summary follows.*
2021-06-24 09:37:39.819832: I tensorflow/core/common_runtime/bfc_allocator.cc:921] Sum Total of in-use chunks: 34.25MiB
2021-06-24 09:37:39.819837: I tensorflow/core/common_runtime/bfc_allocator.cc:923] total_region_allocated_bytes_: 35913728 memory_limit_: 35913728 available bytes: 0 curr_region_allocation_bytes_: 71827456
2021-06-24 09:37:39.819846: I tensorflow/core/common_runtime/bfc_allocator.cc:929] Stats:
Limit: 35913728
InUse: 35913728
MaxInUse: 35913728
NumAllocs: 114
MaxAllocSize: 19636736

2021-06-24 09:37:39.819855: W tensorflow/core/common_runtime/bfc_allocator.cc:424] *****************************************************************************************xxxxxxxxxxx
2021-06-24 09:37:39.819883: W tensorflow/core/framework/op_kernel.cc:1628] OP_REQUIRES failed at constant_op.cc:77 : Resource exhausted: OOM when allocating tensor of shape [128,128] and type float
2021-06-24 09:37:39.819929: E tensorflow/core/common_runtime/executor.cc:642] Executor failed to create kernel. Resource exhausted: OOM when allocating tensor of shape [128,128] and type float
[[{{node embedder/module/bert/encoder/layer_1/attention/output/dense/kernel/adam_v/Initializer/zeros}}]]
Traceback (most recent call last):
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/client/session.py", line 1365, in _do_call
return fn(*args)
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/client/session.py", line 1350, in _run_fn
target_list, run_metadata)
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/client/session.py", line 1443, in _call_tf_sessionrun
run_metadata)
tensorflow.python.framework.errors_impl.ResourceExhaustedError: OOM when allocating tensor of shape [128,128] and type float
[[{{node embedder/module/bert/encoder/layer_1/attention/output/dense/kernel/adam_v/Initializer/zeros}}]]

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/opt/conda/envs/try/lib/python3.7/runpy.py", line 193, in _run_module_as_main
"__main__", mod_spec)
File "/opt/conda/envs/try/lib/python3.7/runpy.py", line 85, in _run_code
exec(code, run_globals)
File "/home/jikaur/language/language/realm/train_realm.py", line 112, in
app.run(main)
File "/opt/conda/envs/try/lib/python3.7/site-packages/absl/app.py", line 312, in run
_run_main(main, args)
File "/opt/conda/envs/try/lib/python3.7/site-packages/absl/app.py", line 258, in _run_main
sys.exit(main(argv))
File "/home/jikaur/language/language/realm/train_realm.py", line 94, in main
exporters=model.get_exporters(params))
File "/home/jikaur/language/language/common/utils/experiment_utils.py", line 169, in run_experiment
eval_spec=eval_spec)
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/training.py", line 473, in train_and_evaluate
return executor.run()
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/training.py", line 613, in run
return self.run_local()
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/training.py", line 714, in run_local
saving_listeners=saving_listeners)
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 370, in train
loss = self._train_model(input_fn, hooks, saving_listeners)
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1161, in _train_model
return self._train_model_default(input_fn, hooks, saving_listeners)
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1195, in _train_model_default
saving_listeners)
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1490, in _train_with_estimator_spec
log_step_count_steps=log_step_count_steps) as mon_sess:
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/training/monitored_session.py", line 584, in MonitoredTrainingSession
stop_grace_period_secs=stop_grace_period_secs)
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/training/monitored_session.py", line 1014, in __init__
stop_grace_period_secs=stop_grace_period_secs)
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/training/monitored_session.py", line 725, in __init__
self._sess = _RecoverableSession(self._coordinated_creator)
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/training/monitored_session.py", line 1207, in __init__
_WrappedSession.__init__(self, self._create_session())
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/training/monitored_session.py", line 1212, in _create_session
return self._sess_creator.create_session()
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/training/monitored_session.py", line 878, in create_session
self.tf_sess = self._session_creator.create_session()
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/training/monitored_session.py", line 647, in create_session
init_fn=self._scaffold.init_fn)
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/training/session_manager.py", line 296, in prepare_session
sess.run(init_op, feed_dict=init_feed_dict)
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/client/session.py", line 956, in run
run_metadata_ptr)
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/client/session.py", line 1180, in _run
feed_dict_tensor, options, run_metadata)
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/client/session.py", line 1359, in _do_run
run_metadata)
File "/opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/client/session.py", line 1384, in _do_call
raise type(e)(node_def, op, message)
tensorflow.python.framework.errors_impl.ResourceExhaustedError: OOM when allocating tensor of shape [128,128] and type float
[[node embedder/module/bert/encoder/layer_1/attention/output/dense/kernel/adam_v/Initializer/zeros (defined at opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:1748) ]]

Original stack trace for 'embedder/module/bert/encoder/layer_1/attention/output/dense/kernel/adam_v/Initializer/zeros':
File "opt/conda/envs/try/lib/python3.7/runpy.py", line 193, in _run_module_as_main
"__main__", mod_spec)
File "opt/conda/envs/try/lib/python3.7/runpy.py", line 85, in _run_code
exec(code, run_globals)
File "home/jikaur/language/language/realm/train_realm.py", line 112, in
app.run(main)
File "opt/conda/envs/try/lib/python3.7/site-packages/absl/app.py", line 312, in run
_run_main(main, args)
File "opt/conda/envs/try/lib/python3.7/site-packages/absl/app.py", line 258, in _run_main
sys.exit(main(argv))
File "home/jikaur/language/language/realm/train_realm.py", line 94, in main
exporters=model.get_exporters(params))
File "home/jikaur/language/language/common/utils/experiment_utils.py", line 169, in run_experiment
eval_spec=eval_spec)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/training.py", line 473, in train_and_evaluate
return executor.run()
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/training.py", line 613, in run
return self.run_local()
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/training.py", line 714, in run_local
saving_listeners=saving_listeners)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 370, in train
loss = self._train_model(input_fn, hooks, saving_listeners)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1161, in _train_model
return self._train_model_default(input_fn, hooks, saving_listeners)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1191, in _train_model_default
features, labels, ModeKeys.TRAIN, self.config)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1149, in _call_model_fn
model_fn_results = self._model_fn(features=features, **kwargs)
File "home/jikaur/language/language/realm/model.py", line 186, in model_fn
use_tpu=params["use_tpu"])
File "opt/conda/envs/try/lib/python3.7/site-packages/bert/optimization.py", line 77, in create_optimizer
zip(grads, tvars), global_step=global_step)
File "opt/conda/envs/try/lib/python3.7/site-packages/bert/optimization.py", line 128, in apply_gradients
initializer=tf.zeros_initializer())
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/ops/variable_scope.py", line 1500, in get_variable
aggregation=aggregation)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/ops/variable_scope.py", line 1243, in get_variable
aggregation=aggregation)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/ops/variable_scope.py", line 567, in get_variable
aggregation=aggregation)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/ops/variable_scope.py", line 519, in _true_getter
aggregation=aggregation)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/ops/variable_scope.py", line 933, in _get_single_variable
aggregation=aggregation)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py", line 258, in __call__
return cls._variable_v1_call(*args, **kwargs)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py", line 219, in _variable_v1_call
shape=shape)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py", line 197, in
previous_getter = lambda **kwargs: default_variable_creator(None, **kwargs)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/ops/variable_scope.py", line 2519, in default_variable_creator
shape=shape)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py", line 262, in __call__
return super(VariableMetaclass, cls).__call__(*args, **kwargs)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py", line 1688, in __init__
shape=shape)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py", line 1818, in _init_from_args
initial_value(), name="initial_value", dtype=dtype)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/ops/variable_scope.py", line 905, in
partition_info=partition_info)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/ops/init_ops.py", line 114, in __call__
return array_ops.zeros(shape, dtype)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/ops/array_ops.py", line 2350, in zeros
output = fill(shape, constant(zero, dtype=dtype), name=name)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/ops/array_ops.py", line 171, in fill
result = gen_array_ops.fill(dims, value, name=name)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_array_ops.py", line 3602, in fill
"Fill", dims=dims, value=value, name=name)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/framework/op_def_library.py", line 794, in _apply_op_helper
op_def=op_def)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/util/deprecation.py", line 507, in new_func
return func(*args, **kwargs)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 3357, in create_op
attrs, op_def, compute_device)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 3426, in _create_op_internal
op_def=op_def)
File "opt/conda/envs/try/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 1748, in __init__
self._traceback = tf_stack.extract_stack()
```


*Not added the summary, can add if required.
Thanks a lot.