LoRA Support

[Quentin-Anthony]

Our finetuning focus has been a bit lacking lately. Let's change that, starting with LoRA!

• Run a finetuning baseline with a Pythia model that saturates VRAM, measure its TFLOPs/VRAM usage in wandb and link it here. Choose a small finetuning dataset for this that requires little compute to converge without being trivial.
• Add prototype LoRA support to gpt-neox
• Compare to baseline in step 1 and ensure TFLOP/VRAM changes make sense
• Compare to baseline in step 1 and ensure loss is maintained

As discussed in Discord @mkerin

[mkerin]

Thanks for creating the feature request! As discussed - I'm happy to take this one on :) I expect to start work on it late this week / early next week.