Skip to content

Commit 4fe4ab6

Browse files
committed
fix(cp): prevent recursion error in wrapper and injection
- Add safety check in CPAttentionWrapper.__getattr__ - Add idempotency check in BaseRingModel.set_cp_adapter
1 parent 23bf72f commit 4fe4ab6

2 files changed

Lines changed: 10 additions & 0 deletions

File tree

src/dnet/core/models/base.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,11 @@ def set_cp_adapter(self, adapter: Any) -> None:
4242
layers = getattr(self, "layers", []) or []
4343
for i, layer in enumerate(layers):
4444
if hasattr(layer, "self_attn"):
45+
# Avoid double-wrapping
46+
if isinstance(layer.self_attn, CPAttentionWrapper):
47+
logger.debug("Layer %d already has CP adapter, skipping wrap", i)
48+
continue
49+
4550
# Wrap existing attention module
4651
layer.self_attn = CPAttentionWrapper(layer.self_attn, adapter)
4752

src/dnet/core/models/cp_layers.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,4 +111,9 @@ def __call__(
111111
return output
112112

113113
def __getattr__(self, name: str):
114+
if name == "base_attn":
115+
# Prevent infinite recursion if base_attn is missing
116+
raise AttributeError(
117+
f"'{type(self).__name__}' object has no attribute 'base_attn'"
118+
)
114119
return getattr(self.base_attn, name)

0 commit comments

Comments
 (0)