adds Granite Liger Kernel option for Granite3.y models #430
Conversation
1cd4e63 to
2b1ad16
Compare
e6ee97a to
89aeca5
Compare
|
Current CI failure seems unrelated to liger_kernel addition |
|
Need to test but this could be very useful if it works correctly |
|
Currently correctness test for the kernels themselves are done in-tree in the liger-kernels repo. They assure correct convergence and logit equivalence after training. Link to test here I've also validated training-dynamics equivalence given identical batches. Link to Jira issue here. There's a roughly 1% raw improvement with batches being identical. This is expected- the real benefits will come from larger possible batch sizes! This kernel set is also missing the most important kernel: I've given the PSAP team a ref to this PR so they can quantify the improved memory headroom and find new max_batch_lens |
RobotSail
left a comment
There was a problem hiding this comment.
@JamesKunstle Thanks for adding this, a few comments
src/instructlab/training/main_ds.py
Outdated
| # Third Party | ||
| from liger_kernel.transformers import apply_liger_kernel_to_granite | ||
| except ImportError: | ||
| apply_liger_kernel_to_granite = lambda *args, **kwargs: None # pylint: disable=C3001 |
There was a problem hiding this comment.
@JamesKunstle We'll want to support other models beyond Granite, and it seems like Liger's API has a number of different functions they expose for common architectures (mistral, llama, etc.).
It seems like they provide a way to map directly from the model_type field on the model into the Liger kernel to apply. Could we please use that here to support more models?
There was a problem hiding this comment.
++ certainly, only limiting to Granite to get PSAP numbers then will expand the integration to other supported architectures.
There was a problem hiding this comment.
@JamesKunstle It seems like Liger actually has a AutoLigerKernelForCausalLM class which inherits from AutoModelForCasualLM:
https://github.com/linkedin/Liger-Kernel/blob/293bf7eec7043c8c34b3cd82975c97e4c2f4254f/src/liger_kernel/transformers/auto_model.py#L15
Here we can just make the apply_liger option general and simply attempt wrapping around any model
src/instructlab/training/main_ds.py
Outdated
| ) | ||
|
|
||
| # this will work for granite-3.y models but not granite-7b because that's a Llama 2 model arch. | ||
| parser.add_argument("--enable-granite-liger-kernel", action="store_true") |
There was a problem hiding this comment.
If we're providing this as an arg through the CLI, we should also expose this in the TrainingArgs config so the SDK can be consistent.
There was a problem hiding this comment.
++ agreed, will do once we've got the Granite numbers from PSAP
|
This pull request has merge conflicts that must be resolved before it can be |
9d05a07 to
ef82e49
Compare
ef82e49 to
910e46d
Compare
|
Test |
|
Looks like GitHub forces you to make an initial comment before the option to reopen becomes available |
|
Never mind, it looks like the reopen and comment button is a lie |
Support for Granite was added in LK v0.5.4, so we can add it as an additional performance option.