Which batch size to use with DataLoader

[g-karthik]

I have a question about the below minimal working example from the DeepSpeed README.

for step, batch in enumerate(data_loader):
    #forward() method
    loss = model_engine(batch)

    #runs backpropagation
    model_engine.backward(loss)

    #weight update
    model_engine.step()

Suppose that deepspeed.initialize() was solely called to wrap around my model, and my client script already creates a data_loader. What batch size should be used when instantiating my data_loader for the above loop to work fine?

If there was no DeepSpeed, I'd initialize my data_loader as follows in the client script:

    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) if args.distributed else None
    train_loader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size, shuffle=(not args.distributed))

where args.train_batch_size above corresponds to DeepSpeed's train_micro_batch_size_per_gpu , i.e., the batch size for each GPU for each step, independent of gradient accumulation steps.

Should I continue to use the micro batch size when instantiating my data_loader in the client script? The documentation is not clear on this.

[g-karthik]

@jeffra @ShadenSmith

[tjruwase]

This is a great question, and is certainly an easy source of confusion unfortunately.

When the client handles data loading issues, the uses of train_micro_batch_size_per_gpu in DeepSpeed are restricted to:

1. Detecting gradient accumulation boundaries in order to initiate gradient reduction and optimizer.step()
2. Computing various latency and throughput statistics.

Thus a mismatch in args.train_batch_size (in your example) and train_micro_batch_size_per_gpu could result in serious training bugs because of usage 1, and less serious system perf errors because of usage 2. So yes, please ensure that args.train_batch_size = train_micro_batch_size_per_gpu.

We are exploring ideas to easily detect these kinds of serious mismatches in args and json config settings, and we appreciate your input. Thanks!

[tjruwase]

By the way, we have an open issue on this, in case you are interested in tracking.

[g-karthik]

@tjruwase Thanks for the info!

To be clear, if my client script has a train_batch_size of 2, and I am running my code on a single node with 8 GPUs, then I need to ensure that my DeepSpeed config has:

gradient_accumulation_steps = X
train_batch_size = 2*8*X

Is this correct?

Also, if I wanted to go to multiple nodes, all having 8 GPUs each, what would I update these parameters to?

[tjruwase]

Based on the code snippet you shared earlier, args.train_batch_size is not the effective batch size of training, but rather is the batch size that fills one GPU memory on each call to the data_loader. Thus, args.train_batch_size is the equivalent of train_micro_batch_size_per_gpu. For example in 1-GPU training, your client script will probably be calling the data_loader gradient_accumulation_steps times to accumulate data samples to aggregate an effective batch size (the equivalent of train_batch_size in the json config) before making an optimizer step to update the model parameters.

Can you clarify that I am aligned with your scenario/thinking?

[g-karthik]

@tjruwase yes that is correct -- in that snippet, args is client args. So args.train_batch_size will be defined by me directly in my client script. The "train_batch_size" in the DeepSpeed config file is not the same as this one.

[g-karthik]

So my thinking is that the client args has the following:

params_file_path
device
local_rank
deepspeed
deepspeed_config

params_file_path will point to a JSON file in which I have "train_batch_size": 2, so my client script will have a DataLoader that is instantiated with this batch size before the batches get enumerated.

deepspeed_config will point to a JSON file in which I have "train_batch_size": 2*8*X (where the 8 refers to the number of GPUs in my single-node) and "gradient_accumulation_steps": X. Does this seem fine?

[tjruwase]

Yes, this seems fine. Please let us know how it goes.

[g-karthik]

Cool, thanks for the confirmation! How would this configuration change when I want to use multiple nodes? Let's say I have N nodes, each with 8 GPUs.

[tjruwase]

You will set train_batch_size in JSON file to 2*N*8*X. Please see our multi-node training documentation.

[g-karthik]

Great, thanks! Actually that link does not say anything about the effective batch size, this one does, although it doesn't elaborate on the difference between the single and multi-node cases from the perspective of the effective batch size.

Also, it talks about using train_step_batch_size to compute train_batch_size, but from its description, it sounds exactly the same as train_micro_batch_size_per_gpu . Are these two the same and just haven't been merged together?

Also, my script does validation set metric computation currently after every epoch of training in a distributed fashion and prints the all-reduced metric. It seems like I should not be using the model_engine wrapper for this metric computation, since the forward() implementation in DeepSpeedLight is specific to training. So while I use model_engine variable for training, I should avoid it during validation and just use the client model variable, is that correct?

[tjruwase]

Also, it talks about using train_step_batch_size to compute train_batch_size, but from its description, it sounds exactly the same as train_micro_batch_size_per_gpu . Are these two the same and just haven't been merged together?

Yes, train_step_batch_size was renamed to train_micro_batch_size_per_gpu . Please pardon our typo.

For now, it is okay to use your client model to do validation since it is not a throughput bottleneck. Ideally, we believe DeepSpeed should support that as well. Perhaps, you can helps us verify that support at a more convenient time for you.

[g-karthik]

@tjruwase but how would model_engine support validation as of now? The DeepSpeedLight forward() method does a check for gradient accumulation steps, which is training-specific and not to be done during validation. So we cannot reuse the model_engine for validation, right?

[tjruwase]

I have not thought through all the requirements, but you could set gradient accumulation steps to 1 in this case. In other words, we could have a separate json config field for validation.

[g-karthik]

Having a separate JSON config for validation would imply instantiating a separate model_engine for validation. Seems a bit much.

Or maybe not -- but it will require significant changes to DeepSpeedLight and the parser.