Pipeline to train encoder-decoder LMs

[akutuzov]

We plan to have a bunch of T5-like language models trained on the HPLT v3 datasets:

1. Monolingual models, following more or less the same list of languages as HPLT v2 BERT models
2. A large multilingual model, aiming at providing a modern alternative to mT5

[oepen]

say, we end up training about 50 monolingual T5s to a decent token budget, plus one larger mT5 on maybe 100 to 200 languages and a substantially larger token budget, can we estimate (a) how many gpu hours will be required, and (b) how much effort (time)?

[MariaFjodorowa]

say, we end up training about 50 monolingual T5s to a decent token budget, plus one larger mT5 on maybe 100 to 200 languages and a substantially larger token budget, can we estimate (a) how many gpu hours will be required, and (b) how much effort (time)?

50 monolingual T5s * 28 hours * 16 nodes * 8 GPUs = 50 * 3 584 = 179 200 gpu hours (with context length 512 and training only without evaluation on downstream tasks).
How large is a substantially larger token budget? If keeping the same amount of tokens for each language, it will be 358 400 - 716 800 for 100-200 languages.

Regarding effort, from previous experience with BERTs, it is not less than a month for 50 monolingual ones only, because unexpected small things always occur (such as LUMI internal errors or evaluation code not taking special symbols from some languages into account). I also still have to write evaluation on downstream tasks, we were talking about changing the sampling strategy etc.

Training the multilingual one for 358 400 gpu hours would take 15 hours if occupying the whole LUMI which I guess we cannot do.

We also need cpu-hours on sharding, training tokenizers, tokenizing.

[akutuzov]

On how many tokens you are currently training monolingual T5s?

[MariaFjodorowa]

'steps'(31,250) x 'GPUs'(128) x 'batch_size'(45) x 'seq_len'(512) = 92 160 000 000

[akutuzov]

mT5 was trained on about 1 trillion tokens. Let's may be aim at minimally 50% of this. Something about 200k steps, if keeping the same number of GPUs. However, I guess the vocabulary size will be larger, so it might be the case we will have to reduce the batch size. With the batch size 45, are you already hitting the RAM limits on LUMI GPUs?

[MariaFjodorowa]

The batch size larger than 45 did not work, yes. While the original mT5 not only has vocabulary of 250 000 subwords, but also was trained with the batch size of 1024 and with context length of 1024. They used TPUs for training.

[MariaFjodorowa]

Also for our old friend Aya TPUs were used

[akutuzov]

So, what is our maximum batch size for something like mT5-XL on LUMI without gradient accumulation?

[MariaFjodorowa]

There is currently no gradient accumulation, so 45.

I am a bit sceptical about training the larger model with the current code just reducing the batch size. Both mt5 and Aya used library t5x supporting tensor parallelism for training, thus handling long sequences.

[akutuzov]

Isn't 1024 in mT5 the global batch size?

[MariaFjodorowa]

They say just "batch size" in the paper. I found only this https://github.com/google-research/text-to-text-transfer-transformer/blob/e081111ccd51df425aa7124e8e311ed6d403d767/t5/models/mtf_model.py#L79 where they say "this is the global batch size and not the per-shard batch size.". Here is their computation of batch size per model replica https://github.com/tensorflow/mesh/blob/e6798a2610a2c2f4c4cd236d8214422cb1ecc00a/mesh_tensorflow/transformer/utils.py#L2562 . Also this docs says about batch size per a whole pod https://cloud.google.com/tpu/docs/training-on-tpu-pods#scaling_batch_size_and_train_steps . Should be probably fine with about 32 per GPU for us then

[akutuzov]

Exactly. 128 * 32 = 4096 global batch size, and we can further tune it with gradient accumulation, if needed.