T5 pretraining grows slower if micro batch size > 1

megatron version: v2.6
pytorch version: 1.10.0+cu111
cluster: 2 nodes, A100*8 for each node
script: pretrain_t5_distributed_with_mp.sh

args:

using world size: 16, data-parallel-size: 4, tensor-model-parallel size: 4, pipeline-model-parallel size: 1
using torch.float16 for parameters ...
Persistent fused layer norm kernel is supported from pytorch v1.11 (nvidia pytorch container paired with v1.11). Defaulting to no_persist_layer_norm=True
------------------------ arguments ------------------------
accumulate_allreduce_grads_in_fp32 .............. False
activations_checkpoint_method ................... None
activations_checkpoint_num_layers ............... 1
adam_beta1 ...................................... 0.9
adam_beta2 ...................................... 0.999
adam_eps ........................................ 1e-08
adlr_autoresume ................................. False
adlr_autoresume_interval ........................ 1000
apply_query_key_layer_scaling ................... True
apply_residual_connection_post_layernorm ........ False
attention_dropout ............................... 0.1
attention_softmax_in_fp32 ....................... False
bert_load ....................................... None
bf16 ............................................ False
bias_dropout_fusion ............................. True
bias_gelu_fusion ................................ True
biencoder_projection_dim ........................ 0
biencoder_shared_query_context_model ............ False
block_data_path ................................. None
clip_grad ....................................... 1.0
consumed_train_samples .......................... 0
consumed_valid_samples .......................... 0
data_impl ....................................... mmap
data_parallel_size .............................. 4
data_path ....................................... ['./processed_data_sample/t5data_text_sentence']
dataloader_type ................................. single
DDP_impl ........................................ local
deallocate_pipeline_outputs ..................... False
decoder_seq_length .............................. 128
distribute_checkpointed_activations ............. False
distributed_backend ............................. nccl
embedding_path .................................. None
empty_unused_memory_level ....................... 0
encoder_seq_length .............................. 512
eod_mask_loss ................................... False
eval_interval ................................... 1000
eval_iters ...................................... 100
evidence_data_path .............................. None
exit_duration_in_mins ........................... None
exit_interval ................................... None
exit_signal_handler ............................. False
ffn_hidden_size ................................. 10240
finetune ........................................ False
fp16 ............................................ True
fp16_lm_cross_entropy ........................... False
fp32_residual_connection ........................ False
global_batch_size ............................... 8
hidden_dropout .................................. 0.1
hidden_size ..................................... 2048
hysteresis ...................................... 2
ict_head_size ................................... None
ict_load ........................................ None
indexer_batch_size .............................. 128
indexer_log_interval ............................ 1000
init_method_std ................................. 0.02
init_method_xavier_uniform ...................... False
initial_loss_scale .............................. 4294967296
kv_channels ..................................... 64
layernorm_epsilon ............................... 1e-05
lazy_mpu_init ................................... None
load ............................................ None
local_rank ...................................... None
log_batch_size_to_tensorboard ................... False
log_interval .................................... 10
log_learning_rate_to_tensorboard ................ True
log_loss_scale_to_tensorboard ................... True
log_memory_to_tensorboard ....................... False
log_num_zeros_in_grad ........................... False
log_params_norm ................................. False
log_timers_to_tensorboard ....................... False
log_validation_ppl_to_tensorboard ............... False
log_world_size_to_tensorboard ................... False
loss_scale ...................................... None
loss_scale_window ............................... 1000
lr .............................................. 0.0001
lr_decay_iters .................................. 1000000
lr_decay_samples ................................ None
lr_decay_style .................................. linear
lr_warmup_fraction .............................. None
lr_warmup_iters ................................. 0
lr_warmup_samples ............................... 0
make_vocab_size_divisible_by .................... 128
mask_prob ....................................... 0.15
masked_softmax_fusion ........................... True
max_position_embeddings ......................... 4096
merge_file ...................................... None
micro_batch_size ................................ 1
min_loss_scale .................................. 1.0
min_lr .......................................... 1e-05
mmap_warmup ..................................... False
no_async_tensor_model_parallel_allreduce ........ False
no_load_optim ................................... None
no_load_rng ..................................... None
no_persist_layer_norm ........................... True
no_save_optim ................................... None
no_save_rng ..................................... None
num_attention_heads ............................. 64
num_layers ...................................... 24
num_layers_per_virtual_pipeline_stage ........... None
num_workers ..................................... 2
onnx_safe ....................................... None
openai_gelu ..................................... False
optimizer ....................................... adam
override_lr_scheduler ........................... False
params_dtype .................................... torch.float16
pipeline_model_parallel_size .................... 1
pipeline_model_parallel_split_rank .............. None
query_in_block_prob ............................. 0.1
rampup_batch_size ............................... None
rank ............................................ 0
reset_attention_mask ............................ False
reset_position_ids .............................. False
retriever_score_scaling ......................... False
retriever_seq_length ............................ 256
sample_rate ..................................... 1.0
save ............................................ None
save_interval ................................... None
scatter_gather_tensors_in_pipeline .............. True
seed ............................................ 1234
seq_length ...................................... 512
sgd_momentum .................................... 0.9
short_seq_prob .................................. 0.1
split ........................................... 1000,0,0
tensor_model_parallel_size ...................... 4
tensorboard_dir ................................. ./log_dir
tensorboard_log_interval ........................ 1
tensorboard_queue_size .......................... 1000
titles_data_path ................................ None
tokenizer_type .................................. BertWordPieceLowerCase
train_iters ..................................... 1000
train_samples ................................... None
use_checkpoint_lr_scheduler ..................... False
use_contiguous_buffers_in_local_ddp ............. True
use_cpu_initialization .......................... None
use_one_sent_docs ............................... False
virtual_pipeline_model_parallel_size ............ None
vocab_extra_ids ................................. 100
vocab_file ...................................... ./t5_vocab.txt
weight_decay .................................... 0.01
world_size ...................................... 16
-------------------- end of arguments ---------------------

with micro_batch_size=1, global_batch_size=8, the throughput seems stable:

iteration       10/    1000 | consumed samples:           80 | elapsed time per iteration (ms): 4196.4 | learning rate: 1.000E-04 | global batch size:     8 | loss scale: 8388608.0 | number of skipped iterations:  10 | number of nan iterations:   0 |
iteration       20/    1000 | consumed samples:          160 | elapsed time per iteration (ms): 3997.8 | learning rate: 1.000E-04 | global batch size:     8 | lm loss: 1.173781E+01 | loss scale: 32768.0 | grad norm: 84.084 | number of skipped iterations:   8 | number of nan iterations:   0 |
iteration       30/    1000 | consumed samples:          240 | elapsed time per iteration (ms): 3892.9 | learning rate: 1.000E-04 | global batch size:     8 | lm loss: 1.171586E+01 | loss scale: 16384.0 | grad norm: 31.831 | number of skipped iterations:   1 | number of nan iterations:   0 |
iteration       40/    1000 | consumed samples:          320 | elapsed time per iteration (ms): 3827.2 | learning rate: 1.000E-04 | global batch size:     8 | lm loss: 7.473515E+00 | loss scale: 16384.0 | grad norm: 7.100 | number of skipped iterations:   0 | number of nan iterations:   0 |
iteration       50/    1000 | consumed samples:          400 | elapsed time per iteration (ms): 3738.2 | learning rate: 1.000E-04 | global batch size:     8 | lm loss: 6.682758E+00 | loss scale: 16384.0 | grad norm: 8.118 | number of skipped iterations:   0 | number of nan iterations:   0 |
iteration       60/    1000 | consumed samples:          480 | elapsed time per iteration (ms): 3853.5 | learning rate: 1.000E-04 | global batch size:     8 | lm loss: 6.173346E+00 | loss scale: 16384.0 | grad norm: 5.354 | number of skipped iterations:   0 | number of nan iterations:   0 |
iteration       70/    1000 | consumed samples:          560 | elapsed time per iteration (ms): 3822.8 | learning rate: 1.000E-04 | global batch size:     8 | lm loss: 5.929058E+00 | loss scale: 16384.0 | grad norm: 7.846 | number of skipped iterations:   0 | number of nan iterations:   0 |

time (ms) | model-and-optimizer-setup: 1504.44 | train/valid/test-data-iterators-setup: 2705.71
time (ms) | forward-compute: 348.18 | backward-compute: 303.54 | backward-params-all-reduce: 2652.54 | backward-embedding-all-reduce: 0.04 | optimizer-copy-to-main-grad: 24.35 | optimizer-unscale-and-check-inf: 858.74 | optimizer: 883.22
time (ms) | forward-compute: 185.87 | backward-compute: 253.68 | backward-params-all-reduce: 2703.11 | backward-embedding-all-reduce: 0.03 | optimizer-copy-to-main-grad: 9.98 |optimizer-unscale-and-check-inf: 606.55 | optimizer-clip-main-grad: 3.61 | optimizer-copy-main-to-model-params: 2.55 | optimizer: 826.70
time (ms) | forward-compute: 264.38 | backward-compute: 256.96 | backward-params-all-reduce: 2567.96 | backward-embedding-all-reduce: 0.05 | optimizer-copy-to-main-grad: 11.13 | optimizer-unscale-and-check-inf: 754.93 | optimizer-clip-main-grad: 13.88 | optimizer-copy-main-to-model-params: 9.93 | optimizer: 814.71
time (ms) | forward-compute: 172.06 | backward-compute: 282.55 | backward-params-all-reduce: 2490.93 | backward-embedding-all-reduce: 0.04 | optimizer-copy-to-main-grad: 10.51 | optimizer-unscale-and-check-inf: 807.80 | optimizer-clip-main-grad: 16.00 | optimizer-copy-main-to-model-params: 11.08 | optimizer: 873.01
time (ms) | forward-compute: 170.71 | backward-compute: 259.05 | backward-params-all-reduce: 2814.66 | backward-embedding-all-reduce: 0.03 | optimizer-copy-to-main-grad: 10.11 | optimizer-unscale-and-check-inf: 421.80 | optimizer-clip-main-grad: 15.47 | optimizer-copy-main-to-model-params: 10.95 | optimizer: 485.84
time (ms) | forward-compute: 205.38 | backward-compute: 252.87 | backward-params-all-reduce: 2631.29 | backward-embedding-all-reduce: 0.03 | optimizer-copy-to-main-grad: 24.72 | optimizer-unscale-and-check-inf: 658.24 | optimizer-clip-main-grad: 16.23 | optimizer-copy-main-to-model-params: 11.00 | optimizer: 737.69
time (ms) | forward-compute: 241.89 | backward-compute: 265.85 | backward-params-all-reduce: 2698.72 | backward-embedding-all-reduce: 0.04 | optimizer-copy-to-main-grad: 10.39 | optimizer-unscale-and-check-inf: 561.97 | optimizer-clip-main-grad: 15.27 | optimizer-copy-main-to-model-params: 10.97 | optimizer: 626.12
time (ms) | forward-compute: 174.71 | backward-compute: 295.77 | backward-params-all-reduce: 2515.41 | backward-embedding-all-reduce: 0.04 | optimizer-copy-to-main-grad: 10.55 | optimizer-unscale-and-check-inf: 744.83 | optimizer-clip-main-grad: 15.62 | optimizer-copy-main-to-model-params: 11.14 | optimizer: 809.86
time (ms) | forward-compute: 178.27 | backward-compute: 270.11 | backward-params-all-reduce: 2630.00 | backward-embedding-all-reduce: 0.03 | optimizer-copy-to-main-grad: 10.03 | optimizer-unscale-and-check-inf: 646.81 | optimizer-clip-main-grad: 15.87 | optimizer-copy-main-to-model-params: 11.00 | optimizer: 711.28

But if I change micro_batch_size to 2 and global_batch_size to 16, the throughput drops during training:

 iteration       10/    1000 | consumed samples:          160 | elapsed time per iteration (ms): 4847.7 | learning rate: 1.000E-04 | global batch size:    16 | loss scale: 8388608.0 | number of skipped iterations:  10 | number of nan iterations:   0 |
 iteration       20/    1000 | consumed samples:          320 | elapsed time per iteration (ms): 5375.5 | learning rate: 1.000E-04 | global batch size:    16 | lm loss: 1.031812E+01 | loss scale: 16384.0 | number of skipped iterations:   9 | number of nan iterations:   0 |
 iteration       30/    1000 | consumed samples:          480 | elapsed time per iteration (ms): 6150.0 | learning rate: 1.000E-04 | global batch size:    16 | lm loss: 1.184227E+01 | loss scale: 16384.0 | grad norm: 16.206 | number of skipped iterations:   0 | number of nan iterations:   0 |
 iteration       40/    1000 | consumed samples:          640 | elapsed time per iteration (ms): 7027.5 | learning rate: 1.000E-04 | global batch size:    16 | lm loss: 7.400169E+00 | loss scale: 16384.0 | grad norm: 7.939 | number of skipped iterations:   0 | number of nan iterations:   0 |
 iteration       50/    1000 | consumed samples:          800 | elapsed time per iteration (ms): 7873.5 | learning rate: 1.000E-04 | global batch size:    16 | lm loss: 6.508469E+00 | loss scale: 16384.0 | grad norm: 4.846 | number of skipped iterations:   0 | number of nan iterations:   0 |
 iteration       60/    1000 | consumed samples:          960 | elapsed time per iteration (ms): 8804.5 | learning rate: 1.000E-04 | global batch size:    16 | lm loss: 6.107175E+00 | loss scale: 16384.0 | grad norm: 5.190 | number of skipped iterations:   0 | number of nan iterations:   0 |
 iteration       70/    1000 | consumed samples:         1120 | elapsed time per iteration (ms): 9476.2 | learning rate: 1.000E-04 | global batch size:    16 | lm loss: 5.975057E+00 | loss scale: 16384.0 | grad norm: 3.755 | number of skipped iterations:   0 | number of nan iterations:   0 |

time (ms) | model-and-optimizer-setup: 1736.32 | train/valid/test-data-iterators-setup: 2674.95
time (ms) | forward-compute: 604.56 | backward-compute: 701.16 | backward-params-all-reduce: 2672.65 | backward-embedding-all-reduce: 0.13 | optimizer-copy-to-main-grad: 21.18 | optimizer-unscale-and-check-inf: 835.71 | optimizer: 857.27
time (ms) | forward-compute: 864.01 | backward-compute: 971.21 | backward-params-all-reduce: 2805.27 | backward-embedding-all-reduce: 0.09 | optimizer-copy-to-main-grad: 12.70 | optimizer-unscale-and-check-inf: 498.17 | optimizer-clip-main-grad: 1.55 | optimizer-copy-main-to-model-params: 1.47 | optimizer: 724.70
time (ms) | forward-compute: 1304.65 | backward-compute: 1428.05 | backward-params-all-reduce: 3067.17 | backward-embedding-all-reduce: 0.04 | optimizer-copy-to-main-grad: 10.33 | optimizer-unscale-and-check-inf: 277.17 | optimizer-clip-main-grad: 15.31 | optimizer-copy-main-to-model-params: 10.97 | optimizer: 341.30
time (ms) | forward-compute: 1664.07 | backward-compute: 1881.06 | backward-params-all-reduce: 3248.30 | backward-embedding-all-reduce: 0.05 | optimizer-copy-to-main-grad: 10.96 | optimizer-unscale-and-check-inf: 161.09 | optimizer-clip-main-grad: 15.23 | optimizer-copy-main-to-model-params: 10.95 | optimizer: 225.78
time (ms) | forward-compute: 2120.17 | backward-compute: 2327.50 | backward-params-all-reduce: 3114.42 | backward-embedding-all-reduce: 0.04 | optimizer-copy-to-main-grad: 10.13 | optimizer-unscale-and-check-inf: 238.10 | optimizer-clip-main-grad: 15.77 | optimizer-copy-main-to-model-params: 10.94 | optimizer: 302.45
time (ms) | forward-compute: 2709.59 | backward-compute: 2703.16 | backward-params-all-reduce: 2872.23 | backward-embedding-all-reduce: 0.04 | optimizer-copy-to-main-grad: 30.98 | optimizer-unscale-and-check-inf: 426.58 | optimizer-clip-main-grad: 15.80 | optimizer-copy-main-to-model-params: 10.95 | optimizer: 511.77
time (ms) | forward-compute: 3070.24 | backward-compute: 3031.89 | backward-params-all-reduce: 2533.92 | backward-embedding-all-reduce: 0.06 | optimizer-copy-to-main-grad: 10.99 | optimizer-unscale-and-check-inf: 765.58 | optimizer-clip-main-grad: 15.78 | optimizer-copy-main-to-model-params: 11.08 | optimizer: 831.14

It seems that forward and backward computing grows slower. Please take a look. Thanks!

NVIDIA / Megatron-LM

T5 pretraining grows slower if micro batch size > 1 #184