Coverage for transformer_lens/benchmarks/component_benchmark.py: 33%

1"""Component-level benchmarks to compare individual model pieces. 

2 

3This module provides benchmarks for comparing individual model components 

4(attention, MLP, embedding, etc.) between HuggingFace and TransformerBridge. 

5""" 

6 

7from typing import Any, Optional 

8 

9from transformer_lens.benchmarks.component_outputs import ComponentBenchmarker 

10from transformer_lens.benchmarks.utils import BenchmarkResult, BenchmarkSeverity 

11 

12 

13def benchmark_all_components( 

14 bridge, 

15 hf_model, 

16 atol: float = 1e-4, 

17 rtol: float = 1e-4, 

18 reference_model: Optional[Any] = None, 

19) -> BenchmarkResult: 

20 """Comprehensive benchmark of all model components. 

21 

22 This function systematically tests every component in the model using the 

23 architecture adapter to find and compare equivalent components. 

24 

25 Args: 

26 bridge: The TransformerBridge model 

27 hf_model: The HuggingFace model to compare against 

28 atol: Absolute tolerance for comparison 

29 rtol: Relative tolerance for comparison 

30 reference_model: Optional reference model (unused, for API consistency) 

31 

32 Returns: 

33 BenchmarkResult summarizing all component tests 

34 """ 

35 try: 

36 # Set up component testing (e.g., sync rotary_emb references for Gemma models, eager attention) 

37 # This must be called before creating the ComponentBenchmarker 

38 bridge.adapter.setup_component_testing(hf_model, bridge_model=bridge) 

39 

40 # Create benchmarker 

41 benchmarker = ComponentBenchmarker( 

42 bridge_model=bridge, 

43 hf_model=hf_model, 

44 adapter=bridge.adapter, 

45 cfg=bridge.cfg, 

46 atol=atol, 

47 rtol=rtol, 

48 ) 

49 

50 # Skip modality towers for multimodal models — they require image or 

51 # audio inputs that isolated text-based component testing cannot 

52 # provide. Vision components are validated separately in Phase 7. 

53 skip_components = [] 

54 if getattr(bridge.cfg, "is_multimodal", False): 54 ↛ 55line 54 didn't jump to line 55 because the condition on line 54 was never true

55 skip_components = [ 

56 "vision_encoder", 

57 "vision_projector", 

58 "audio_encoder", 

59 "audio_projector", 

60 ] 

61 if getattr(bridge.cfg, "is_audio_model", False): 61 ↛ 63line 61 didn't jump to line 63 because the condition on line 61 was never true

62 # Audio preprocessing needs waveform input; validated in Phase 8 

63 skip_components.extend(["audio_feature_extractor", "feat_proj", "conv_pos_embed"]) 

64 

65 # Run comprehensive benchmark 

66 report = benchmarker.benchmark_all_components(skip_components=skip_components) 

67 

68 # Convert to BenchmarkResult format 

69 if report.failed_components == 0: 69 ↛ 83line 69 didn't jump to line 83 because the condition on line 69 was always true

70 return BenchmarkResult( 

71 name="all_components", 

72 severity=BenchmarkSeverity.INFO, 

73 passed=True, 

74 message=f"All {report.total_components} components produce equivalent outputs", 

75 details={ 

76 "total_components": report.total_components, 

77 "pass_rate": report.pass_rate, 

78 "component_types": report.get_component_type_summary(), 

79 }, 

80 ) 

81 else: 

82 # Get failure details 

83 failures_by_severity = report.get_failure_by_severity() 

84 

85 # Determine overall severity 

86 if failures_by_severity["critical"]: 

87 severity = BenchmarkSeverity.ERROR 

88 elif failures_by_severity["high"]: 

89 severity = BenchmarkSeverity.DANGER 

90 else: 

91 severity = BenchmarkSeverity.WARNING 

92 

93 # Create failure message 

94 failure_summary = [] 

95 for sev in ["critical", "high", "medium", "low"]: 

96 count = len(failures_by_severity[sev]) 

97 if count > 0: 

98 failure_summary.append(f"{count} {sev}") 

99 

100 message = ( 

101 f"{report.failed_components}/{report.total_components} components failed " 

102 f"({', '.join(failure_summary)})" 

103 ) 

104 

105 # Collect failed component details 

106 failed_details = {} 

107 for result in report.component_results: 

108 if not result.passed: 

109 failed_details[result.component_path] = { 

110 "max_diff": result.max_diff, 

111 "mean_diff": result.mean_diff, 

112 "severity": result.get_failure_severity(), 

113 "error": result.error_message, 

114 } 

115 

116 return BenchmarkResult( 

117 name="all_components", 

118 passed=False, 

119 severity=severity, 

120 message=message, 

121 details={ 

122 "total_components": report.total_components, 

123 "passed_components": report.passed_components, 

124 "failed_components": report.failed_components, 

125 "pass_rate": report.pass_rate, 

126 "failures": failed_details, 

127 }, 

128 ) 

129 

130 except Exception as e: 

131 return BenchmarkResult( 

132 name="all_components", 

133 passed=False, 

134 severity=BenchmarkSeverity.ERROR, 

135 message=f"Error running comprehensive component benchmark: {str(e)}", 

136 details={"exception": str(e)}, 

137 )