Dynamic/variable batch size support

[ecly]

For the model I am training, I am relying on a custom Sampler, that returns variable batch sizes. My task at hand is translation, where I following Attention is all you need (2017) create batches based on total token count in a batch, which given the variable length input, results in batches of varying numbers of examples (examples here being one source/target text translation pair).

For regular DDP based training, this worked fine, by simply creating a distributed version of this sampler, to split the variable size batch into sub-batches based on the GPU rank. For DeepSpeed however, I am forced to provide either train_micro_batch_size_per_gpu or train_batch_size, both my current understanding tells me are based on the number of examples in the batch.

As the number of examples in my batch varies for each batch, and I just want to configure the accumulation based on batch count, rather than batch size, I'm not sure how to achieve this with DeepSpeed's configuration.

Am I misunderstanding the impact of the configuration variables, missing some other configuration, or is this not possible to achieve at the moment?

[aniruddhakal]

I'm also super interested in knowing more about this. Happy to lend a hand, so to see it available faster!

[tjruwase]

@ecly, apologies that this request somehow slipped through. I wonder what solution you ended up with and whether you are still interested in DeepSpeed support?

@aniruddhakal, thanks for bumping this back to our attention. Is it okay to wait @ecly to respond in order to decide next steps?

[ecly]

@tjruwase we ended up for the most part just using pure DDP in PyTorch. We did have moderate success using Fairscale which supported the variable batch sizes out of the box, but didn't find any value from doing so with our model size of ~200M parameters. The problem is very much still relevant to us, as we'd still like to further adopt DeepSpeed, but this is a blocker that makes it non-trivial for us to adopt.

[wangleiofficial]

@tjruwase we ended up for the most part just using pure DDP in PyTorch. We did have moderate success using Fairscale which supported the variable batch sizes out of the box, but didn't find any value from doing so with our model size of ~200M parameters. The problem is very much still relevant to us, as we'd still like to further adopt DeepSpeed, but this is a blocker that makes it non-trivial for us to adopt.

Hello，how to use dynamic batch in ddp, can you give an example. due to the dynamic batch size, the number of batches allocated by different ranks is inconsistent. This is what makes different ranks unable to communicate in DDP. Are there other solutions？

[tjruwase]

@ecly, I am also interested in your dynamic batch in ddp. If you can share some client code with us, that would help with DeepSpeed support. Thanks!

[wangleiofficial]

@tjruwase infinibatch may be a good choice for the dynamic batch in ddp. Notice, Dataset with the DistributedSampler may be better than infinibatch for validation set.

[ecly]

Hey @tjruwase and @wangleiofficial

As my original question is starting to be a bit old, I think I perhaps need to retest on a newer version of DeepSpeed to confirm that it's still the case. But nonetheless, I'll share a bit more details below.

The idea is just that we limit batches in the number of total tokens in the batch (to provide a similar amount of signal for learning in each minibatch), rather than the number of examples in the context of training Transformers on text inputs of varying length (in our case for Machine Translation). The code for our MaxTokensBatchSampler that we use as a batch_sampler with the PyTorch DataLoader, is similar in nature to the one used in fairseq: https://github.com/facebookresearch/fairseq/blob/b5a039c292facba9c73f59ff34621ec131d82341/fairseq/data/data_utils.py#L282

We adapt it for DDP with only a tiny bit of code:

class DistributedMaxTokensBatchSampler(DistributedSampler, MaxTokensBatchSampler):

    def __init__(self, dataset: TranslationDataset, batch_max_tokens: int, **kwargs):
        DistributedSampler.__init__(self, dataset)
        MaxTokensBatchSampler.__init__(self, dataset, batch_max_tokens, **kwargs)

    def __iter__(self) -> Iterator[List[int]]:
        iterator = MaxTokensBatchSampler.__iter__(self)
        return itertools.islice(iterator, self.rank, None, self.num_replicas)

    def __len__(self):
        return len(self.batches) // self.num_replicas

This approach can effectively result in batches with different numbers of examples, but where they are similarly sized in terms of the number of tokens. It works out of the box with DDP from our experience.

[HsunGong]

Keep tracking this issue

[bm-synth]

@ecly @aniruddhakal @tjruwase If not too late: I started developing dynamic batch size and the corresponding LR scaling on deepspeed in PR 7104. For the time being you have the Torch and DeepSpeed implementation for FSDP and DDP parallelism.

[tjruwase]

@bm-synth, it is not too late. I believe my colleague @conglongli already responded on #5237. Thanks!

[ShomyLiu]

@bm-synth Thank you for your contribution to this important feature. Since it's not yet merged, I have a question regarding my specific use case.

I'm working with datasets of fixed but different lengths:

1. Dataset A: sequence length = 2048, batch size = 8
2. Dataset B: sequence length = 4096, batch size = 4

Note: The longer sequences use a smaller batch size, but the total token count per batch remains constant across both datasets.

My training setup: In each training step, the dataloader selects only one type of data (either A or B) to train across all GPUs.

Given this simplified scenario, would it be appropriate to set micro_batch_size=8 (i.e., the larger of the two batch sizes) in the DeepSpeed configuration?

Additional context:

• This setup ensures that the total number of tokens processed in each batch is consistent, regardless of which dataset is chosen.
• I'm particularly interested in understanding if this configuration might cause any issues or if there's a more optimal way to handle varying sequence lengths in this context.

Thank you in advance for your guidance on this matter.

[bm-synth]

Hi @ShomyLiu thank you. I'll try to finish the PR asap.

in your case, you dont need to have a "setup" that automatically batches samples. The dataloader_for_variable_batch_size will do it for you (just set max_tokens_per_batch=16384 in the call to batch_by_size. If you already have that, fine. call lr_scheduler_for_variable_batch_size to get your adaptive LR scheduler, and that will update the LR at every step, taking into account the dataloader of your setup.

The micro_batch_size is just the reference batch size to which your LR in the config refers to. As an example: in your case, if you pick lr_scaling_method=linear, micro_batch_size=8 and lr=1e-3 (ie appropriate for dataset A), then when you use dataset B, because your micro_batch_size is 4 (not 8), you LR will be changed internally to lr=2e-3. I.e. the micro_batch_size and lr are just the reference pair, and the LR will change on batch based on the new micro batch size.

I would recommend you to not do the batching yourself (aka the "setup") and call dataloader_for_variable_batch_size instead to generate a dataloader, to avoid issues.

[ShomyLiu]

@bm-synth Thank you for such a detailed explanation and reply. I appreciate your help.
I think I understand the main point for the variable LR.
If I decide to use a constant learning rate and my own custom data loader, it seems that as long as the number of gradient accumulations is fixed, and the variable micro batches are used directly.
From reviewing the source code of Deepspeed, the variable micro batches could affect the only gradient accumulation boundaries in order to initiate gradient reduction.
So if the the number of gradient accumulations is fixed, there are no such issues above?
Is my understanding correct?
Thanks again for your thorough response. Your explanation has been very helpful in clarifying these concepts for me.