Getting closer

[BenjaminBossan]

It was necessary to flatten the LoRA weights for 3d MoE, as LoRA always expected 2d weights (being nn.Linear).

Therefore, I added new conversion ops: PermuteDims and FlattenDims. I also added a convenience function _block_diag_3d, it could probably be optimized.

Furthermore, I now have separate branches for down_proj and gate_up_proj as they are different enough that treating them in one branch is too confusing.

[BenjaminBossan]

Current error:

MixtralForCausalLM LOAD REPORT from: peft-internal-testing/mixtral-pre-v5-lora
Key                                                              | Status   | Details                                                                                  
-----------------------------------------------------------------+----------+------------------------------------------------------------------------------------------
model.layers.{0, 1}.mlp.experts.lora_A.default.weight            | MISMATCH | Reinit due to size mismatch ckpt: torch.Size([64, 3584]) vs model:torch.Size([64, 1024]) 
model.layers.{0, 1}.mlp.experts.lora_B.default.weight            | MISMATCH | Reinit due to size mismatch ckpt: torch.Size([64, 1024]) vs model:torch.Size([3584, 64]) 
model.layers.{0, 1}.mlp.experts.base_layer.lora_A.default.weight | MISMATCH | Reinit due to size mismatch ckpt: torch.Size([128, 1024]) vs model:torch.Size([64, 7168])
model.layers.{0, 1}.mlp.experts.base_layer.lora_B.default.weight | MISMATCH | Reinit due to size mismatch ckpt: torch.Size([128, 7168]) vs model:torch.Size([1024, 64])

So there isn't much missing to align the shapes.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[github-actions]

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=43327&sha=ba6769