Unloading multiple loras: norms do not return to their original values

[christopher5106]

When unloading from multiple loras on flux pipeline, I believe that the norm layers are not restored here.

Shouldn't we have:

        if len(transformer_norm_state_dict) > 0:
            original_norm_layers_state_dict = self._load_norm_into_transformer(
                transformer_norm_state_dict,
                transformer=transformer,
                discard_original_layers=False,
            )
            if not hasattr(transformer, "_transformer_norm_layers"):
                 transformer._transformer_norm_layers = original_norm_layers_state_dict

[sayakpaul]

Should it not already take care of it?

diffusers/src/diffusers/loaders/lora_pipeline.py

Line 1894 in 464374f

transformer.load_state_dict(transformer._transformer_norm_layers, strict=False)

What am I missing?

Additionally, the following test does ensure its effectivity:

diffusers/tests/lora/test_lora_layers_flux.py

Line 661 in 28f48f4

def test_lora_unload_with_parameter_expanded_shapes(self):

[christopher5106]

Ah, is it possible to call load_lora_weights() multiple times on a pipeline to load multiple weights ? does it unload in between to restore original weights ?

[sayakpaul]

If you don’t call unload_lora_weights() it won’t be called automatically.

[christopher5106]

so in that case of multiple calls to load_lora_weights(), the attribute _transformer_norm_layers become overwritten by the norms of the previously loaded lora ?

[sayakpaul]

If you're loading Control LoRA and want to keep it with others, the norm layer values that came with the Control LoRA will remain like so. Otherwise, the effectiveness of Control LoRA won't fully be there.

Or am I misinterpreting the core use case here?

[christopher5106]

From what I understand, you make the assumption that only Control Lora have trained norms, right ?

From what i see in my code is:

for lora, adapter in loras:
   pipe.load_lora_weights(local_weights_cache, adapter_name=adapter_name)

that means if two loras come with trained norm layers, we loose original weights. I believe we should make it generic because we might forget this assumption.

[sayakpaul]

From what I understand, you make the assumption that only Control Lora have trained norms, right ?

Well, Control LoRA is a bit of a special case in that its state dicts have the exact same norm params as their non-LoRA variants. More specifically, these norm layer params differ from the base Flux.1 Dev model and there were taken from the non-LoRA control variants of Flux (so for the Depth Control LoRA, that would be https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev and for the Canny Control LoRA, that would be https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev/). The norm params are not LoRA params.

But on the other hand, it's totally possible to also target the norm layers for applying LoRA (which is something we see often in the community).

So, I don't think there's a need to change anything here.

[christopher5106]

but my wonder is about the following:

I load one lora with norm layers, let's say lora1_norms:

• it saves the original layer weights (let's say original_norms) in _transformer_norm_layers

I load a second lora with norm layers, let's s say lora2_norms:

• it saves the overwritten layer weights in _transformer_norm_layers

so now in _transformer_norm_layers I have lora1_norms

when I apply unload_lora_weights(), does it revert to original_norms or to lora1_norms ? what did I miss ?

[sayakpaul]

when I apply unload_lora_weights(), does it revert to original_norms or to lora1_norms ? what did I miss ?

It unloads all the LoRA overwritten params including lora1 and lora2.

More notes:
https://huggingface.co/docs/diffusers/main/en/api/pipelines/flux#note-about-unloadloraweights-when-using-flux-loras

[christopher5106]

Sorry I don't see what I miss

If I call two times load_lora_weights with 2 different loras lora1 and lora2 that have norm layers, that means to me that two times it's calling _load_norm_into_transformer() that returns the overwritten layers. But to me, where it takes these values is from the current transformer state dict

diffusers/src/diffusers/loaders/lora_pipeline.py

Line 1669 in 69f919d

overwritten_layers_state_dict[key] = transformer_state_dict[key].clone()

and if these norms have been overwritten by lora1 already, when you call unload_loraweights, it will restore the lora1 norm with

diffusers/src/diffusers/loaders/lora_pipeline.py

Line 1894 in 464374f

transformer.load_state_dict(transformer._transformer_norm_layers, strict=False)

but not the original norms in the model.

[sayakpaul]

Yeah you're right! But it's also not very common to have structures for norm layer params. We can prioritize that as in when such a LoRA comes in. For now, we can add a comment to make ourselves aware of it.

[christopher5106]

As you want ;)
But you see my point now :-)
I believe it's a bit wider problem than norm layers since load_state_dict() and state_dict() can load whatever tensors for whatever layers in our models

[sayakpaul]

But it's also weird for LoRA state dicts to have arbitrary non-LoRA params. So, that is there.

[christopher5106]

bfl could create a new trend that loras with bias and norms become a way to have better loras, no?