Issues: Dynamic Batching

[srush]

(Posting a couple issues to get features upstreamed from OpenNMT-py, cc @da03)

For the transformer model, we need some improvements to dynamic batching. In particular the batch_size_fn interface has some issues. Let me give an example.

We need batches of 4096 tokens (including padding).

• We can't really do this with with batch_size_fn because while it lets us count the total number of tokens, it doesn't let us account for padding (max size in the batch). One bad example either causes tons of padding, or a huge batch and an OOM.

Our current terrible hack:

        global max_src_in_batch, max_tgt_in_batch

        def batch_size_fn(new, count, sofar):
            global max_src_in_batch, max_tgt_in_batch
            if count == 1:
                max_src_in_batch = 0
                max_tgt_in_batch = 0
            max_src_in_batch = max(max_src_in_batch,  len(new.src) + 2)
            max_tgt_in_batch = max(max_tgt_in_batch,  len(new.tgt) + 1)
            src_elements = count * max_src_in_batch
            tgt_elements = count * max_tgt_in_batch
            return max(src_elements, tgt_elements)

• Iterator uses this line to buffer data for batching:

for p in batch(data, batch_size * 100, batch_size_fn): https://github.com/pytorch/text/blob/master/torchtext/data/iterator.py#L271

Unfortunately even though there is a 100 here, it doesn't help because if we are counting padding in batch_size_fn then on long example will make every other sentence take a ton of space. Think we need control.

Our current hack (don't use batch_size_fn for buffering):

for p in torchtext.data.batch(data, self.batch_size * 100):

• Minor: Batching use sort for two different purposes. One to find the batches themselves, and the other for the order in which the batch is created. I would like to be able to have a batch_construction_sort to find sentences of the same length and then an batch_sort for each in batch. For example: in MT I would like to sort by a weighted src x tgt len in batch_construction (to minimize padding), but then have the batch itself sorted by src len to make cudnn work.

[jekbradbury]

1. This is what we currently use to perform dynamic batching including vs not including padding in the size:

def dyn_batch_with_padding(new, i, sofar):
    prev_max_len = sofar / (i - 1) if i > 1 else 0
    if args.distillation:
        return max(len(new.src), len(new.trg), len(new.dec), prev_max_len) * i
    else:
        return max(len(new.src), len(new.trg),  prev_max_len) * i

def dyn_batch_without_padding(new, i, sofar):
    if args.distillation:
        return sofar + max(len(new.src), len(new.trg), len(new.dec))
    else:
        return sofar + max(len(new.src), len(new.trg))

It's worse than yours, because it assumes that the maximum of the padded size of the source and the padded size of the target is a good proxy for the average of the two (instead it's a fairly loose upper bound).
I suppose what we should do is provide a generic state-based API:

def batch_size_fn(new, state=None)
    if state is None: state = <initial state>
    state = <update state from new>
    size_proxy = <function of state returning numeric value>
    return size_proxy, state

and let the user decide what that state should represent.

2. Would control over the value that's currently 100 be enough, or is there something more fundamental that needs to change?

3. Agreed. They should certainly default to being the same, but the user should be able to provide two different sort keys.

[srush]

Cool makes sense.

For 1) I think a general state would be nice. You could have it be functional or have it be an object pointer. with a __call__.

For 2,3) it seems like overloading batching to mean both standard mini-batches, and big buffering batches is maybe just not the correct abstraction. Maybe that should always just buffer that next M examples.

[jekbradbury]

Oh of course, yes, that's definitely wrong! Was a holdover from before dynamic batching, when all batch did was buffer the next M.

[srush]

Any update on this? Still have to hack around it in practice.

[ecly]

Curious what the status on this is? It would certainly be nice to have. Also just curious if anyone have any recent success with nice implementations of this.

[patelrajnath]

In the current updated version, I think it is not possible to use batch_size_fn. Can anyone suggest how to do token-based dynamic batching with the current version which uses Pytorch Dataloader. I tried to use sampler but no success.

[ecly]

@patelrajnath for my usecase I've implemented a subclass of torch.utils.data.Sampler that simple generates lists of indices, corresponding to dynamically generated batches, based on the amount of tokens within the batch.

For each call of __iter__ I do the following:

• Shuffle a copy of a List containing Tuple[int, Tuple[int, int]] which maps Dataset indices to the src_len, tgt_len tuples.
• Split this list for entire dataset into sublists, or "pools" (since we will sort afterwards, this is necessary if you want shuffling between epochs).
• Sort each pool by src_len then tgt_len
• For each pool generate batches (which can be of any size) that are efficiently packed (due to sorting), by simply filling them up until the batch_size * max(src_len) + batch * max(tgt_len) is add your desired max batch tokens.

[patelrajnath]

@ecly Thanks! I'll try that.

[parmeet]

Related Issue #1295

[patelrajnath]

@ecly
I implemented this - https://github.com/makeesyai/makeesy-deep-learning/blob/main/seq2seq/torchtext_iter.py

Please suggest if we can improve something.

[ecly]

@patelrajnath I'll instead suggest that you have a look at huggingface who has a similar implementation (which I just found out recently), which you can use to compare or get inspiration:
https://github.com/huggingface/transformers/blob/c2cd02ac625cd0ab64cf42124ad71ce9158fb67c/src/transformers/trainer_pt_utils.py#L501