[Kernels] Add swiglu and geglu MLP across models - #48335
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
On NOTE 2 — there's no guard today, and the failure mode is silent rather than loud. The hub layers hardcode the activation ( Three concrete ones on this branch:
Might be worth a generic check that each decorated forward actually matches the canonical Separately, a heads-up on the training path: it maps to |
|
Yes, exactly why I want to have a kernels native condition/guard option before merging this. WIP Inkling is a real regression, but also fine to drop t5gemmas - inkling at least could be refactored to be able to use the swiglu. Edit: fixed inkling and t5gemma (no inheritance) If someone pins a wrong version then that's on them tbh and kernels also throws warnings if the not latest version is selected |
|
[For maintainers] Suggested jobs to run (before merge) run-slow: afmoe, aimv2, aria, axk1, axk2, bamba, bitnet, cohere, cohere2, cohere2_moe, cohere_compass, csm, cwm, deepseek_ocr2, deepseek_v2, deepseek_v3 |
CI recapDashboard: View test results in Grafana |
…nel dispatch transformers is adding @use_kernel_forward_from_hub(..., condition=...) to LlamaMLP and the Gemma-family GEGLU MLPs (huggingface/transformers#48335), backed by a kernels-community layer that itself maps to LigerTiledSwiGLUMLP/LigerSiLUMulFunction. kernels.kernelize() walks modules by class and unconditionally rebinds `forward` on every instance of a class carrying `kernel_layer_name` -- including one already instance-patched by apply_liger_tiled_mlp -- even when no kernel mapping is registered, silently falling back to the class's own forward. It records no marker of having done so, so the only way to avoid the race is calling order: kernelize() must run before apply_liger_tiled_mlp on an already-built model, never after. Warn at patch time when this applies, and add regression tests that reproduce the silent revert (simulating the future LlamaMLP decoration against a plain use_kernel_forward_from_hub call) and confirm the correct order sticks.
3 participants
As per title, adds the tiled mlp kernels for swiglu / geglu across the models cc @kashif @qgallouedec
Quoting Quenting on the effects:
NOTE:
use_kernel_forward_from_hubkernels#796 for proper guarding the kernelization