When I try the below I get the error ValueError: some parameters appear in more than one parameter group. However, inspecting the model it is not clear to me what is the overlapping module.
The only possiblity as to why it might think so may be because lm_head and transformer.wte have parameters named weight. I'm wondering if this name is what is causing this error.
I am doing this so that I can have the lower layers "moving slowly" compared to the upper layers. Happy to hear if there is an alternative way to do these discriminative learning rates where I don't have overlapping parameters (if any).
import torch
from transformers import AutoModelForCausalLM
language_model = AutoModelForCausalLM.from_pretrained("gpt2")
FREEZE_LAYERS = 2
caption_params = [
{"params": language_model.lm_head.parameters() , "lr": 1e-4},
{"params": language_model.transformer.ln_f.parameters() , "lr": 1e-4},
{"params": language_model.transformer.h[FREEZE_LAYERS:].parameters() , "lr": 5e-5},
{"params": language_model.transformer.wte.parameters() , "lr": 1e-5},
]
optimizer = torch.optim.Adam(caption_params)