Issue merging a Lora model to a SANA transformer

[frutiemax92]

System Info

peft=0.14.0

Who can help?

@BenjaminBossan @sayakpaul

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder
• My own task or dataset (give details below)

Reproduction

from diffusers import SanaPipeline, SanaPAGPipeline, SanaTransformer2DModel
from peft import PeftModel

transformer = SanaTransformer2DModel.from_pretrained("frutiemax/twistedreality-sana-1600m-1024px")
print(transformer)
peft_model = PeftModel.from_pretrained(transformer, '0')
model = peft_model.merge_and_unload()

Expected behavior

I've trained a Lora model with PEFT on a SANA checkpoint. I can train and inference using the PEFT model. However, when I try to merge the Lora to the base checkpoint, I encounter a shape mismatch. I've attached the Lora model with a rank 4.

0.zip

[sayakpaul]

Please make the code fully reproducible.

[frutiemax92]

Updated the code with missing imports.

[BenjaminBossan]

Could you please show how you initialized the model (training code not necessary) and how you saved it? Also, do you know if the frutiemax/twistedreality-sana-1600m-1024px corresponds to the official Sana model? What is different?

[frutiemax92]

I've investigated this a bit more and found out only the target module 'conv_depth' makes it crash, PEFT with the other modules work fine.

from diffusers import SanaTransformer2DModel
from peft import LoraConfig, get_peft_model

transformer = SanaTransformer2DModel.from_pretrained("Efficient-Large-Model/Sana_1600M_1024px_diffusers", subfolder='transformer')
lora_config = LoraConfig(r=4, target_modules=['conv_depth'], lora_alpha=4)
model = get_peft_model(transformer, lora_config)
model = model.merge_and_unload()
model.save_pretrained('merged_model')

[frutiemax92]

I also want to add that calling the forward function with lokr and loha also crash on the conv_depth module, but not with the regular lora algorithm.

[BenjaminBossan]

Thanks for the snippet. The LoRA adapter does not have the right shape because we're not honoring the groups argument of this Conv2d layer. Usually, it's 1, so it doesn't matter, but here it's 11200. This issue was already reported in #2153 but still awaiting a PR.

I also want to add that calling the forward function with lokr and loha also crash on the conv_depth module, but not with the regular lora algorithm.

Most likely, the same issue is the cause here. I'm not sure why forward works for LoRA but I would not rely on the result being correct.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

[BenjaminBossan]

not stale

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

[BenjaminBossan]

There was an attempt to add groups support in #2403. @frutiemax92 it might help us if you could share how you trained the LoRA.

[frutiemax92]

Only the code I've provided is needed to replicate the problem.

[BenjaminBossan]

It's not about reproducing the error but finding a solution that produces the same output as is expected from training. Did you basically create a PEFT model with

transformer = SanaTransformer2DModel.from_pretrained("Efficient-Large-Model/Sana_1600M_1024px_diffusers", subfolder='transformer')
lora_config = LoraConfig(r=4, target_modules=['conv_depth'], lora_alpha=4)
model = get_peft_model(transformer, lora_config)

and then train it with a vanilla PyTorch training loop? Did you use one of the many training frameworks out there?

[frutiemax92]

I've used my own training code:
https://github.com/frutiemax92/YAT

It is a vanilla pytorch training loop with WebDatasets, PEFT and EMA.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.