Question about the calculation of train_batch_size

[formath]

In the doc https://www.deepspeed.ai/docs/config-json/#batch-size-related-parameters, train_batch_size is calculated in this way:

train_batch_size =  train_micro_batch_size_per_gpu * gradient_accumulation * number of GPUs

However, in pipeline mode, rightly it should be like this:

train_batch_size = train_micro_batch_size_per_gpu * gradient_accumulation * number of GPUs / number of stages

What is the difference between those?

[ShadenSmith]

Hi @formath, great question 😸. In both cases the last term represents the degree of data parallelism used in training (since each data-parallel replica will have its own data pipeline contributing to the train_batch_size.

With pure data parallelism, that's simply the number of GPUs. If we are also training with pipeline (model) parallelism, the resulting degree of data parallelism is the number of GPUs divided by the number of pipeline stages.