A few questions on training about palm-rlhf-pytorch

[conceptofmind]

Hi, I've been planning to train this model, I have a tpu pod(v3-128) through trc, which should equate to ~ 5 tb of ram and 2 tb of vram, I had a few questions about how to begin training the model.

1. What would be the appropriate dataset to train on? I was currently considering using the pile for pre training, but gathering human feedback for rlhf still seems like a challenge.
2. How large of a model would you recommend?
3. I saw you mentioned flash attention, are there any drawbacks to using it, because it seems to be practically the best attention

Thanks for all of your implementations, they have been really helpful to learn from

I am currently doing a distributed training run and will be open-sourcing all of the weights for a PaLM-rlhf-pytorch model. I will open a PR when training finishes.

To answer these questions:

1. The Pile, or its deduplicated version, is currently the best bet for a large open-source dataset. The Pile with the NEOX tokenizer is over ~300B tokens. I am releasing a massive instruction FLAN dataset soon with the help of the original Google authors. You could combine both of these datasets. There are a few different human feedback datasets such as the hh dataset by Anthropic on the hub.
2. The size of the model completely depends on the amount of computing you have and how many tokens you want to train for. You would have to see how long it would take to train a model of x size with x number of batches on your TPU pods. What sequence length you plan on using as well as other factors such as TFLOPs.
3. You can't use Flash Attention with TPUs.

I will update as training progresses.

Best,

Enrico

[TheRealAakash]

Hi Enrico,
Thanks for your response, on the note of flash attention not being possible on tpus, does this imply that tpu context size/efficiency will be substantially behind gpus for the foreseeable future? Or is there an alternative to get similar improvements with tpus? Memory management and partitioning has been what I struggle with the most in trying to train on tpus with Jax.

Thanks,
Aakash

[conceptofmind]

Hi Enrico, Thanks for your response, on the note of flash attention not being possible on tpus, does this imply that tpu context size/efficiency will be substantially behind gpus for the foreseeable future? Or is there an alternative to get similar improvements with tpus? Memory management and partitioning has been what I struggle with the most in trying to train on tpus with Jax.

Thanks, Aakash

TPUs are already pretty highly efficient at scale. You can benchmark Flash Attention vs Jax if you are interested in the exact speed tests.