LORA + FSDP: ValueError: FlatParameter requires uniform requires_grad

[minionKP]

Environment

To reproduce

Steps to reproduce the behavior:

1. add lora configs to llm-foundry/scripts/train/finetune_example/mpt-7b-arc-easy--gpu.yaml

lora:
  args:
    r: 16
    lora_alpha: 32
    lora_dropout: 0.05
    target_modules: ['Wqkv']

2. run composer train.py finetune_example/mpt-7b-arc-easy--gpu.yaml

Expected behavior

FSDP + LORA fine-tuning

Config file

No other changes to example config file

[codestar12]

Hi @minionKP. We do not yet have LoRA working with FSDP. To use LoRA use DDP. We expect to add support eventually, but do not have a solid timeline.

[minionKP]

Want to do LORA on MPT-30B model. From another issue #413 looks like QLORA or doing prepare_model_for_int8_training from PEFT does not work. Thanks for communicating the priority.

[floschne]

I think this issues not only arises with LoRA but with any mixed setup of frozen and unfrozen parameters. Is there any solution in sight? I know one could write a custom auto-wrap policy when using lightning or pure PyTorch but that seems to be complicated especially because it creates un-uniformly sized FSDPUnits :-/

[dakinggg]

Hi @floschne, LoRA + FSDP support is almost merged! You can follow along here if you like: #886. That PR does currently work if you install composer from source.