Fix Large-BS Dataloader Bug

[Quentin-Anthony]

Currently, if train_iters * seqlen * gradient_accumulation_steps * micro_batch_size * world_size > 2147483647 our dataloader's sample_idx overflows leading to the cryptic error:

  File "torch/utils/data/_utils/collate.py", line 141, in default_collate
    return torch.stack(batch, 0, out=out)
RuntimeError: stack expects each tensor to be equal size, but got [2049] at entry 0 and [198705720] at entry 16

Which will happen once the torch dataloader reaches the overflowed sample_idx. Simply storing to np.int64 instead of np.int32 will do the job, but will waste memory and disk. I added a simple switch between the default int32 and new int64 dataset builders and tested it works.

[Quentin-Anthony]

I went the easy route and just created two separate dataset building functions (build_sample_idx_int32 and build_sample_idx_int64), then switch between them depending on whether the number of samples will overflow with int32.

Now we won't waste memory with int64 for the majority of runs that won't overflow with int32. Tested with both overflowing and non-overflowing cases without issue.