Coverage for transformer_lens/utilities/attn_implementation.py: 100%

1"""attn_implementation. 

2 

3Shared helper for forcing eager attention on a loaded HuggingFace model. 

4""" 

5 

6from __future__ import annotations 

7 

8from typing import Any, List 

9 

10 

11def force_eager_attention(model: Any, *, per_layer: bool = False) -> None: 

12 """Switch a pre-loaded model to eager attention so attention hooks can fire. 

13 

14 Prefers the public ``set_attn_implementation`` API; exotic wrapped models can 

15 reject it, so failures fall back to writing ``config._attn_implementation`` 

16 (including nested multimodal ``text_config``). ``per_layer=True`` also stamps 

17 every submodule's ``self_attn.config`` — some models keep per-layer config 

18 copies that the top-level write never reaches. 

19 

20 Best-effort by design: never raises, silently no-ops on objects exposing 

21 neither the public API nor a config. 

22 """ 

23 handled = False 

24 if hasattr(model, "set_attn_implementation"): 

25 try: 

26 model.set_attn_implementation("eager") 

27 handled = True 

28 except Exception: 

29 pass # Exotic wrapped models can reject the public API; write the config instead. 

30 if not handled: 

31 config = getattr(model, "config", None) 

32 if config is not None and hasattr(config, "_attn_implementation"): 

33 config._attn_implementation = "eager" 

34 # Nested multimodal configs carry their own attn implementation. 

35 text_config = getattr(config, "text_config", None) 

36 if text_config is not None: 

37 text_config._attn_implementation = "eager" 

38 if per_layer: 

39 for layer in _layer_candidates(model): 

40 if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "config"): 

41 layer.self_attn.config._attn_implementation = "eager" 

42 

43 

44def _layer_candidates(model: Any) -> List[Any]: 

45 """Modules that might own a per-layer attention config. 

46 

47 Real models walk ``modules()`` so nested stacks (``model.language_model``, 

48 HRM's L/H modules) are covered; plain-object test fakes fall back to the 

49 conventional ``model.model.layers`` chain. 

50 """ 

51 modules = getattr(model, "modules", None) 

52 if callable(modules): 

53 try: 

54 return list(modules()) 

55 except TypeError: 

56 return [] # Mock-style modules() returning a non-iterable. 

57 lm = getattr(model, "model", None) 

58 layers = getattr(lm, "layers", None) if lm is not None else None 

59 return list(layers) if layers is not None else []