Task-conditioned LoRA + MoE

[ryspark]

This PR adds two new pieces of functionality: LoRA and MoE adapters conditioned on task-specific embeddings. These are meant to be learned during finetuning.

LoRA adapters: drop-in replacement for the attention output projection weights, conditioned on downstream tasks. This is meant to be used during finetuning. From this paper it seems like finetuning only these projection weights gets you most of the way there compared to full finetuning.

• Given embeddings E for each downstream task, the TaskLoRALinear layer computes a LoRA update (ie two matrices A: D x r, B: r x D) whose product is added to the original projection weights. Both A and B are computed via an MLP on top of E. This MLP is shared across tasks but different per-layer, and is directly learned during finetuning.

MoE adapters: instead of computing FFN(x) in each Transformer block, adds a soft MoE adapter so that the pre-LayerScale output is Linear(FFN(x) + MoE(x)).

• Optionally MoE adapters compute expert combine weights (i.e., deciding which experts to use per token) by conditioning on batch-level task embeddings instead of on token-level embeddings.

• The actual MoE implementation is mostly from this reference implementation and can be found in helios.nn.moe.

To accommodate these changes, there are a few extra arguments added to EncoderConfig, Encoder, FlexiHeliosBase, etc. all the way down to the base helios.nn.attention.Attention layers. Additionally, there is a new argument task_emb added to the forward pass of Encoder. I considered subclassing Encoder (i.e. something like EncoderWithTaskEmbeds) but decided it was simpler to just add a few extra arguments directly.