fix: utilities to post process checkpoint for LoRA

[Ssukriti]

Description of the change

utility function to post process checkpoint after LoRA tuning to convert to format required by vLLM.
This will need to be called at end of LoRA tuning to allow inferencing on LoRA, for models for which we have added new tokens.

Since it is fast enough to load adapters.safetensors , added it as a post-processing function.

This PR adds a script that can be called after tuning to do the processing.

Related issue number

https://github.ibm.com/ai-foundation/watson-fm-stack-tracker/issues/1210
Details on vLLM issue: vllm-project/vllm#2816 (comment)

Context: Embedding vectors for new tokens need to be placed in new_embeddings.safetensors and lm_head.weight should be deleted from adapter_model.safetensors as per vllm-project/vllm#2816 (comment)

How to verify the PR

1. Function called on llama model which was LoRA tuned - post-processed and being tested on vLLM
2. Verified PR with every unique model architecture we support

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

[kmehant]

Couple of comments, thanks

[github-actions]

Thanks for making a pull request! 😃
One of the maintainers will review and advise on the next steps.

[aluu317]

I added a commit to add unit test in test_merge_model_utils.py. Steps to run it (since Github Action skips it due to not having cuda);

• Start a sleep pod with volume mount to our fmaas-integration-tests COS bucket.
• Change the value in test file just for this run:

DUMMY_TUNED_LLAMA_WITH_ADDED_TOKENS = "/testing/tuning/output/llama_dummy_LoRA/checkpoint-35"

• Copy test file into pod:

oc cp tests/utils/test_merge_model_utils.py angel-sleep-pod-sft-trainer:/home/tuning/.local/lib/python3.11/site-packages/tuning/utils/test_merge_model_utils.py

• In sleep pod, pip install pytest. Then run:

[tuning@angel-sleep-pod-sft-trainer app]$ cd /home/tuning/.local/lib/python3.11/site-packages/tuning/utils
[tuning@angel-sleep-pod-sft-trainer utils]$ pytest test_merge_model_utils.py 
======================================= test session starts ========================================
platform linux -- Python 3.11.7, pytest-8.3.3, pluggy-1.5.0
rootdir: /home/tuning/.local/lib/python3.11/site-packages/tuning/utils
collected 1 item                                                                                   

test_merge_model_utils.py .                                                                  [100%]

======================================== 1 passed in 3.50s =========================================

[willmj]

Note these tests are failing because they are running old unit tets (see changes in previous commit)
Test fail (only 10 args):

What it should be running (11 args):

[Ssukriti]

this was just done to ensure we add and resize embeddings in same function

[Ssukriti]

Returning additional metadata from train() containing information of newly added tokens. @kmehant @ashokponkumar let me know if this change is ok with you - it might be disruptive if using SDK and if you were relying on 1 return value. We can make this a major release with the change to clarify API changed.

Alternatively, I can write the file artifacts I want inside 'output_dir' and later can copy it to save_model_dir as well, if return value is an issue. Then I can avoid returning the dict. I just thought this return value could be extendable to any other metadata we need in future

[ashokponkumar]

If we remove custom addition of tokens from fms-hf-tuning, we might not need to return this added tokens metadata.

[Ssukriti]

These files are all just dummy LoRA artifacts needed for unit tests

[anhuong]

Great work all! Few questions and changes, thanks for adding docs! I know you talked about adding some unit tests for the tokenizer_data_utils which would be good as well

[anhuong]

Just to verify, new_embedding_size could be:

• the new size of the embeddings after new tokens are added
• new size of embeddings after resized to multiple_of
• or the same size if none of the above occurs

[Ssukriti]

its size of embeddings at end of tuning , whichever case it is. Will be helpful to decipher weight names based on size

[anhuong]

Should we check that num_new_tokens > 0 otherwise don't need to write this file? This is for the case where the model already has all the tokens it needs so no new tokens or embedding resize is happening.

[Ssukriti]

no , writing it in all cases to differentiate between "tuning was not done on latest code / no tokens were added". The file will always exist, it will contain value 0 if no tokens added

[Ssukriti]

added a check in post-process script to exit early if num_added_tokens == 0 though https://github.com/foundation-model-stack/fms-hf-tuning/pull/338/files#diff-443a0a5c8690e1a9251b2805f8e5850c97381d40b47cd90c5716eeaca84706d9R61

[Ssukriti]

@anhuong @willmj I have added sufficient unit tests pertaining to this change. Lets wait on @ashokponkumar to check the PR as well.

[ashokponkumar]

Few clarifications..

[ashokponkumar]

Sorry for my ignorance.

But just wondering, why not just remove this logic that is specific to llama tokernizer or GPT2 tokenizer out of fms-hf-tuning, we just assume the base model just has it? This will remove the need for all these additional tokens. For example, I am not sure why unk token is being added. My undertanding was BPE based tokenizers never end up creating unk tokens.

Or we can make these tokens adding a externally flag, that allows adding custom tokens to any tokernizer. This would make fms-hf-tuning more generic.

[Ssukriti]

I feel there is often some confusion here, this piece of code was taken from openinstruct https://github.com/allenai/open-instruct/blob/main/open_instruct/finetune.py#L635 , if you see they have mentioned that though the tokens are already in model, they are not added as special_tokens. This can have some implications in quality as models treat special tokens different from regular tokens - do not split it etc. Hence this change is just making those tokens special tokens instead and not adding any new tokens.

Num_added_tokens will just eb the pad token

[Ssukriti]

pad token is always needed to be added for any model architecture - that is a open source requirement by SFT Trainer, without which training will not proceed. Hence we always add pad token for all architectures here https://github.com/foundation-model-stack/fms-hf-tuning/blob/main/tuning/sft_trainer.py#L278 (except granite which already has pad token)

Hence minimum 1 token is always added and hence post processing is needed

[ashokponkumar]

If we remove custom addition of tokens from fms-hf-tuning, we might not need to return this added tokens metadata.

[Ssukriti]

@ashokponkumar as explained on slack, we cannot remove addition of pad token as it is a rwquirement and known in open source. Without which we will run into this error https://stackoverflow.com/questions/70544129/transformers-asking-to-pad-but-the-tokenizer-does-not-have-a-padding-token

Most models in open source do not have pad token - new llama3.1, llama3, llama, allam, mixtral, mistral . Hence we have to add minimal 1 token for all these architectures, which we do in generic manner 'if pad token is None, set it'

This PR is thus doing the post-processing needed to handle addition of any token for LoRA inferencing on vLLM. Without this PR, LoRA inference on vLLM does not work for any of above architectures.

Even if we remove other tokens, like unk etc - which can be done in following PR and issue , the change is still needed for pad token. Hence better to keep the change generic and return number of added tokens from code

[anhuong]

LGTM 🚀

changes addressed and reviewer not available