feat: add liger kernel with fused cross entropy loss

[anhuong]

No description provided.

[anhuong]

Question: do we need thisadd_start_docstrings_to_model_forward header that exists in https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L1133 but is different for each model type?

[fabianlim]

I thinik dont need

[anhuong]

The forward methods are very similar for llama, granite, and mistral however the above line as you can see is within granite here: https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/granite/modeling_granite.py#L1102 --> how to handle this difference?

[fabianlim]

the naive method is to copy and modify multiple forward functions for the fused kernel. The better solution is the other design

[anhuong]

Added a different forward method for mixtral since it was different than the others in line 550-561 and in the object returned is of different type. Is adding a separate forward method the best way to handle this?

[anhuong]

I'm guessing granite here refers to teh GraniteForCausalLM models and the granite with llama architecture models would be handled within the llama.py correct? I'm not sure if the liger kernels work for any of the granite models, do you know?

[fabianlim]

thats the problem with this solution, because we copy and paste the forward functions wholesale. so if the forwards are different this doesnt work

[anhuong]

Will update with benchmark data as I gather them, currently I do see the new kernel being used:

***** FMS AccelerationFramework *****
INFO:framework.py:***** FMS AccelerationFramework *****
Active Plugin: FastKernelsAccelerationPlugin. Python package: fms_acceleration_foak. Version: 0.2.1.dev0.
INFO:framework.py:Active Plugin: FastKernelsAccelerationPlugin. Python package: fms_acceleration_foak. Version: 0.2.1.dev0.
***************** Module Forwards Patching *************
INFO:framework.py:***************** Module Forwards Patching *************
Rule: llama-fused-lce Module:                           Class: LlamaForCausalLM Num:  1
INFO:framework.py:Rule: llama-fused-lce 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 *************
INFO:framework.py:***************** Accelerator Patching *************
Currently training with a batch size of: 2
max_steps is given, it will override any value given in num_train_epochs
/home/tuning/.local/lib/python3.11/site-packages/trl/trainer/sft_trainer.py:421: UserWarning: You passed `packing=True` to the SFTTrainer/SFTConfig, and you are training your model with `max_steps` strategy. The dataset will be iterated until the `max_steps` are reached.
  warnings.warn(
***** Running training *****
  Num examples = 2,364
  Num Epochs = 1
  Instantaneous batch size per device = 2
  Total train batch size (w. parallel, distributed & accumulation) = 4
  Gradient Accumulation steps = 1
  Total optimization steps = 100
  Number of trainable parameters = 1,703,936

Running with llama3-8b model, I ran with no framework, foak-fast-kernels (fast_loss, fast_rsm_layernorm, fast_rope_embeddings), and foak-fast-kernels-liger (fused_linear_loss, fast_rsm_layernorm, fast_rope_embeddings) and saw that memory decreased further by using the fused cross entropy loss but that train runtime was the same as using fast kernels with cross entropy loss.

Findings:

Framework	Num GPUs	per_device_batch	train_mem_gpu_peaked_delta	mem_nvidia_mem_reserved	mem_peak_torch_mem_alloc_in_bytes	train_runtime	train_samples_per_second	train_steps_per_second	train_tokens_per_second
No framework	2	2	17GB	31KB	25GB	253.9422	1.575	0.394	3225.931
No framework	4	1	9.5GB	19.8KB	13.5GB	144.9639	2.759	0.69	2825.53
Foak-fast-kernels	2	2	12.8GB	31.5KB	20.8GB	225.3918	1.775	0.444	3634.559
Foak-fast-kernels	4	1	7.4GB	17.8KB	11.5GB	134.7904	2.968	0.742	3038.792
Foak-fast-kernels-liger	2	2	8.6GB	19.6KB	16.7GB	225.7871	1.772	0.443	3628.197
Foak-fast-kernels-liger	4	1	5.6GB	15.8KB	9.6GB	135.1398	2.96	0.74	3030.935

• train_mem_gpu_peaked_delta
  • 25% speedup from base to foak-fast-kernels
  • 49% speedup from base to foak-fast-kernels-liger
• mem_peak_torch_mem_alloc_in_bytes
  • 16.8% speedup from base to foak-fast-kernels
  • 33% speedup from base to foak-fast-kernels-liger

[fabianlim]

I think this is problematic because fast_loss and fused_linear_loss both switch the loss, so it causes weird combinations

Maybe its better to convert fast_loss to accept both boolean and a string

fast_loss: True means legacy fast_loss: 'fused' means fused

[fabianlim]

or use fast_loss = True / False and fast_loss_impl we can have the deafult (non-fused) and "fused"

[anhuong]

But that's why we are only allowing the user to set one, if one is true, the other must be false. That prevents this issue right?

[fabianlim]

I just feel that if already having True/False switches is already abit annoying as you mentioned before. Now having constraints between them might be even more confusing

[fabianlim]

@anhuong it is a bit unexpected that you are not seeing speedups, as it was reported previously for the same model (alibet different seq lengths are batch sizes) that there are significant speedups #76

Even though the bench parameters that you tried are different, in ours we actually pack the sequences to full 2048, so we expect that this packing will end up with similar throughputs

Once difference I can see are that yours are multi gpu runs, can we first look at the issue using 1 gpu first.

[fabianlim]

@anhuong i was able to repro the liger benches by running them like this

cd $HOME/Liger-Kernel/examples/huggingface
pip install -r requirements.txt
bash run_benchmarks.sh

and doing some hacking to enable / disable the fused_cross entropy loss. I ran this with two GPUs.

Benches of FMS Accel

I have also tested this PR on FMS Accel on a set of different parameters, basically more aggressive batch sizes.

It is possible that the speedups did not show before the non-fused version OOMs

benchmarks.csv

[anhuong]

For reference, using the fms-acceleration benchmarks I tried:

• I removed the memory probes and only set memory logging to nvidia and tried with higher batch sizes but the benchmarking is still showing improved memory but the same train_runtime as running with the previous kernels. I had to run with the mistral model as the llama model with larger batch sizes was hitting OOM
  • GPU=1, batch_size=16, train_runtime=1803 (none)--> 1573 (foak-fast-kernels) -->1570 (foak-fast-kernels)-liger
  • GPU=2, batch_size=32, train_runtime=1792-->1563-->1560
  • GPU=4, batch_size=64, train_runtime=1863-->1632-->1629
  • GPU=4, batch_size=64, FT, train_runtime=2169-->1929-->1927
• Llama3-8b, 4 GPUs, max_seq_length=512, dataset_num_samples=10000
  • Lora, batch_size=32 (per_device=8) — 149 —> 131 --> 133
  • Lora, batch_size=64 (per_device=16) — 132 --> 115 --> 116
  • Lora, batch_size=100 (per_device=25) — 128 --> 111 --> 110
  • FT, batch_size=32 — 191 --> 172 --> 179
  • FT, batch_size=64 — 167 --> 149 --> 152
  • FT, batch_size=100 — 158 --> 141 --> 158

As you can see the train_runtime improves going from 0 kernels to fast kernels but from fast kernels to liger fused kernel showed minimal or no speedup as Fabian showed as well.

As noted Fabian, I had to decrease the dataset size and max_seq_length in order to run full benchmarks which is what Aaron had done in the issue,

[fabianlim]

@anhuong but isnt this a speedup for llama3?

GPU=4, batch_size=64, FT, train_runtime=80777-->78182-->72640

We expect more speedup for larger vocabularly, in which case llama3 will see more benefit than mistral

[anhuong]

Sorry that was memory metrics, wrote down the wrong info. I did see a small speedup in FT with batch_size=100 adding above nevermind it actually took longer :(

[fabianlim]

@anhuong I did quite a few changes to this PR

• refactor: I moved the liger code to be under the fused_ops folder, since its more of that
• benches: I updated the scenarios-liger.yaml benches to better feel what it should be tested.
• configs: i changed fast_loss to now accept True, False or fused_ce_liger

New Benches