[WIP] Fused CEL #457
[WIP] Fused CEL #457
Conversation
|
Thanks Jerome and fantastic work - will check this! |
|
This is very much a WIP. The PR description isn't rendering nicely -- here's a more readable version (located in the test folder). Aside from integrating with unsloth's custom language models, need to check how kernel interplays with grad accum and grad checkpointing. Also will likely need to rebase to tidy up the commit history. |
|
Additional benchmarking for Benchmark config: {
"max_steps": 50,
"dtype": "bfloat16",
"model_id": "meta-llama/Meta-Llama-3-8B",
"batch_size": 2,
"max_seq_len": 512,
"packing": true,
"grad_accum_steps": 1,
"load_in_4bit": true,
"use_lora": true,
"fused_cel_n_loop_iters": [
1,
2,
4,
8
]
}QLoRA config:
Some observations:
|
|
@jeromeku Thanks for testing again! Hmm weird on the training loss being noticeable higher hmmm that is really really weird I can understand why the VRAM reduction are less pronounced, but unsure on the training loss issues |
|
Trying to figure out why the divergence when using When using |
Efficient Fused Cross Entropy Loss
Memory-efficient cross entropy implementation that only materializes the derivatives of the language modeling head layer without storing the logits and chunks the computation of the logits such that the full logits tensor is never realized.
This is a direct adaptation of this repo.
Contents
Overview
In short:
dXhereafter), and the derivative with respect to the logits projection weights (dWhereafter) are computed in chunksSee the original repo for an excellent explanation of the design.
Changes
The following changes were made to the original kernel:
3-Dlanguage modeling tensors with the required shapes of the kernel.losstofloat32, which in the original kernel was initialized to the autocasted / in-feat dtype.torch.cuda.amp.{custom_fwd,custom_bwd}to theautograd.Function.All changes are enumerated in
unsloth/kernels/fused_cel.py.Additionally, adapter layers and configs in
fused_cel.pyenable integration withtransformersandunsloth.Tests
See
tests/test_CEL.pyfor correctness checks.The comments in the tests describe numerical edge cases.
Benchmarks
Following are results from preliminary benchmarking / testing on a
L4NVIDIA GPU for a smallllama-likemodel with and without thefused CELlayer.The takeaway is that the memory efficiency claims of the original
repoare evident, with overall memory usage lower, decreasing linearly with the number of loop iterations.Can be reproduced by passing the provided options to
benchmark_hf_test_cel.py(run with--helpto see all options).Below is the overall config, followed by
training losses/grad normsand overalltraining metricsforfloat32andbfloat16.Test config:max_steps=50model_id=hf-internal-testing/tiny-random-LlamaForCausalLMbatch_size=2max_seq_len=256packing=Truegrad_accum_steps=1load_in_4bit=Falseuse_lora=Falsefused_cel_n_loop_iters=[1, 2, 4]float32Training metricsforfloat32:bfloat16Training metricsforbfloat16Next Steps
FastLanguageModelLoRAandQLoRAconfigs