google-deepmind / google-deepmind/gemma

Fix infinite recursion in LoRA module repr when inside interceptor

Open
#473 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
5.7k
Forks
1k
Avg merge
10h 33m
Merged PRs (30d)
2

Description

## Problem

The LoRA wrapper classes (`LoRADense`, `LoRAEinsum`, `LoRADenseGeneral`) caused **infinite recursion** when `repr()` was called inside a `ModuleInterceptor` context. This occurred because the default `__repr__` implementation called `repr(self.wrapped)`, which could re-trigger the interceptor and create a recursion loop.

### Previous State
- Calling `repr()` on LoRA modules inside interceptors caused a **stack overflow**
- Made debugging and inspection of LoRA-wrapped modules difficult
- A TODO comment indicated this was a known issue

## Solution

Added custom `__repr__` methods to all three LoRA wrapper classes that:
- Use `type(self.wrapped).__name__` instead of `repr(self.wrapped)` to avoid recursion
- Display useful information (rank, wrapped module type, and key attributes)
- Safely access simple attributes without triggering interceptors

## Implementation

- **`LoRADense.__repr__()`**
- Shows rank and wrapped `Dense` features

- **`LoRAEinsum.__repr__()`**
- Shows rank and wrapped `Einsum` `einsum_str` / shape

- **`LoRADenseGeneral.__repr__()`**
- Shows rank and wrapped `DenseGeneral` features / axis / `batch_dims`

PR opened for changes: [#474 ](https://github.com/google-deepmind/gemma/pull/474)

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.