Coverage for transformer_lens/model_bridge/buffer_restore.py: 92%

1"""Repair non-persistent buffers destroyed by meta-device loading. 

2 

3transformers 5.x replaces every non-persistent buffer with ``torch.empty_like`` 

4(``modeling_utils._move_missing_keys_from_meta_to_device``) and only restores 

5rotary tables for modules exposing ``original_inv_freq``. Remote code predating 

6that attribute silently keeps uninitialized memory — internlm2 loads a table of 

7zeros, which collapses RoPE to the identity transform. 

8 

9The defect is invisible to the phase benchmarks: they compare the bridge against 

10an HF reference loaded the same way, so both sides are wrong together and agree. 

11Shared by adapters (``prepare_model``) and the benchmark's HF reference so the 

12two can never drift apart. 

13""" 

14 

15from typing import Any 

16 

17import torch 

18 

19 

20def _is_valid_inv_freq(inv_freq: torch.Tensor) -> bool: 

21 """True when the table looks like real rotary frequencies. 

22 

23 Valid tables are finite, strictly decreasing, and lie in (0, 1] — including 

24 the scaled variants (linear / NTK / llama3), whose tables keep that shape. 

25 Uninitialized memory (zeros, denormals, huge magnitudes) does not. 

26 """ 

27 if inv_freq.numel() == 0 or not bool(torch.isfinite(inv_freq).all()): 

28 return False 

29 values = inv_freq.float() 

30 if not bool((values > 0).all()) or float(values.max()) > 1.0 + 1e-6: 

31 return False 

32 return bool((values[:-1] > values[1:]).all()) if values.numel() > 1 else True 

33 

34 

35def restore_rotary_inv_freq(hf_model: Any) -> int: 

36 """Recompute invalid rotary ``inv_freq`` tables in place; returns the count. 

37 

38 Only tables failing :func:`_is_valid_inv_freq` are touched, so legitimately 

39 scaled tables are never clobbered. 

40 """ 

41 restored = 0 

42 for module in hf_model.modules(): 

43 if "RotaryEmbedding" not in type(module).__name__: 

44 continue 

45 inv_freq = getattr(module, "inv_freq", None) 

46 dim = getattr(module, "dim", None) 

47 base = getattr(module, "base", None) 

48 if not isinstance(inv_freq, torch.Tensor) or dim is None or base is None: 

49 continue 

50 if _is_valid_inv_freq(inv_freq): 

51 continue 

52 recomputed = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim)) 

53 module.inv_freq = recomputed.to(device=inv_freq.device, dtype=inv_freq.dtype) 

54 # Older remote code caches cos/sin from inv_freq at init and indexes the 

55 # cache directly (`self.cos_cached[:seq_len]`), so it must be rebuilt, 

56 # never cleared — setting it to None raises on the next forward. 

57 rebuild = getattr(module, "_set_cos_sin_cache", None) 

58 if callable(rebuild) and hasattr(module, "cos_cached"): 

59 try: 

60 rebuild( 

61 seq_len=getattr(module, "max_seq_len_cached", None) 

62 or getattr(module, "max_position_embeddings", 2048), 

63 device=inv_freq.device, 

64 dtype=module.cos_cached.dtype, 

65 ) 

66 except Exception: 

67 # Signature varies across remote forks; force a rebuild on the 

68 # next forward instead, which every such implementation guards 

69 # with `if seq_len > self.max_seq_len_cached`. 

70 module.max_seq_len_cached = 0 

71 if hasattr(module, "original_inv_freq"): 71 ↛ 72line 71 didn't jump to line 72 because the condition on line 71 was never true

72 module.original_inv_freq = module.inv_freq 

73 restored += 1 

74 return restored