patching LigerRMSNorm with partial arguments given would cause error in _init_weights() #739

🐛 Describe the bug

huggingface/transformers#37070 brought a clean solution to correctly init weights for most composite models.

However, it took a isintance(module, XXXRMSNorm) to fill initial values into RMSNorm modules, which can lead to a TypeError if we patch RMSNorm with partial(LigerRMSNorm, ...).

Take gemma for example:
If we make a partial object to match the default arguments for GemmaRMSNorm

, it will break isinstance(module, GemmaRMSNorm)

TypeError: isinstance() arg 2 must be a type, a tuple of types, or a union

Solution

Instead of creating a partial object, we can make a module with desired default arguments like gemma3 did.

By the way, we probably want to move these modules under transformers/rms_norm.py instead of a new module.

Reproduce

No response

Versions

liger_kernel==cd6ec32
transformers==4.52.4