fix: cross entropy for transformers>4.45

[anhuong]

• Checks for transformers version and creates new custom loss function for llama, granite, mistral, and mixtral models.
• Adds shard_checkpoint function from transformers as it is missing in later versions

Tested and saw cross-entropy switching based on transformers version correctly. Benchmark singular images posted below here is a single image comparing the previous benchmark to the new one. Overall the train_loss is on par, memory use is on par/a little higher, and train_tokens_per_second is larger/slower.

closes: #98

[anhuong]

Had some questions when going through the code, thanks so much Fabian. Also once the new benchmark is complete I will add the results to the scripts/benchmarks/refs with the CSV and the requirements that shows the updated deps.

[anhuong]

What is -100 ignore_index, I see that ignore_index is the target value that is ignored and does not contribute to the input gradient, but for CausalLMLoss what is at index -100?

[fabianlim]

-100 is used extensively throughout HF, while they provide some means for user to change it, almostly nobody will bother to change it

It is the label that is at -100. For a label with that value, we will ignore that token's contribution to the loss

[anhuong]

Can you describe the difference between Fast_CrossEntropyLoss and FastCrossEntropyLoss?

[fabianlim]

• Fast_CrossEntropyLoss is the autograd function. we inherit this from unsloth
• FastCrossEntropyLoss is a specialization of torch.nn.CrossEntropyLoss that serves as a convinienced, implemted using Fast_CrossEntropyLoss

[anhuong]

Would we need to create a similar FastForCausalLMLoss for liger kernel cross entropy?

[fabianlim]

yes I think we will have a new function for liger cross entropy with the same API. then its a plug and play. But it should be used only if the transformer versioin is advanced

[anhuong]

Thoughts on calling this granite-custom-crossent-loss instead to be specific that the custom loss is for cross entropy

[fabianlim]

I feel custom-loss is ok, because it mostly refers to that we are using the new custom loss feature.

[anhuong]

I also noticed in the new transformers version there is a lot of slowness after loading on the checkpoint on log line:

The new embeddings will be initialized from a multivariate normal distribution that has old embeddings' mean and covariance. As described in this article: https://nlp.stanford.edu/~johnhew/vocab-expansion.html. To disable this, use `mean_resizing=False`

Is this something we want to do? Is it expected that loading the embeddings will take a long time to run?

[fabianlim]

@anhuong the slownesss i feel its due to the recent changes in fms-hf-tuning, and is there because now we will resize the embedding layer if there is special tokens. Previously we didnt use to do that.

[anhuong]

Able to verify enabling the correct cross-entropy is triggeed based on transformers version for granite, llama, mistral, and mixtral models.

With transformers==4.48

***************** Module Forwards Patching *************
Rule: llama-custom-loss Module:                           Class: LlamaForCausalLM Num:  1
INFO:framework.py:Rule: llama-custom-loss Module:                           Class: LlamaForCausalLM Num:  1
Rule: llama-rms       Module: input_layernorm           Class: LlamaRMSNorm    Num: 32
INFO:framework.py:Rule: llama-rms       Module: input_layernorm           Class: LlamaRMSNorm    Num: 32
Rule: llama-rms       Module: model                     Class: LlamaRMSNorm    Num:  1
INFO:framework.py:Rule: llama-rms       Module: model                     Class: LlamaRMSNorm    Num:  1
Rule: llama-rms       Module: post_attention_layernorm  Class: LlamaRMSNorm    Num: 32
INFO:framework.py:Rule: llama-rms       Module: post_attention_layernorm  Class: LlamaRMSNorm    Num: 32
Rule: llama-rope      Module:                           Class: LlamaForCausalLM Num:  1
INFO:framework.py:Rule: llama-rope      Module:                           Class: LlamaForCausalLM Num:  1
***************** Accelerator Patching *************

With transformers=4.45

***************** Module Forwards Patching *************
Rule: llama-cross-ent Module:                           Class: LlamaForCausalLM Num:  1
INFO:framework.py:Rule: llama-cross-ent Module:                           Class: LlamaForCausalLM Num:  1
Rule: llama-rms       Module: input_layernorm           Class: LlamaRMSNorm    Num: 32
INFO:framework.py:Rule: llama-rms       Module: input_layernorm           Class: LlamaRMSNorm    Num: 32
Rule: llama-rms       Module: model                     Class: LlamaRMSNorm    Num:  1
INFO:framework.py:Rule: llama-rms       Module: model                     Class: LlamaRMSNorm    Num:  1
Rule: llama-rms       Module: post_attention_layernorm  Class: LlamaRMSNorm    Num: 32
INFO:framework.py:Rule: llama-rms       Module: post_attention_layernorm  Class: LlamaRMSNorm    Num: 32
Rule: llama-rope      Module:                           Class: LlamaForCausalLM Num:  1
INFO:framework.py:Rule: llama-rope      Module:                           Class: LlamaForCausalLM Num:  1
***************** Accelerator Patching *************

[anhuong]

After transformers v4.46, this method no longer exists in in transformers so I copied it in here to start. The new method to migrate to as per the warning message in the original function says to migrate to split_torch_state_dict_into_shards as noted in the TODO item here. This method was similar but requires more investigation on the difference - https://github.com/huggingface/huggingface_hub/blob/main/src/huggingface_hub/serialization/_torch.py#L302

[fabianlim]

ok this is fine for now

[anhuong]

I ran the benchmarks and found the failure in auto-gptq that I commented on the fix above. Rerunning benchmarks for auto-gptq. I also ran benchmarks for granite3.1 model but for comparison I had to run against the granite-gptcode model. Do we want to update to update to the granite3.1 model? It did run successfully with it as well.

Here are the charts showing the comparison against a100_80gb.csv without including the auto-gptq failed runs shows the train_loss is on par, memory use is on par/a little higher, and train_tokens_per_second is larger/slower

[fabianlim]

these look quite decent

[fabianlim]

LGTM the benches look good.

One question is that the new benches you ran. rather than have it in a seperate file, it is better to replace the previous a100_80gb and update the requirements, but did you run all the cases or only a subset?

[anhuong]

I will rename the benchmarks and requirements, I had run all of the benchmarks except for auto-gptq due to the error that came up so I ran those separately after the fix. I also did not run the baseline-bnb but running now separately, what is the purpose of this benchmark? updated the image in the description and will include individual images here. I will add all of them to the benchmark. I updated the above description with the summary results that include auto-gptq, Here are the individual images. Overall they continue to look good, the only outlier identified was

framework_config	peft_method	model_name_or_path	num_gpus	per_device_train_batch_size	reference	metric	new
accelerated-peft-bnb	lora	bigcode/gpt_bigcode-santacoder	2	2	6840.355	train_tokens_per_second	8548.922
accelerated-peft-bnb-foak	lora	bigcode/gpt_bigcode-santacoder	2	2	10345.932	train_tokens_per_second	11994.044

[anhuong]

I have replaced the benchmark and requirements with my full runs

[fabianlim]

I also did not run the baseline-bnb but running now separately, what is the purpose of this benchmark? updated the image in the description and will include individual images here.
the purpose of this is to have a baseline so we can compare the accelerations. The baseline could change also due to different transformer versions

[anhuong]

Makes sense, the benchmark I added has the complete benchmark including the baseline. It matches the original on the number of runs. With this, I will merge in this change

[fabianlim]

sounds good!