fix: remove lm_head for granite with llama arch models

[Ssukriti]

Description of the change

• Before copying the last checkpoint created to the users's output dir, we load the checkpoint, test to see if it has llama arch and that the lm_head weight matches the embed_weight, and if so delete the lm_head.weight.
  • Loads model with AutoModelForCausalLM.from_pretrained() if is a fine tuned model, otherwise loads base model and adapter with PeftModel.from_pretrained().
  • After deleting lm_head.weight, saves the model to output_dir with model.save_pretrained()

Related issue number

We found that for granite models with llama architecture (such as granite-3b-code) has lm_head weights tied to embed weights and when tuning with accelerate and FSDP for multi-GPU, the lm_head weight was created as a duplicate of the embed weight. Due to this unexpected weight being created, vLLM is unable to load the model. To fix this, we are deleting the lm_head weight.

How to verify the PR

Ran updated accelerate_launch.py script in cluster with below configs:

granite-3b-code-base, fine tuning, multi-GPU

Successfully fine tuned model and removed lm_head. The tuned model was able to be loaded on vLLM (where verified previous lm_head error). Verified that lm_head is removed from checkpoint's model.safetensors.index.json.

tail tuning logs:

{'train_runtime': 585.5663, 'train_samples_per_second': 7.395, 'train_steps_per_second': 0.939, 'train_loss': 0.4583773977106268, 'epoch': 10.0}
100%|████████████████████████████████████████████████████████████████████████████████████| 550/550 [09:45<00:00,  1.06s/it]
You are attempting to use Flash Attention 2.0 with a model not initialized on GPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████| 2/2 [00:16<00:00,  8.17s/it]
INFO:root:Removing lm_head from checkpoint

granite-3b-code, LoRA tuning, multi-GPU

Successfully LoRA tuned model and removed lm_head. The tuned model was able to be loaded on vLLM (where verified previous lm_head error).

tail tuning logs:

{'train_runtime': 281.5458, 'train_samples_per_second': 15.379, 'train_steps_per_second': 1.954, 'train_loss': 0.9379253248734908, 'epoch': 10.0}
100%|████████████████████████████████████████████████████████████████████████████████████| 550/550 [04:41<00:00,  1.95it/s]
DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): huggingface.co:443
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "HEAD /ibm-granite/granite-3b-code-base/resolve/main/config.json HTTP/11" 200 0
You are attempting to use Flash Attention 2.0 with a model not initialized on GPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████| 2/2 [00:12<00:00,  6.02s/it]
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "HEAD /ibm-granite/granite-3b-code-base/resolve/main/generation_config.json HTTP/11" 200 0
INFO:root:Removing lm_head from checkpoint

granite-3b-code-base, prompt tuning, single-GPU

Successfully prompt tuned and as we know single GPU does not result in the lm_head issue and so it does not delete lm_head tuning runs.

llama-13b-base, fine tuned, multi-GPU

Successfully fine tuned and as expected, does not result in lm_head being deleted. Tuning logs show no message that lm_head is deleted.

Was the PR tested

• I have added >=1 unit test(s) for every new method I have added.
• I have ensured all unit tests pass

[Ssukriti]

maybe just deleting the lm_head.weight should suffice?

[anhuong]

I saved both the model and tokenizer because I think the tokenizer is needed for vLLM to load the model and to match the previous checkpoint output. The LoRA tuned output looks the same but fine tuned output looks a little different:

new output:

$ ls /data/anhuong/tuning_output/granite-3b-code-ft-test-remove-lm-head
config.json		model-00001-of-00002.safetensors  special_tokens_map.json  training_logs.jsonl
generation_config.json	model-00002-of-00002.safetensors  tokenizer.json	   vocab.json
merges.txt		model.safetensors.index.json	  tokenizer_config.json

previous output:

$ ls /data/anhuong/tuning_output/granite-3b-code-test-accelerate-orig_params-transformers_v4.42
config.json			  model.safetensors.index.json	scheduler.pt		 training_args.bin
generation_config.json		  optimizer.bin			special_tokens_map.json  training_logs.jsonl
merges.txt			  pytorch_model_fsdp.bin	tokenizer.json		 vocab.json
model-00001-of-00002.safetensors  rng_state_0.pth		tokenizer_config.json
model-00002-of-00002.safetensors  rng_state_1.pth		trainer_state.json

the model was still able to be loaded by vLLM but wanted to make note of this.

[anhuong]

In addition, I was able to check for fine tuned models that lm_head.weight was deleted by checking file model.safetensors.index.json, but is there an easy way to check this for LoRA tuned adapters?

[anhuong]

Finally, when I load the fine tuned model up in python, you do see lm_head.weight still even though it's deleted in model.safetensors.index.json.

>>> use_flash_attn = os.getenv("use_flash_attn", True)
>>> ft_model = "/data/anhuong/tuning_output/granite-3b-code-ft-test-remove-lm-head"
>>> ft_model = AutoModelForCausalLM.from_pretrained(ft_model, attn_implementation="flash_attention_2" if use_flash_attn else None, torch_dtype=bfloat16 if use_flash_attn else None)
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████| 2/2 [00:33<00:00, 16.94s/it]
>>> ft_model.lm_head
Linear(in_features=2560, out_features=49152, bias=False)
>>> ft_model.lm_head.weight
Parameter containing:
tensor([[-0.0022,  0.0032,  0.0302,  ..., -0.0132, -0.0620,  0.0070],
        [ 0.0152,  0.0287, -0.0002,  ..., -0.0028,  0.0100,  0.0159],
        [ 0.0132, -0.0197,  0.0447,  ...,  0.0247,  0.0732, -0.0236],
        ...,
        [ 0.0067, -0.0435, -0.0234,  ..., -0.0294,  0.0168,  0.0226],
        [ 0.0129, -0.0044,  0.0471,  ...,  0.0024,  0.0066, -0.0032],
        [ 0.0057, -0.0280, -0.0093,  ...,  0.0039,  0.0347,  0.0199]],
       dtype=torch.bfloat16, requires_grad=True)
>>> ft_model.lm_head.weight.untyped_storage().data_ptr() == ft_model.model.embed_tokens.weight.untyped_storage().data_ptr()True
>>> ft_model.lm_head.weight.untyped_storage().data_ptr() == base_model.lm_head.weight.untyped_storage().data_ptr()
False
>>> ft_model.lm_head.weight.untyped_storage().data_ptr() == base_model.model.embed_tokens.weight.untyped_storage().data_ptr()
False

[anhuong]

I had assumed that model.model.model.embed_tokens.weight.untyped_storage().data_ptr() was the same as referencing

params_dict = dict(model.named_parameters())
params_dict.get("base_model.model.model.embed_tokens.weight")

Let me know if this assumption is incorrect. I moved away from params_dict in order to not have to differentiate between base_model and model

[alex-jw-brooks]

Can you add a docstring to this function, just explaining that the adapter config is for peft models since this might be confusing to people that have only used transformers models?

[alex-jw-brooks]

It would be better to use os.path.join here since we use it in most places and it's easier to port to different platforms :)

Suggested change

	adapter_config_path = f"{last_checkpoint_path}/adapter_config.json"
	adapter_config_path = os.path.join(last_checkpoint_path, "adapter_config.json")

[alex-jw-brooks]

is there a reason that this needs to happen again? I think this is called by the launched sft trainer command, so shouldn't the resizing already be handled?

[anhuong]

Unit tests for LoRA and PT were failing when trying to load the checkpoint with error:

RuntimeError: Error(s) in loading state_dict for LlamaForCausalLM:
	size mismatch for model.embed_tokens.weight: copying a param with shape torch.Size([32008, 64]) from checkpoint, the shape in current model is torch.Size([32000, 64]).

so then I had to resize the tokens to be able to load the checkpoint.

[alex-jw-brooks]

Iiiiinteresting. I guess the resize information must not be preserved in the config and reloading with the base class resets it to the base size or something 🤔

[anhuong]

I also verified that I get this same error when manually running tuning on tiny llama model with same params as unit test and then trying to load the checkpoint.

[Ssukriti]

huh? is this happening on current main due to some other changes in dependencies or something?

why is this change related to the lm_head deletion here?

[Ssukriti]

if something changed, we should not just patch it by doing such things for final checkpoint only. Though product uses final checkpoint only for now , rest of research can use any checkpoint . so our load and inference scripts should work with any checkpoint and we have to understand what changed

[anhuong]

I didn't think this was a patch, I thought this was always needed to load a model we tuned such as what we run in run_inference. Since we update the tokenizer in our sft_trainer then when we load the tuned checkpoint again we need to update the tokenizer as well

[anhuong]

I didn't realize this was a new change added in a few weeks ago - https://github.com/foundation-model-stack/fms-hf-tuning/pull/227/files -- that resizes the embeddings to be a multiple of 8 with new tokens instead of the size of the tokenizer which required this addition....hmmm