Coverage for transformer_lens/model_bridge/generalized_components/bloom_mlp.py: 77%

1"""BLOOM-specific MLP bridge component. 

2 

3BLOOM MLP requires a special 'residual' argument that standard MLPBridge doesn't handle. 

4This custom component passes the residual argument through to the original component. 

5""" 

6from typing import Any, Dict, Optional 

7 

8import torch 

9 

10from transformer_lens.model_bridge.generalized_components.base import ( 

11 GeneralizedComponent, 

12) 

13from transformer_lens.model_bridge.generalized_components.mlp import MLPBridge 

14 

15 

16class BloomMLPBridge(MLPBridge): 

17 """MLP bridge for BLOOM models that handles residual connections. 

18 

19 BLOOM MLP has a unique forward signature that requires: 

20 - hidden_states (first positional arg) 

21 - residual (keyword arg): The residual connection tensor 

22 

23 This bridge ensures the residual argument is properly passed through. 

24 """ 

25 

26 def __init__( 

27 self, 

28 name: Optional[str], 

29 config: Optional[Any] = None, 

30 submodules: Optional[Dict[str, GeneralizedComponent]] = None, 

31 ): 

32 """Initialize the BLOOM MLP bridge. 

33 

34 Args: 

35 name: The name of the component in the model 

36 config: Optional configuration 

37 submodules: Dictionary of submodules to register (e.g., dense_h_to_4h, dense_4h_to_h) 

38 """ 

39 super().__init__(name, config, submodules or {}) 

40 

41 def set_original_component(self, original_component: torch.nn.Module) -> None: 

42 super().set_original_component(original_component) 

43 # The Megatron-TP replay path (pretraining_tp>1 + slow_but_exact) computes 

44 # dense_4h_to_h via F.linear on weight slices, bypassing the module call the 

45 # out-projection hooks attach to. Force the module path; the only difference 

46 # is fp summation order. 

47 if getattr(original_component, "slow_but_exact", False): 

48 setattr(original_component, "slow_but_exact", False) 

49 

50 def forward(self, *args: Any, **kwargs: Any) -> Any: 

51 """Forward pass through BLOOM MLP with hooks. 

52 

53 BLOOM MLP requires these arguments: 

54 - hidden_states (first positional arg) 

55 - residual (second positional arg) 

56 

57 Args: 

58 *args: Input arguments (hidden_states, residual) 

59 **kwargs: Additional keyword arguments (if any) 

60 

61 Returns: 

62 Output tensor from BLOOM MLP 

63 """ 

64 if self.original_component is None: 64 ↛ 65line 64 didn't jump to line 65 because the condition on line 64 was never true

65 raise RuntimeError( 

66 f"Original component not set for {self.name}. Call set_original_component() first." 

67 ) 

68 

69 # Apply hook_in to hidden_states (first positional argument) 

70 if len(args) > 0 and isinstance(args[0], torch.Tensor): 70 ↛ 73line 70 didn't jump to line 73 because the condition on line 70 was always true

71 hooked_input = self.hook_in(args[0]) 

72 args = (hooked_input,) + args[1:] 

73 elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor): 

74 kwargs["hidden_states"] = self.hook_in(kwargs["hidden_states"]) 

75 

76 # BLOOM MLP requires residual as second positional arg 

77 # The original BLOOM block passes it, so we just pass everything through 

78 # No need to validate since the original component will handle it 

79 

80 output = self.original_component(*args, **kwargs) 

81 

82 output = self.hook_out(output) 

83 

84 return output