## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Fix#880
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: <BLANK>
- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
---------
Signed-off-by: Tcc0403 <[email protected]>
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
This PR fixes#309.
The original implementation injects Medusa heads to FSDP-unwrapped model
variables, which are not accessed from the state_dict (FSDP-wrapped). To
fix this bug, I've changed the method to inject Medusa heads with model
loader, and the method to access the state_dict through APIs from
accelerate.
Minor updates:
- Update dependencies to the latest.
- Fix scripts to run Medusa example.
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: A100-80G-NVLINK 8 cards
- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
Co-authored-by: Shao Tang <[email protected]>
## Summary
fixes#701. `strict=False` is default either way
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: NVIDIA GeForce MX230
- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
Co-authored-by: Shivam Sahni <[email protected]>
This reverts commit 5476b20d3b.
The PR has an undefined `prepare_dataset` function which causes
checkstyle to fail
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: <BLANK>
- [ ] run `make test` to ensure correctness
- [ ] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
---------
Signed-off-by: Austin Liu <[email protected]>
Co-authored-by: Austin Liu <[email protected]>
Co-authored-by: Shao Tang <[email protected]>
Co-authored-by: hebiao064 <[email protected]>
Fixes#487
I've chosen to remove the deprecated parameter as previously mentioned
in the issue. The sequence length for the training dataset can be
specified using the HF datasets library as mentioned
[here](https://huggingface.co/docs/trl/main/en/sft_trainer#datasets).
Please let me know if any further rectification is required and I will
make the necessary changes.
cc: @Tcc0403
---------
Co-authored-by: Yun Dai <[email protected]>
Co-authored-by: Shao Tang <[email protected]>
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
1. Fix Qwen2VL mrope tests. Transformers removed BC with
Qwen2VLRotaryEmbedding which is why we need this change
2. Fix mllama tests to not use cross entropy. Using Cross Entropy causes
logits mismatch which is expected as we do inplace operation in cross
entropy kernel.
3. Runs make checkstyle --- there seems to be a bug in checkstyle
workflow where unformatted checkstyle seems to pass tests. Changed the
makefile to throw error in such cases.
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: <BLANK>
- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
To fix https://github.com/linkedin/Liger-Kernel/issues/458
1. Move orpo trainer to trainer/
2. Remove trl from `[dev]` now so the test represents the actual user
environment.
3. Disable Qwen VL test for transformers>=4.47.0 for now and create
issue https://github.com/linkedin/Liger-Kernel/issues/461
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: <BLANK>
- [ ] run `make test` to ensure correctness
- [ ] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
This PR adds support for the following:
1. LigerORPOTrainer: a wrapper on top of [HuggingFace ORPO
Trainer](https://github.com/huggingface/trl/blob/main/trl/trainer/orpo_trainer.py)
to use LigerORPOLoss module.
2. We also provide an example for using LigerORPOTrainer in
`examples/alignment/run_orpo.py`
3. Change FusedLinearPreference base class' forward function to return
additional metrics to align our implementation with HF ORPO Trainer
4. Additional refactor to avoid torch compile recompilations --
accumulate_chunk function now calls accumulate_helper which is torch
compiled solely and input_chunk/target_chunk/target dimension 1 (seq
len) is explicitly marked as dynamic to avoid recompilations
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: <BLANK>
- [ ] run `make test` to ensure correctness
- [ ] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
## Summary
Replica of #396
Adds xpu support so all tests, benchmarks etc. run on XPUs or Intel
GPUs.
## Details
infer_device() function is moved to a separate file and in any file
where previously "cuda" was needed, infer_device is imported and "cuda"
is replaced with return value of a call to infer_device()
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
A100 80GB PCIe, RTX 3060, Intel Data Center GPU Max 1550
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: <BLANK>
- [x] run `make test` to ensure correctness
- [ ] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
---------
Co-authored-by: Shao Tang <[email protected]>
## Summary
Add a script that will allow users to run the examples on Modal, even if
they don't have access to a 4xA100 80Gb node at work.
Also fix hparams for tuning qwen2-vl on a 4xA100 80Gb node
## Testing Done
- Hardware Type:
Modal launched from macbook, script runs on modal on 4x80Gb A100s
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
## Summary
Example demonstrating how to use `SFTTrainer` to finetune Qwen2-VL on a
multimodal dataset using Liger-Kernel.
## Testing Done
- Hardware Type: 2x RTX 4090
Ran on two RTX 4090s, however I hit CUDA OOMs after about 10 steps (even
with per-device batch_size=2, adafactor, shorter max_seq_length...)
Would love some help from a LinkedIn employee with access to A100s/H100s
at work to help verify.
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
---------
Co-authored-by: Byron Hsu <[email protected]>
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Currently, there are two errors on Medusa training examples:
1. When we use Liger Kernel, the first head (`model.medusa_head[0]`) is
not trained.
2. When we don't use Liger Kernel, the logits of the last head
(`medusa_logits[-1]`) is ignored.
This PR fixes these errors.
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: A100 80GB 8 GPUs
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
There is a bug of incorrect weight sharing between layers for each
Medusa head. Since `nn.Module` is Python reference, the original source
code creates a list containing references to the same weights. This PR
fixes the bug.
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: A100-80G-PCIe
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
## Summary
- Added Hugging Face training benchmarking script used for tech report
- Writes files to
`/results/${MODEL_TYPE}_use_liger_${USE_LIGER}_batch_size_${BATCH_SIZE}_rep_${i}.log`
## Testing Done
- Ran benchmarking script
- Hardware Type: A100
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Switch to Qwen2 0.5B instruct model with 1024 max seq length for L40
training (control max memory to < 48G)
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: <BLANK>
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Add gemma 7b it benchmark
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
Add gemma 7b it benchmark
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
NA
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: <BLANK>
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
---------
Co-authored-by: jaszhu <[email protected]>
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Add missing tf_keras to req.txt
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
Add missing tf_keras to req.txt , otherwise pip install -r
requirement.txt and then run the script will result in
```
ValueError: Your currently installed version of Keras is Keras 3, but this is not yet supported in Transformers. Please install the backwards-compatible tf-keras package with `pip install tf-keras`.
```
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
Done
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: <BLANK>
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
Co-authored-by: jaszhu <[email protected]>
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Add gemma lightning example for single L40 GPU
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: <BLANK>
- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
---------
Co-authored-by: Shao Tang <[email protected]>
Co-authored-by: Byron Hsu <[email protected]>
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Add code to save medusa heads and model
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
Add code to save medusa heads and model, when medusa and
medusa_only_head are both enabled, we save only heads. Otherwise the
whole model will be saved.
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
Head module is saved:
```
medusa_lm_head.safetensors
```
Whole model is saved:
```
config.json model-00001-of-00005.safetensors model-00004-of-00005.safetensors special_tokens_map.json training_args.bin
generation_config.json model-00002-of-00005.safetensors model-00005-of-00005.safetensors tokenizer.json
model-00003-of-00005.safetensors model.safetensors.index.json tokenizer_config.json
```
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: <BLANK>
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
---------
Co-authored-by: jaszhu <[email protected]>
## Summary
- Adding AutoLigerKernelForCausalLM to automatically apply relevant
liger kernels to corresponding model to address #70
- Args passed to auto model that are part of
`apply_liger_kernel_to_{model_type}` supported args are passed through
- Refactored the monkey_patch.py so the mapping and patching functions
are all in the same file so it's easier to not miss updating. Also added
test to ensure the patching functions match the auto model mapping.
Example:
```python
model = AutoLigerKernelForCausalLM.from_pretrained(
custom_args.model_name,
trust_remote_code=True,
use_cache=False,
torch_dtype=torch.bfloat16,
cross_entropy=True,
fused_linear_cross_entropy=False,
)
> INFO:liger_kernel.transformers.monkey_patch:Applying Liger kernels for model type: llama with kwargs: {'cross_entropy': True, 'fused_linear_cross_entropy': False}
```
## Testing Done
- Run huggingface example
```python
assert modeling_llama.LlamaRMSNorm != LigerRMSNorm
model = AutoLigerKernelForCausalLM.from_pretrained(
custom_args.model_name,
trust_remote_code=True,
use_cache=False,
torch_dtype=torch.bfloat16,
# These args will get passed to the appropriate apply_liger_kernel_to_* function
# to override the default settings
cross_entropy=True,
fused_linear_cross_entropy=False,
)
assert modeling_llama.LlamaRMSNorm == LigerRMSNorm
```
- Hardware Type: A100
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
```
$ make checkstyle
flake8 .; flake8_status=$?; \
isort .; isort_status=$?; \
black .; black_status=$?; \
if [ $flake8_status -ne 0 ] || [ $isort_status -ne 0 ] || [ $black_status -ne 0 ]; then \
exit 1; \
fi
Skipped 3 files
All done! ✨🍰✨
54 files left unchanged.
$ make test
pytest --disable-warnings test/ --ignore=test/convergence
==================================================================================================================================== test session starts ====================================================================================================================================
platform linux -- Python 3.10.14, pytest-7.1.2, pluggy-1.0.0
rootdir: /home/jobuser/Liger-Kernel
plugins: lipy-config-base-32.0.27, lipy-fabric-35.2.19, lipy-test-8.0.66, datadir-1.3.1
collected 117 items
test/transformers/test_auto_model.py . [ 0%]
test/transformers/test_cross_entropy.py .......................................................... [ 50%]
test/transformers/test_fused_linear_cross_entropy.py ...... [ 55%]
test/transformers/test_geglu.py ........ [ 62%]
test/transformers/test_monkey_patch.py ..... [ 66%]
test/transformers/test_rms_norm.py ................ [ 80%]
test/transformers/test_rope.py ............ [ 90%]
test/transformers/test_swiglu.py ........ [ 97%]
test/transformers/test_trainer_integration.py . [ 98%]
test/triton/test_triton_monkey_patch.py .. [100%]
=================================================================================================================================== 117 passed in 50.75s ====================================================================================================================================
$ make test-convergence
HF_DATASETS_OFFLINE=1 pytest --disable-warnings test/convergence
==================================================================================================================================== test session starts ====================================================================================================================================
platform linux -- Python 3.10.14, pytest-7.1.2, pluggy-1.0.0
rootdir: /home/jobuser/Liger-Kernel
plugins: lipy-config-base-32.0.27, lipy-fabric-35.2.19, lipy-test-8.0.66, datadir-1.3.1
collected 8 items
test/convergence/test_mini_models.py ...... [ 75%]
test/convergence/test_mini_models_no_logits.py .. [100%]
==================================================================================================================================== 8 passed in 47.97s =====================================================================================================================================
```
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
This PR adds qwen2 7b benchmarks
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: A100-80G-PCIe
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
---------
Co-authored-by: shisahni <[email protected]>
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Same problem still appears after #78
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: <BLANK>
- [ ] run `make test` to ensure correctness
- [ ] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Fix example bugs:
1) lightning data loading
original loaded data has format below so should be flattened to extract
nested features.
```
DatasetDict({
train: Dataset({
features: ['train'],
num_rows: 99842
})
})
```
2) medusa sklearn learn dependency: public scikit-learn should use
full-name
misc: readme clarify test A100 clip HBM size, add deepspeed dependency
in lightning requirements
Thanks to https://github.com/t-vi who found out the issues.
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: <BLANK>
- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence
## Summary
- There are four images I believe are unused in
`examples/medusa/docs/images/`
- These images were replaced by updated benchmark results in this PR:
https://github.com/linkedin/Liger-Kernel/pull/59/
- Furthermore these images have filenames that only differ by one letter
case, causing issues on cloning the repo on case-insensitive filesystems
like those used on Windows and MacOS. (Linux not affected) You would see
an error like this:
```
git clone [email protected]:pchng/Liger-Kernel.git
Cloning into 'Liger-Kernel'...
remote: Enumerating objects: 735, done.
remote: Counting objects: 100% (342/342), done.
remote: Compressing objects: 100% (213/213), done.
remote: Total 735 (delta 211), reused 213 (delta 121), pack-reused 393 (from 1)
Receiving objects: 100% (735/735), 8.18 MiB | 11.39 MiB/s, done.
Resolving deltas: 100% (306/306), done.
warning: the following paths have collided (e.g. case-sensitive paths
on a case-insensitive filesystem) and only one from the same
colliding group is in the working tree:
'examples/medusa/docs/images/Memory_Stage1_num_head_3.png'
'examples/medusa/docs/images/memory_Stage1_num_head_3.png'
'examples/medusa/docs/images/Memory_Stage1_num_head_5.png'
'examples/medusa/docs/images/memory_Stage1_num_head_5.png'
'examples/medusa/docs/images/Memory_Stage2_num_head_3.png'
'examples/medusa/docs/images/memory_Stage2_num_head_3.png'
'examples/medusa/docs/images/Memory_Stage2_num_head_5.png'
'examples/medusa/docs/images/memory_Stage2_num_head_5.png'
```
## Testing Done
- Ran `grep -r "memory_Stage" ./*`, didn't find references to the
deleted image files.
- The relevant `README.md` file still looks OK after removing these
unused images:
https://github.com/pchng/Liger-Kernel/blob/50345cfb881388ce4de44210e08bd8c43cf1ed04/examples/medusa/README.md
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Update benchmark result of Medusa for batch size = 6 setup
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
Update benchmark result of Medusa for batch size = 6 setup, and link the
plot for throughput and memory
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
Done
<!--
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- [x ] run `make test` to ensure correctness
- [x ] run `make checkstyle` to ensure code style
- [x ] run `make test-convergence` to ensure convergence
Co-authored-by: jaszhu <[email protected]>
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
This PR adds example of using liger kernel for finetuning with lightning
trainer.
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
Worked on on-prem but need some lightning studio credits to test the
deps and configs outside :-0
# Summary
Add medusa patching code and example job with memory efficient Liger
Kernel
Medusa is a simple framework that democratizes the acceleration
techniques for LLM generation with multiple decoding heads.
[[repo](https://arxiv.org/abs/2401.10774)],
[[paper](https://arxiv.org/abs/2401.10774)]
During training, Medusa requires adding \(k\) decoding heads to the
hidden states right before the regular LM head \(h_t\). The \(k\)-th
head is used to predict the token in the \((t + k + 1)\)-th position of
the next tokens (the original language model head is used to predict the
\((t + 1)\)-th position).
The Liger fused CE kernel is highly effective in this scenario,
eliminating the need to materialize logits for each head, which usually
consumes a large volume of memory due to the extensive vocabulary size
(e.g., for LLaMA-3, the vocabulary size is 128k). The introduction of
multiple heads can easily lead to OOM (Out of Memory) issues. However,
thanks to the efficient Liger fused CE, which calculates the gradient in
place and doesn't materialize the logits, we have observed very
effective results. This efficiency opens up more opportunities for
multi-token prediction research and development.
# List of changes
* Add self-contained medusa training flow with corresponding fsdp config
* Add README.md with benchmark result
* Add patching code to apply multi-head to backbone model
### TODO ([JasonZhu1313](https://github.com/JasonZhu1313))
* [P0] Clean up the code and remove hard-code coefficient in patching
code
* [P1] parallelize the llm head inference code
# Testing done
Convergence test
Liger kernel is very effective in memory usage for medusa training:
---------
Co-authored-by: jaszhu <[email protected]>
Co-authored-by: ByronHsu <[email protected]>
Tcc0403