Support for patching post-model initialization

[shimizust]

Summary

• Currently, calling the patching APIs after the model has been initialized will only partially patch with Liger kernels. For example, the following will still be patched:
  • Model forward() method (e.g. modeling_llama.LlamaForCausalLM.forward = lce_forward)
  • module functions (e.g. modeling_llama.apply_rotary_pos_emb = liger_rotary_pos_emb)
    but not any modules that were already instantiated and set as instance variables on the model:
  • For example: modeling_llama.LlamaRMSNorm = LigerRMSNorm will not affect existing LlamaRMSNorm instances
• This means that integrations with HF Trainer and SFTTrainer only partially work. In the case of SFTTrainer, the current integration only works fully if the user passes a path to the model (Adds experimental Liger support to SFT script huggingface/trl#1992), and SFTTrainer handles calling AutoLigerKernelForCausalLM.from_pretrained. However, both HF Trainer and SFTTrainer allow passing the model instance directly, in which case we would need a way of patching the model post-init.

API

from liger_kernel.transformers import _apply_liger_kernel_to_instance

llama_model = AutoModelForCausalLM.from_pretrained("/path/to/llama", ...)
_apply_liger_kernel_to_instance(model=llama_model)

# can also pass in model-specific args that will get passed into the correct apply_liger_kernel_to_{model_type}
_apply_liger_kernel_to_instance(model=llama_model, rope=False)

Required changes post-PR:

• Update HF Trainer and SFTTrainer to pass in model instead of model_type to _apply_liger_kernel(model=model)

Testing Done

• Tested HF example with no patching, pre-init patching, post-init patching using existing method, post-init patching of model instance variables showing that post-init instance patching results in same performance as pre-init patching:
• Added unit tests that model instances are actually patched correctly and that each patching API supports model instance patching.

Llama

Scenario	Memory Usage (GB)	Throughput (tokens/sec)
No patching	68.3	9161
Patch post-init (existing method)	39.9	10939
Patch pre-init	38.4	12313
Patch post-init (instance patching)	38.3	12409

Mistral

Scenario	Memory Usage (GB)	Throughput (tokens/sec)
No patching	34.5	10976
Patch post-init (existing method)	34.5	10812
Patch pre-init	34.5	12286
Patch post-init (instance patching)	34.5	12086

Mixtral
OOM on testing setup, but confirmed patched post-model init loaded correctly and started training

Gemma

Scenario	Memory Usage (GB)	Throughput (tokens/sec)
No patching	52.8	7199
Patch post-init (existing method)	40.7	9479
Patch pre-init	40.7	9209
Patch post-init (instance patching)	40.7	9981

Gemma2

Scenario	Memory Usage (GB)	Throughput (tokens/sec)
No patching	66.9	12256
Patch post-init (existing method)	51.0	14288
Patch pre-init	50.8	17084
Patch post-init (instance patching)	46.8	15893

• Note: Gemma2 pre-init and post-init (instance patching) are not converging, seeing nan gradnorm. Need to investigate separately.

Qwen2

Scenario	Memory Usage (GB)	Throughput (tokens/sec)
No patching	41.2	10785
Patch post-init (existing method)	36.8	10491
Patch pre-init	36.8	11908
Patch post-init (instance patching)	36.8	12068

Qwen2 VL

• Patched, but couldn't test since not yet released in latest transformers

Phi3

Scenario	Memory Usage (GB)	Throughput (tokens/sec)
No patching	49.4	15601
Patch post-init (existing method)	49.4	16259
Patch pre-init	42.8	20383
Patch post-init (instance patching)	48.6	18922

• Note: For phi3, post-init patching seems to consistently not match pre-init patching for some reason

• Hardware Type:
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[ByronHsu]

awesome bug catches!

[ByronHsu]

i would prefer separate this as another api. maybe apply_liger_kernel_on_instance. Would love the thoughts from @yundai424 @qingquansong

[yundai424]

seconded on separating instance-patch to a different API. Otherwise apply_liger_kernel_to_llama(model=) is just a liiitle bit cumbersome because the model itself has told what model family is it 🤔

[shimizust]

Makes sense, I can make it a separate API

[qingquansong]

Would it be good to still unified as one entry _apply_liger_kernel(model: Union[str, PreTrainedModel]) and can have apply_liger_kernel_on_instance to be used in side (but feel like not necessarily needed to have a separate instance API? 🤔

[qingquansong]

To confirm my understanding, both _apply_liger_kernel_to_instance and _apply_liger_kernel will be used in the HF api underlying based on different cases and we name it with underline to not recommend for users to directly use it (compared to using single model patch function)? 🤔

[shimizust]

To confirm my understanding, both _apply_liger_kernel_to_instance and _apply_liger_kernel will be used in the HF api underlying based on different cases and we name it with underline to not recommend for users to directly use it (compared to using single model patch function)? 🤔

@qingquansong Actually the HF Trainer API would be switched to using _apply_liger_kernel_to_instance. The _apply_liger_kernel is now meant to only be used internally by AutoLigerKernelForCausalLM.

The previous discussion with @ByronHsu was not to expose these methods publicly hence the leading _, although maybe we should re-export this from transformers/init.py to make internal refactoring less disruptive down the road.