1#!/usr/bin/env python3
2
3import os
4import sys
5import torch
6import transformers
7import json
8import textwrap
9import numpy as np
10from pathlib import Path
11
12
13def get_model_name_from_env_path(env_path_name):
14 model_path = os.getenv(env_path_name)
15 if not model_path:
16 print(f"Error: {env_path_name} environment variable not set")
17 sys.exit(1)
18
19 if not os.path.exists(model_path):
20 print(f"Error: Model file not found: {model_path}")
21 sys.exit(1)
22
23 name = os.path.basename(os.path.normpath(model_path))
24 if name.endswith(".gguf"):
25 name = name[:-5]
26
27 return name
28
29
30def summarize(tensor: torch.Tensor, name: str, max_seq: int = 3, max_vals: int = 3):
31 """
32 Print a tensor in llama.cpp debug style.
33
34 Supports:
35 - 2D tensors (seq, hidden)
36 - 3D tensors (batch, seq, hidden)
37 - 4D tensors (batch, seq, heads, dim_per_head) via flattening heads × dim_per_head
38
39 Shows first and last max_vals of each vector per sequence position.
40 """
41 t = tensor.detach().to(torch.float32).cpu()
42
43 # Determine dimensions
44 if t.ndim == 3:
45 _, s, _ = t.shape
46 elif t.ndim == 2:
47 _, s = 1, t.shape[0]
48 t = t.unsqueeze(0)
49 elif t.ndim == 4:
50 _, s, _, _ = t.shape
51 else:
52 print(f"Skipping tensor due to unsupported dimensions: {t.ndim}")
53 return
54
55 ten_shape = t.shape
56
57 print(f"ggml_debug: {name} = (f32) ... = {{{ten_shape}}}")
58 print(" [")
59 print(" [")
60
61 # Determine indices for first and last sequences
62 first_indices = list(range(min(s, max_seq)))
63 last_indices = list(range(max(0, s - max_seq), s))
64
65 # Check if there's an overlap between first and last indices or if we're at the edge case of s = 2 * max_seq
66 has_overlap = bool(set(first_indices) & set(last_indices)) or (max_seq * 2 == s)
67
68 # Combine indices
69 if has_overlap:
70 # If there's overlap, just use the combined unique indices
71 indices = sorted(list(set(first_indices + last_indices)))
72 separator_index = None
73 else:
74 # If no overlap, we'll add a separator between first and last sequences
75 indices = first_indices + last_indices
76 separator_index = len(first_indices)
77
78 for i, si in enumerate(indices):
79 # Add separator if needed
80 if separator_index is not None and i == separator_index:
81 print(" ...")
82
83 # Extract appropriate slice
84 vec = t[0, si]
85 if vec.ndim == 2: # 4D case: flatten heads × dim_per_head
86 flat = vec.flatten().tolist()
87 else: # 2D or 3D case
88 flat = vec.tolist()
89
90 # First and last slices
91 first = flat[:max_vals]
92 last = flat[-max_vals:] if len(flat) >= max_vals else flat
93 first_str = ", ".join(f"{v:12.4f}" for v in first)
94 last_str = ", ".join(f"{v:12.4f}" for v in last)
95
96 print(f" [{first_str}, ..., {last_str}]")
97
98 print(" ],")
99 print(" ]")
100 print(f" sum = {t.sum().item():.6f}\n")
101
102
103def debug_hook(name):
104 def fn(_m, input, output):
105 if isinstance(input, torch.Tensor):
106 summarize(input, name + "_in")
107 elif isinstance(input, (tuple, list)) and len(input) > 0 and isinstance(input[0], torch.Tensor):
108 summarize(input[0], name + "_in")
109 if isinstance(output, torch.Tensor):
110 summarize(output, name + "_out")
111 elif isinstance(output, (tuple, list)) and len(output) > 0 and isinstance(output[0], torch.Tensor):
112 summarize(output[0], name + "_out")
113
114 return fn
115
116
117def setup_rope_debug(model_module_path: str, function_name: str = "apply_rotary_pos_emb"):
118 """
119 Apply monkey patch to dump RoPE activations for debugging.
120
121 Args:
122 model_module_path: Path to the model module (e.g., "transformers.models.apertus.modeling_apertus")
123 function_name: Name of the RoPE function to patch (default: "apply_rotary_pos_emb")
124
125 Example:
126 from utils.common import setup_rope_debug
127 setup_rope_debug("transformers.models.apertus.modeling_apertus")
128 """
129 import importlib
130
131 # Import the module and get the original function
132 module = importlib.import_module(model_module_path)
133 orig_rope = getattr(module, function_name)
134
135 # Set torch print options for better debugging
136 torch.set_printoptions(threshold=float('inf'))
137 torch.set_printoptions(precision=6, sci_mode=False)
138
139 def debug_rope(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
140 # log inputs
141 summarize(q, "RoPE.q_in")
142 summarize(k, "RoPE.k_in")
143
144 # call original
145 q_out, k_out = orig_rope(q, k, cos, sin, position_ids, unsqueeze_dim)
146
147 # log outputs
148 summarize(q_out, "RoPE.q_out")
149 summarize(k_out, "RoPE.k_out")
150
151 return q_out, k_out
152
153 # Patch it
154 setattr(module, function_name, debug_rope)
155 print(f"RoPE debug patching applied to {model_module_path}.{function_name}")
156
157
158def save_output_data(data, tokens, prompt, model_name, type_suffix="", output_dir="data"):
159 """
160 Save output data (logits/embeddings), tokens, and prompt to files.
161
162 Args:
163 data: numpy array of floats (logits or embeddings)
164 tokens: list or array of token IDs
165 prompt: string containing the input prompt
166 model_name: name of the model
167 type_suffix: optional suffix like "-embeddings" (default: "")
168 output_dir: directory to save files (default: "data")
169
170 Creates the following files in output_dir:
171 - pytorch-{model_name}{type_suffix}.bin
172 - pytorch-{model_name}{type_suffix}.txt
173 - pytorch-{model_name}{type_suffix}-prompt.txt
174 - pytorch-{model_name}{type_suffix}-tokens.bin
175 """
176 data_dir = Path(output_dir)
177 data_dir.mkdir(exist_ok=True)
178 base_path = data_dir / f"pytorch-{model_name}{type_suffix}"
179
180 # Convert and flatten logits/embeddings
181 data = data.cpu().numpy() if isinstance(data, torch.Tensor) else np.asarray(data)
182 data = data.flatten() if data.ndim > 1 else data
183
184 # Save logits/embedding files
185 data.astype(np.float32).tofile(f"{base_path}.bin")
186 print(f"Data saved to {base_path}.bin")
187
188 with open(f"{base_path}.txt", "w") as f:
189 f.writelines(f"{i}: {value:.6f}\n" for i, value in enumerate(data))
190 print(f"Data saved to {base_path}.txt")
191
192 # Convert and flatten tokens
193 tokens = tokens.cpu().numpy() if isinstance(tokens, torch.Tensor) else np.asarray(tokens)
194 tokens = tokens.flatten() if tokens.ndim > 1 else tokens
195
196 # Save token binary file
197 tokens.astype(np.int32).tofile(f"{base_path}-tokens.bin")
198 print(f"Tokens saved to {base_path}-tokens.bin")
199
200 # Save prompt file
201 with open(f"{base_path}-prompt.txt", "w") as f:
202 f.write(f"prompt: {prompt}\n")
203 f.write(f"n_tokens: {len(tokens)}\n")
204 f.write(f"token ids: {', '.join(str(int(tid)) for tid in tokens)}\n")
205 print(f"Prompt saved to {base_path}-prompt.txt")
206
207
208def compare_tokens(original, converted, type_suffix="", output_dir="data"):
209 data_dir = Path(output_dir)
210
211 # Read tokens from both models
212 tokens1_file = data_dir / f"{original}{type_suffix}-tokens.bin"
213 tokens2_file = data_dir / f"{converted}{type_suffix}-tokens.bin"
214
215 if not tokens1_file.exists():
216 print(f"Error: Token file not found: {tokens1_file}")
217 return False
218
219 if not tokens2_file.exists():
220 print(f"Error: Token file not found: {tokens2_file}")
221 return False
222
223 tokens1 = np.fromfile(tokens1_file, dtype=np.int32)
224 tokens2 = np.fromfile(tokens2_file, dtype=np.int32)
225
226 print(f"\nComparing tokens between:")
227 print(f" Original : {original} ({len(tokens1)} tokens)")
228 print(f" Converted: {converted} ({len(tokens2)} tokens)")
229
230 if len(tokens1) != len(tokens2):
231 print(f"\n❌ Token count mismatch: {len(tokens1)} vs {len(tokens2)}")
232 return False
233
234 if np.array_equal(tokens1, tokens2):
235 print(f"\n✅ All {len(tokens1)} tokens match!")
236 return True
237
238 mismatches = np.where(tokens1 != tokens2)[0]
239 print(f"\n❌ Found {len(mismatches)} mismatched tokens:")
240
241 num_to_show = min(len(mismatches), 10)
242 for idx in mismatches[:num_to_show]:
243 print(f" Position {idx}: {tokens1[idx]} vs {tokens2[idx]}")
244
245 if len(mismatches) > num_to_show:
246 print(f" ... and {len(mismatches) - num_to_show} more mismatches")
247
248 return False
249
250
251def show_version_warning(current_version, model_version):
252 if not model_version:
253 return False
254
255 try:
256 from packaging.version import parse, InvalidVersion
257 try:
258 return parse(current_version) < parse(model_version)
259 except InvalidVersion:
260 return current_version != model_version
261 except ImportError:
262 return current_version != model_version
263
264def get_model_transformers_version(model_path):
265 if not model_path:
266 return None
267
268 config_path = Path(model_path) / "config.json"
269 if not config_path.is_file():
270 return None
271
272 try:
273 with open(config_path, "r", encoding="utf-8") as f:
274 config = json.load(f)
275 return config.get("transformers_version")
276 except (IOError, json.JSONDecodeError) as e:
277 print(f"Warning: Could not read or parse {config_path}: {e}", file=sys.stderr)
278 return None
279
280def exit_with_warning(message, model_path):
281 print(message)
282
283 if model_path and transformers is not None:
284 model_transformers_version = get_model_transformers_version(model_path)
285 transformers_version = transformers.__version__
286 if show_version_warning(transformers_version, model_transformers_version):
287 warning_message = f"""
288 =====================================================================
289 Verification failure might be due to a transformers version mismatch:
290
291 Current transformers version: {transformers_version}
292 Model's required version : {model_transformers_version}
293
294 Consider installing the version specified by the model's config:
295 pip install transformers=={model_transformers_version}
296 =====================================================================
297 """
298 print(textwrap.dedent(warning_message))
299 sys.exit(1)