adds Granite Liger Kernel option for Granite3.y models #430
Conversation
mergify
bot
added
ci-failure
dependencies
JamesKunstle
force-pushed
the
jkunstle/granite-liger-kernel
branch
2 times, most recently
from
1cd4e63 to
2b1ad16
Compare
JamesKunstle
force-pushed
the
jkunstle/granite-liger-kernel
branch
from
2b1ad16 to
e6ee97a
Compare
Signed-off-by: James Kunstle <[email protected]>
JamesKunstle
force-pushed
the
jkunstle/granite-liger-kernel
branch
from
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 |
| # 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.
| @@ -974,6 +995,9 @@ def run_training(torch_args: TorchrunArgs, train_args: TrainingArgs) -> None: | |||
| "The last checkpoint will be saved as 'last_epoch'." | |||
| ), | |||
| ) | |||
|
|
|||
| # 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
Support for Granite was added in LK v0.5.4, so we can add it as an additional performance option.