1"""
  2model_inspector.py
  3------------------
  4Utilities for inspecting HuggingFace transformer models: enumerating weights,
  5capturing raw CUDA pointers, registering activation hooks, and building a
  6full memory map — ready for use with custom CUDA kernels or OpenGL interop.
  7
  8Usage:
  9    from model_inspector import ModelInspector
 10
 11    inspector = ModelInspector(model)
 12    inspector.print_architecture()
 13    inspector.print_parameters()
 14    inspector.print_cuda_pointers()
 15    inspector.inspect_layer(0)
 16    inspector.register_hooks()
 17    inspector.run_forward(tokenizer, "Hello world")
 18    inspector.print_activations()
 19    inspector.print_memory_map()
 20
 21    # Get a raw pointer for a CUDA kernel or GL copy:
 22    ptr = inspector.get_param_ptr("model.layers.0.self_attn.q_proj.weight")
 23    tensor = inspector.get_activation("layer_0")
 24"""
 25
 26from __future__ import annotations
 27
 28import warnings
 29import matplotlib.pyplot as plt
 30import torch
 31from typing import Any
 32
 33try:
 34    from .am_control_model_exp import *
 35except Exception:
 36    import traceback
 37    traceback.print_exc()
 38    warnings.warn("Failed to import control model")
 39
 40# ─────────────────────────────────────────────────────────────────────────────
 41# ModelInspector
 42# ─────────────────────────────────────────────────────────────────────────────
 43
 44class ModelInspector:
 45    """
 46    Wraps a HuggingFace CausalLM (or any nn.Module) and exposes utilities for
 47    inspecting weights, buffers, activations, raw CUDA pointers, and memory.
 48
 49    Parameters
 50    ----------
 51    model : torch.nn.Module
 52        A loaded HuggingFace model (or any PyTorch model on CUDA).
 53    """
 54
 55    def __init__(self, model: torch.nn.Module) -> None:
 56        self.model = model
 57        self._hooks: list[torch.utils.hooks.RemovableHook] = []
 58        self._activations: dict[str, dict[str, Any]] = {}
 59
 60    # ─────────────────────────────────────────────────────────────────────────
 61    # Architecture
 62    # ─────────────────────────────────────────────────────────────────────────
 63
 64    def print_architecture(self) -> None:
 65        """Print the top-level module tree."""
 66        print("\n" + "═" * 80)
 67        print("  MODEL ARCHITECTURE")
 68        print("═" * 80)
 69        print(self.model)
 70        print()
 71
 72    # ─────────────────────────────────────────────────────────────────────────
 73    # Parameters & buffers
 74    # ─────────────────────────────────────────────────────────────────────────
 75
 76    def print_parameters(self, cuda_only: bool = False) -> None:
 77        """
 78        Print every named parameter (trainable weight).
 79
 80        Parameters
 81        ----------
 82        cuda_only : bool
 83            When True, only show tensors that live on a CUDA device.
 84        """
 85        print( "Parameters" )
 86        print("\n" + "─" * 110)
 87        print(f"  {'NAME':55s}  {'SHAPE':28s}  {'DTYPE':12s}  DEVICE")
 88        print("─" * 110)
 89        for name, param in self.model.named_parameters():
 90            if cuda_only and param.device.type != "cuda":
 91                continue
 92            print(f"  {name:55s}  {str(tuple(param.shape)):28s}  "
 93                  f"{str(param.dtype):12s}  {param.device}")
 94        print()
 95
 96    def print_buffers(self, cuda_only: bool = False) -> None:
 97        """
 98        Print every named buffer (non-trainable tensors, e.g. RoPE cos/sin).
 99
100        Parameters
101        ----------
102        cuda_only : bool
103            When True, only show tensors that live on a CUDA device.
104        """
105        print( "Buffers" )
106        print("\n" + "─" * 110)
107        print(f"  {'NAME (BUFFER)':55s}  {'SHAPE':28s}  {'DTYPE':12s}  DEVICE")
108        print("─" * 110)
109        for name, buf in self.model.named_buffers():
110            if cuda_only and buf.device.type != "cuda":
111                continue
112            print(f"  {name:55s}  {str(tuple(buf.shape)):28s}  "
113                  f"{str(buf.dtype):12s}  {buf.device}")
114        print()
115
116    # ─────────────────────────────────────────────────────────────────────────
117    # Raw CUDA pointers
118    # ─────────────────────────────────────────────────────────────────────────
119
120    def print_cuda_pointers(self) -> None:
121        """
122        Print the raw device pointer for every CUDA parameter and buffer.
123        These pointers can be passed directly to custom CUDA kernels or used
124        with cudaMemcpy for OpenGL / Vulkan interop.
125        """
126        print( "Pointers" )
127        print("\n" + "─" * 120)
128        print(f"  {'NAME':60s}  {'SHAPE':25s}  {'DTYPE':12s}  DATA_PTR (hex)")
129        print("─" * 120)
130        for name, tensor in list(self.model.named_parameters()) + list(self.model.named_buffers()):
131            if tensor.device.type == "cuda":
132                print(f"  {name:60s}  {str(tuple(tensor.shape)):25s}  "
133                      f"{str(tensor.dtype):12s}  {hex(tensor.data_ptr())}")
134        print()
135
136    def get_param_ptr(self, param_name: str) -> int:
137        """
138        Return the raw CUDA device pointer (as a Python int / void*) for a
139        named parameter or buffer.
140
141        Example
142        -------
143        >>> ptr = inspector.get_param_ptr("model.layers.0.self_attn.q_proj.weight")
144        >>> my_cuda_ext.my_kernel(ptr, rows, cols)
145
146        Parameters
147        ----------
148        param_name : str
149            Dot-separated parameter path, e.g. "model.layers.0.self_attn.q_proj.weight".
150
151        Returns
152        -------
153        int
154            The device pointer as a Python integer.
155        """
156        # Search parameters first, then buffers
157        lookup = dict(self.model.named_parameters())
158        lookup.update(dict(self.model.named_buffers()))
159        if param_name not in lookup:
160            raise KeyError(f"No parameter or buffer named '{param_name}'")
161        tensor = lookup[param_name]
162        if tensor.device.type != "cuda":
163            raise RuntimeError(f"'{param_name}' is on {tensor.device}, not CUDA")
164        return tensor.data_ptr()
165
166    # ─────────────────────────────────────────────────────────────────────────
167    # Layer-level deep-dive (LLaMA / Mistral / GPT-NeoX style)
168    # ─────────────────────────────────────────────────────────────────────────
169
170    def inspect_layer(self, layer_idx: int) -> None:
171        try:
172            layer = self.model.model.layers[layer_idx]
173        except (AttributeError, IndexError) as exc:
174            print(f"[inspect_layer] Could not access layer {layer_idx}: {exc}")
175            return
176
177        print(f"\n{'═' * 80}")
178        print(f"  LAYER {layer_idx} DETAIL")
179        print("═" * 80)
180
181        # Dynamically discover and print every submodule (e.g., self_attn, shared_mlp, norm)
182        for name, module in layer.named_children():
183            self._print_submodule(name, module)
184        print()
185
186    def _print_submodule(self, label: str, module: torch.nn.Module) -> None:
187        """Helper: print all parameters in an arbitrary sub-module."""
188        print(f"\n  [{label}]")
189        print(f"    {'param':30s}  {'shape':25s}  {'dtype':12s}  data_ptr")
190        print("    " + "─" * 82)
191        for name, param in module.named_parameters():
192            ptr = hex(param.data_ptr()) if param.is_cuda else "cpu"
193            print(f"    {name:30s}  {str(tuple(param.shape)):25s}  "
194                  f"{str(param.dtype):12s}  {ptr}")
195
196    def visualize_weights(self, layer_idx: int = 0, cmap: str = "viridis") -> None:
197        """
198        Visualizes the weights and biases of a specific layer as 2D slices.
199        Higher-dimensional tensors (e.g., 3D or 4D) are sliced into 2D grids.
200        """
201        try:
202            # Targeting the specific layer in the Granite model
203            layer = self.model.model.layers[layer_idx]
204        except (AttributeError, IndexError):
205            print(f"Layer {layer_idx} not found.")
206            return
207
208        # Collect all parameters within this specific layer
209        params = {n: p for n, p in layer.named_parameters()}
210        
211        if not params:
212            print(f"No parameters found in layer {layer_idx}.")
213            return
214
215        num_params = len(params)
216        fig, axes = plt.subplots(num_params, 1, figsize=(10, 10 * num_params))
217        
218        if num_params == 1: axes = [axes]
219
220        for ax, (name, param) in zip(axes, params.items()):
221
222            data = param.detach().cpu().float()
223            
224            # Handle different dimensionalities
225            if data.dim() == 1:
226                # 1D (Biases/Norms) -> Reshape to a 2D row for visualization
227                data = data.unsqueeze(0)
228            elif data.dim() > 2:
229                # Multi-dimensional -> Take the first 2D slice [0, 0, :, :]
230                # This is common in MoE routers or complex head projections
231                while data.dim() > 2:
232                    data = data[0]
233
234            im = ax.imshow(data.numpy(), aspect='auto', cmap=cmap)
235            ax.set_title(f"{name} | Shape: {tuple(param.shape)}")
236            fig.colorbar(im, ax=ax)
237
238        plt.tight_layout()
239        plt.show()
240
241    def visualize_activations(self, tag: str, tokenizer: Any = None, prompt: str = None) -> None:
242        """
243        Visualizes the captured activation for a given module tag.
244        If tokenizer and prompt are provided, labels the Y-axis with tokens.
245        """
246        if tag not in self._activations:
247            print(f"No activation found for tag: '{tag}'. Did you run_forward()?")
248            return
249
250        # Extract the tensor from the stored activation data
251        data = self._activations[tag]["tensor"].detach().cpu().float()
252
253        # Squeeze batch dimension if it's 1: (1, Seq, Hidden) -> (Seq, Hidden)
254        if data.dim() == 3 and data.size(0) == 1:
255            data = data.squeeze(0)
256        
257        # If it's still > 2D (e.g. multi-head attention states), take the last 2 dims
258        while data.dim() > 2:
259            data = data[0]
260
261        fig, ax = plt.subplots(figsize=(12, 0.5 * data.size(0) + 2))
262        im = ax.imshow(data.numpy(), aspect='auto', cmap='magma')
263        
264        # Labeling tokens if possible
265        if tokenizer and prompt:
266            tokens = tokenizer.tokenize(prompt)
267            # Handle special tokens if the tokenizer adds them (like <s>)
268            if len(tokens) < data.size(0):
269                tokens = ["<bos>"] + tokens
270            
271            ax.set_yticks(range(len(tokens)))
272            ax.set_yticklabels(tokens)
273            ax.set_ylabel("Tokens")
274        else:
275            ax.set_ylabel("Sequence Position")
276
277        ax.set_xlabel("Hidden Dimension (Features)")
278        ax.set_title(f"Activations: {tag}\nShape: {tuple(data.shape)}")
279        
280        fig.colorbar(im, ax=ax, label="Activation Value")
281        plt.tight_layout()
282        plt.show()
283
284    # ─────────────────────────────────────────────────────────────────────────
285    # Forward hooks — capture live activations
286    # ─────────────────────────────────────────────────────────────────────────
287
288    def register_hooks(self) -> None:
289
290        self.remove_hooks()
291
292        def _make_hook(tag: str):
293            def _hook(module, inp, output):
294                tensor = output[0] if isinstance(output, tuple) else output
295                if not isinstance(tensor, torch.Tensor):
296                    return
297                self._activations[tag] = {
298                    "shape"    : tuple(tensor.shape),
299                    "dtype"    : str(tensor.dtype),
300                    "device"   : str(tensor.device),
301                    "data_ptr" : hex(tensor.data_ptr()) if tensor.is_cuda else "cpu",
302                    "tensor"   : tensor, 
303                }
304            return _hook
305
306        # Recursively hook every module
307        for name, module in self.model.named_modules():
308            # Skip the top-level model to avoid redundant giant hooks
309            if module is self.model or name == "":
310                continue
311                
312            # Only hook leaf modules (the actual operations like Linear, RMSNorm)
313            # to avoid duplicating tensors from parent container modules.
314            if not list(module.children()):
315                h = module.register_forward_hook(_make_hook(name))
316                self._hooks.append(h)
317
318        print(f"[register_hooks] Registered {len(self._hooks)} granular hooks.")
319
320    def remove_hooks(self) -> None:
321        """Remove all registered forward hooks."""
322        for h in self._hooks:
323            h.remove()
324        self._hooks.clear()
325
326    def run_forward(self, tokenizer, text: str) -> torch.Tensor:
327        """
328        Tokenize ``text``, run a single forward pass, and return the raw
329        logits tensor.  Hooks (if registered) will populate ``_activations``.
330
331        Parameters
332        ----------
333        tokenizer : transformers tokenizer
334        text : str
335
336        Returns
337        -------
338        torch.Tensor   — logits of shape [1, seq_len, vocab_size]
339        """
340        device = next(self.model.parameters()).device
341        inputs = tokenizer(text, return_tensors="pt").to(device)
342        with torch.no_grad():
343            outputs = self.model(**inputs, use_cache=False)
344        return outputs.logits
345
346    def inspect_activations(self) -> None:
347        for name, info in self._activations.items():
348            self.visualize_activations( name )
349            break
350
351    def print_activations(self) -> None:
352        """
353        Print a summary of every activation captured by the forward hooks.
354        Must call ``register_hooks()`` + ``run_forward()`` first.
355        """
356        print( "Activations" )
357        if not self._activations:
358            print("[print_activations] No activations captured yet. "
359                  "Call register_hooks() then run_forward().")
360            return
361
362        print("\n" + "─" * 100)
363        print(f"  {'HOOK NAME':20s}  {'SHAPE':35s}  {'DTYPE':12s}  DATA_PTR (hex)")
364        print("─" * 100)
365        for name, info in self._activations.items():
366            print(f"  {name:20s}  {str(info['shape']):35s}  "
367                  f"{info['dtype']:12s}  {info['data_ptr']}")
368        print()
369
370    def get_activation(self, hook_name: str) -> torch.Tensor:
371        """
372        Retrieve a captured activation tensor by hook name.
373
374        Example
375        -------
376        >>> tensor = inspector.get_activation("layer_0")
377        >>> my_cuda_ext.process(tensor.data_ptr(), tensor.numel())
378
379        Parameters
380        ----------
381        hook_name : str
382            One of the keys printed by ``print_activations()``.
383
384        Returns
385        -------
386        torch.Tensor  — the live activation (zero-copy)
387        """
388        if hook_name not in self._activations:
389            raise KeyError(
390                f"No activation for '{hook_name}'. "
391                f"Available: {list(self._activations.keys())}"
392            )
393        return self._activations[hook_name]["tensor"]
394
395    # ─────────────────────────────────────────────────────────────────────────
396    # Memory map
397    # ─────────────────────────────────────────────────────────────────────────
398
399    def print_memory_map(self, top_n: int = 0) -> None:
400        """
401        Print a full memory map of all parameters, sorted by size descending.
402
403        Parameters
404        ----------
405        top_n : int
406            If > 0, show only the top N largest tensors.
407        """
408        rows = []
409        total_bytes = 0
410        for name, param in self.model.named_parameters():
411            nb = param.numel() * param.element_size()
412            total_bytes += nb
413            rows.append({
414                "name"       : name,
415                "shape"      : tuple(param.shape),
416                "dtype"      : str(param.dtype),
417                "mb"         : nb / 1e6,
418                "data_ptr"   : hex(param.data_ptr()) if param.is_cuda else "cpu",
419                "contiguous" : param.is_contiguous(),
420            })
421
422        rows.sort(key=lambda r: r["mb"], reverse=True)
423        if top_n:
424            rows = rows[:top_n]
425
426        print( "Memory Map" )
427        print("\n" + "─" * 135)
428        print(f"  {'NAME':55s}  {'SHAPE':22s}  {'DTYPE':12s}  "
429              f"{'MB':8s}  {'PTR':18s}  CONTIG")
430        print("─" * 135)
431        for r in rows:
432            print(f"  {r['name']:55s}  {str(r['shape']):22s}  {r['dtype']:12s}  "
433                  f"{r['mb']:8.2f}  {r['data_ptr']:18s}  {r['contiguous']}")
434
435        print(f"\n  Total parameter memory: {total_bytes / 1e9:.3f} GB\n")
436
437    # ─────────────────────────────────────────────────────────────────────────
438    # Convenience: wrap a raw CUDA pointer as a PyTorch tensor (zero-copy)
439    # ─────────────────────────────────────────────────────────────────────────
440
441    @staticmethod
442    def from_raw_ptr(
443        ptr: int,
444        shape: tuple[int, ...],
445        dtype: torch.dtype = torch.float32,
446    ) -> torch.Tensor:
447        """
448        Wrap an external CUDA device pointer (e.g. from an OpenGL buffer mapped
449        via cudaGraphicsResourceGetMappedPointer, or from Vulkan external memory)
450        as a PyTorch tensor — zero-copy.
451
452        This is the reverse of ``get_param_ptr()``: use it when you want PyTorch
453        to operate on memory that was allocated outside of PyTorch.
454
455        Example (in a C++ extension after GL/Vulkan interop):
456            void* gl_ptr = ...;                           // from cudaGraphics*
457            auto t = torch::from_blob(gl_ptr, {n}, opts); // in C++
458
459        Or from Python with a ctypes / CUDA-Python pointer:
460            t = ModelInspector.from_raw_ptr(raw_ptr, (1024,), torch.float16)
461
462        Parameters
463        ----------
464        ptr   : int           Raw CUDA device pointer as a Python int.
465        shape : tuple[int]    Desired tensor shape.
466        dtype : torch.dtype   Element type (must match what the memory holds).
467
468        Returns
469        -------
470        torch.Tensor — a non-owning view; the caller is responsible for lifetime.
471        """
472        import ctypes
473        storage = torch.cuda.UntypedStorage.from_file(  # type: ignore[attr-defined]
474            # Workaround: use ctypes + from_blob equivalent via storage pointer
475            # For real use, call torch::from_blob in a C++ extension instead.
476            "", False, 0
477        )
478        # Preferred path: expose via torch._utils._rebuild_tensor_v2 or
479        # a thin C++ extension.  Here we document the pattern:
480        raise NotImplementedError(
481            "from_raw_ptr is best implemented in a C++ extension using "
482            "torch::from_blob(ptr, shape, options).  "
483            "See the docstring for the equivalent C++ snippet."
484        )
485
486
487# ─────────────────────────────────────────────────────────────────────────────
488# Quick demo (run: python model_inspector.py)
489# ─────────────────────────────────────────────────────────────────────────────
490
491if __name__ == "__main__":
492
493    import argparse
494    from transformers import AutoModelForCausalLM, AutoTokenizer
495
496    parser = argparse.ArgumentParser(description="Inspect a HuggingFace model")
497    parser.add_argument("--model_path", default="../separable-alignment/transformer/g4-zepo-192-7-checkpoints/step_0001000", help="Path or HF model id")
498    parser.add_argument("--prompt", default="1211012011010", help="Prompt for forward pass")
499    parser.add_argument("--layer", type=int, default=0, help="Layer index to deep-dive")
500    parser.add_argument("--top_n", type=int, default=1000, help="Top N tensors in memory map")
501
502    args = parser.parse_args()
503
504    device = "cuda" if torch.cuda.is_available() else "cpu"
505    print(f"Loading model from '{args.model_path}' on {device} …")
506
507    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
508    model = AutoModelForCausalLM.from_pretrained(args.model_path, device_map=device)
509
510    model.eval()
511
512    inspector = ModelInspector(model)
513
514    inspector.print_architecture()
515    inspector.print_parameters(cuda_only=True)
516
517    # inspector.print_buffers(cuda_only=True)
518    # inspector.print_cuda_pointers()
519    # inspector.inspect_layer(args.layer)
520    # inspector.visualize_weights()
521
522    inspector.register_hooks()
523    inspector.run_forward(tokenizer, args.prompt)
524    inspector.print_activations()
525    inspector.inspect_activations()
526    inspector.print_memory_map(top_n=args.top_n)
527
528    inspector.remove_hooks()