Skip to content

Commit 01a66f1

Browse files
authored
[Fix] Handle missing 'elementwise_affine' in RMSNorm extra_repr for patched models (#990)
## Summary Fixes an AttributeError encountered in `LigerRMSNorm.extra_repr` when models are patched `in-place` (e.g., using `apply_liger_kernel_to_model`). In my previous PR (#989 ), I added `elementwise_affine` to `extra_repr` to improve layer visibility. However, when layers are replaced via monkey patching, the `LigerRMSNorm` constructor is typically skipped, leaving the instance without the `elementwise_affine` attribute. ## Testing Done I verified the fix locally. The issue is reproducible via `test/transformers/test_monkey_patch.py`. After applying this fix, `test/transformers/test_monkey_patch.py` passes successfully. - Hardware Type: NVIDIA A100-SXM4-80GB - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence
1 parent 87d7057 commit 01a66f1

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

src/liger_kernel/transformers/rms_norm.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ def forward(self, hidden_states):
4646
)
4747

4848
def extra_repr(self):
49-
return f"elementwise_affine={self.elementwise_affine}, weight_shape={tuple(self.weight.shape) if self.weight is not None else None}, eps={self.variance_epsilon}, offset={self.offset}, in_place={self.in_place}, row_mode={self.row_mode}"
49+
return f"weight_shape={tuple(self.weight.shape) if self.weight is not None else None}, eps={self.variance_epsilon}, offset={self.offset}, in_place={self.in_place}, row_mode={self.row_mode}"
5050

5151

5252
class LigerRMSNormForGemma(LigerRMSNorm):

0 commit comments

Comments
 (0)