feat: add IA3 prompt tuning

[maw501]

This PR adds $(IA)^3$ tuning per Few-Shot Parameter-Efficient Fine-Tuning is Better and Cheaper than In-Context Learning.

The main change is the addition of two new classes, ParallelMLPIA3 and ParallelSelfAttentionIA3. Both of these simply create a new scaling vector then modify the forward pass to perform the (re)scaling. Note that the forward pass is otherwise unchanged except for where we apply $(IA)^3$.

This has been tested using model parallelism (i.e. tensor parallelism in DeepSpeed terminology) on the Stability cluster.

Notes

• There are no hyperparameters for the method.
• We rescale each attention head (the paper is unclear here) independently.

TODO

Pipeline parallelism is currently failing with $(IA)^3$. To replicate the issue:

• pipe-parallel-size : 2 , model-parallel-size : 1, ia3_tuning: True. Note: num_gpus not set. 👎
• pipe-parallel-size : 2 , model-parallel-size : 1, ia3_tuning: False. Note: num_gpus not set. 👍