Train Embedding Layer with LoRA

[RADxEEn]

Hey, so I know there was a similar question like it. But I am facing a problem
So, I not only want to update the layer, but also update the embedding layer. So for it I use
modules_to_save=["embed_tokens"]

but after training there should be weight change between base model embed_tokens_weight and modules_to_save Embedding weight. But when I check on their mean this is 0, that means it didn't update my embedding layer weight but, it updated the layers weight. So, I guess somehow my embedding layer is not connecting or participating in training.
merged_model_f_n.base_model.model.model.lm_head.original_moduleSS.weight-merged_model_f_n.base_model.model.model.lm_head.modules_to_save.default.weight

tensor([[0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        ...,
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.]], device='cuda:0', dtype=torch.float16,
       grad_fn=<SubBackward0>)

PeftModel(
  (base_model): LoraModel(
    (model): PeftModelForCausalLM(
      (base_model): LoraModel(
        (model): MistralForCausalLM(
          (model): MistralModel(
            (embed_tokens): ModulesToSaveWrapper(
              (original_module): Embedding(67840, 4096)
              (modules_to_save): ModuleDict(
                (default): Embedding(67840, 4096)
              )
            )

Will be happy to accept any suggestion. Thanks

[BenjaminBossan]

This is quite strange and I could not reproduce the issue. Your suggestion to add modules_to_save=["embed_tokens"] is correct. However, this should create a copy of the original weights, so they should not be all zeros. Could you please provide a complete code example to reproduce the issue?

[RADxEEn]

So I have model structure suppose one is base_model and the other one is merged_model which is being merged with base and adapter model after training.

base_model.model.model.lm_head.original_moduleSS.weight-.base_model.model.model.lm_head.modules_to_save.default.weight

which should be give some difference between the weight. But as giving same weight, so after taking difference, it is giving 0.

[BenjaminBossan]

Sorry, I still don't quite understand how you calculated that, maybe you could provide more code?

However, the layers weight has changed

So is it possible that the update worked correctly?

[RADxEEn]

Suppose the structure is like below for both code (base model and merged_model_after_train).
So for base model we can get the embedding weight by using

MistralForCausalLM(
  (model): MistralModel(
    (embed_tokens): Embedding(67840, 4096)
    (layers): ModuleList(
      (0-31): 32 x MistralDecoderLayer(
        (self_attn): MistralAttention(
          (q_proj): Linear4bit(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear4bit(in_features=4096, out_features=1024, bias=False)
          (v_proj): Linear4bit(in_features=4096, out_features=1024, bias=False)
          (o_proj): Linear4bit(in_features=4096, out_features=4096, bias=False)
          (rotary_emb): MistralRotaryEmbedding()
        )
        (mlp): MistralMLP(
          (gate_proj): Linear4bit(in_features=4096, out_features=14336, bias=False)
          (up_proj): Linear4bit(in_features=4096, out_features=14336, bias=False)
          (down_proj): Linear4bit(in_features=14336, out_features=4096, bias=False)
          (act_fn): SiLU()
        )
        (input_layernorm): MistralRMSNorm()
        (post_attention_layernorm): MistralRMSNorm()
      )
    )
    (norm): MistralRMSNorm()
  )
  (lm_head): Linear(in_features=4096, out_features=67840, bias=False)
)

model.embed_tokens.weight which will generate the tensor of the weight for embedding module.
same for the merged model.

Then if I take the difference between these two tensors, it should not be 0, as it should be updating, like the self_attention layer module.

If I calculate the difference for self_attention layer between two model, i have found that the difference is giving me non zero tensor, which is not the case for embedding layer weight, even though I use modules_to_save.

So, what I am thinking is that, I missing something or is the merging the adapter with the base model is not doing properly?

Here I am giving the merging code

from peft import PeftModel
adapter_path="/content/bn_adapter"
merged_model_f_n=PeftModel.from_pretrained(base_model,adapter_path)
merged_model_f=merged_model_f_n.merge_and_unload()

[BenjaminBossan]

So one thing I see is that you're using 4bit quantization. Is the quantization also applied to the embedding layer? Since this should make it so that the embedding layer cannot be trained.

[RADxEEn]

Yes the quantization also applied in the embedding layer. Anyway, the issue
has been resolved. The problem is with the colab. As they have a tendency
to optimize the vram, although i am calling with different variable,somehow
they use the cache to store the data. For this, if I call the models, I
don’t know but it automatically updates the base model similar to the
merged model. For this, i am getting zero. Now in kaggle, the problem has
been solved. Thank you.