PawKanarek
commented
7 months ago I modified the code a little bit to make some sanity checks.
def train():
gemma2it = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it") # sanity check model
tokenizer = AutoTokenizer.from_pretrained("NousResearch/gemma-2b-it-tokenizer")
model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", torch_dtype=torch.bfloat16)
dataset = load_dataset("pawkanarek/poke_test", split="train")
lora_config = LoraConfig(r=8, target_modules=["k_proj", "v_proj"], task_type="CAUSAL_LM")
fsdp_config = {"fsdp_transformer_layer_cls_to_wrap": ["GemmaDecoderLayer"], "xla": True, "xla_fsdp_v2": True, "xla_fsdp_grad_ckpt": True}
trainer = SFTTrainer(
model=model,
train_dataset=dataset,
tokenizer = tokenizer,
args=TrainingArguments(
per_device_train_batch_size=64,
num_train_epochs=4,
output_dir="output/trained_model",
optim="adafactor",
dataloader_drop_last = True, # Required for SPMD.
fsdp="full_shard",
fsdp_config=fsdp_config,
),
peft_config=lora_config,
max_seq_length=2048,
)
# 1
trainer.train()
print("comparing gemma2it with trainer.model")
compare_weights(gemma2it, trainer.model) # different GemmaForCausalLM:2506172416 params vs SpmdFullyShardedDataParallel:3031123968 params
# 2
merged_model = trainer.model.merge_and_unload()
print("comparing gemma2it with merged_model")
compare_weights(gemma2it, merged_model) # different GemmaForCausalLM:2506172416 params vs GemmaForCausalLM:3030460416 params
# 3
print("saving merged_model")
merged_model.to("cpu")
merged_model.save_pretrained("output/merged_model")
compare_weights(gemma2it, merged_model) # different GemmaForCausalLM:2506172416 params vs GemmaForCausalLM:3030460416 params
# 4
print("comparing loaded merged_model from disk with in-memory merged_model")
loaded_merged_model = AutoModelForCausalLM.from_pretrained("output/merged_model")
compare_weights(merged_model, loaded_merged_model) # different GemmaForCausalLM:3030460416 params vs GemmaForCausalLM:2506172416 params
# 5
print("comparing gemma2it with loaded merged_model from disk")
compare_weights(gemma2it, loaded_merged_model) # models GemmaForCausalLM and GemmaForCausalLM are the sameI added some sanity checks with base, untouched gemma2it model, and some mid-step comparison:
- Check if model after training
trainer.model, differs from the basegemma2it: yes, the are different in number of parameters - that implies that training was succesfull - Check if trained model after merge
merged_model, differs from the basegemma2it: yes, the are different in numer of parameters - that implies that merging was succesfull - Saving merged model and check if model after save
merged_modeldifffers from the basegemma2it: yes, the are different in number of parameters - that implies that saving does nothing to parameters - Loading merged model from the disk
loaded_merged_modeland check if it differs from themerged_modelbefore saving - YES THEY ARE DIFFERENT :( - that implies that there is something wrong with loading the model (or saving) 4.1. This warning popped when loading model from the disk:Some weights of the model checkpoint at output/merged_model were not used when initializing GemmaForCausalLM: ['model.layers.0._orig_module.input_layernorm.weight', (...) 'model.layers.9._orig_module.self_attn.v_proj.weight'] - This IS expected if you are initializing GemmaForCausalLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model). - This IS NOT expected if you are initializing GemmaForCausalLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model). Some weights of GemmaForCausalLM were not initialized from the model checkpoint at output/merged_model and are newly initialized: ['model.layers.0.input_layernorm.weight', (...) 'model.layers.9.self_attn.v_proj.weight'] (...)You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. - Check if merged model from disk
loaded_merged_modeldiffers from the basegemma2it: no, they are the same... - that implies that all my training was worthless...
Looks like there is something fishy with my code when saving / loading model from the disk... I'll update if i notice what's wrong. I will check why my weights are saved to something called _orig_module.
zorrofox
amyeroberts
shub-kris
moficodes
alanwaketan
System Info
torch.version='2.3.0.dev20240307' torch_xla.version='2.3.0+git46e2230' peft.version='0.9.0' trl.version='0.7.12.dev0' Python 3.10.13
Who can help?
@ArthurZucker , @younesbelkada, @muellerzr, @pacman100
Information
Tasks
examplesfolder (such as GLUE/SQuAD, ...)Reproduction
Hello, I have a problem with training the
gemma-2b-itmodel on Google TPU v3-8. My goal is to train it with the peft lora adapter, and then save it as a standalone model.For merging base model with lora adapter I was following the guide: https://huggingface.co/docs/trl/main/en/use_model Training code is based on this blog post: https://huggingface.co/blog/gemma-peft
The problem is that the training takes a while (for 300k rows in a data loader it might take even 8 hours) but after training the model seems… untrained. The interference output looks almost identical to the output of the base model.
Furthermore, when I check for the weights of the trained and original models then they appear to be identical.
I also consistently encounter the following error message, while loading saved model:
Below is the minimal working code that trains and saves the model.
And this is the output
I'm stuck so, I'm asking for help. I tried many combinations of the
PeftModel.merge_and_unload(),saving_pretrained(), andtrainer.save_model()and nothing seems to work. Every idea to push this issue forward will be appreciated. Thanks.Expected behavior
Training trains the model.