PyTorch Hooks Injection System

Relevant source files

The following files were used as context for generating this wiki page:

• neuraldbg.py
• infrastructure_planning/CODEBASE_GUIDE.md
• README.md
• plans/reasoning_engine_analysis.md
• ROADMAP.md

The PyTorch Hooks Injection System is the primary observation mechanism of the NeuralDBG engine. It allows the system to non-invasively monitor neural network training dynamics by attaching callbacks to model layers. This system is designed to extract Semantic Events without requiring modifications to the user's existing model code or training loops.

By operating at module boundaries, the injection system remains compiler-aware and compatible with torch.compile. It captures raw tensor statistics during forward and backward passes, which are then used for Causal Inference to identify root causes of training failures such as vanishing or exploding gradients.

Sources: README.md, infrastructure_planning/CODEBASE_GUIDE.md

Core Architecture and Injection Logic

The system utilizes the NeuralDbg class as a context manager to manage the lifecycle of hook injection. When a model is wrapped within this context, the engine traverses the model's architecture and attaches hooks to every layer.

Hook Lifecycle Management

The injection system follows a strict lifecycle to ensure that hooks do not persist longer than necessary, preventing memory leaks and performance degradation.

This diagram illustrates the automated lifecycle of hook attachment and removal within the context manager.
Sources: neuraldbg.py:40-60, infrastructure_planning/CODEBASE_GUIDE.md

Key Components

Component	Description	Implementation Detail
_install_hooks	Iterates through all named modules and registers PyTorch hooks.	module.register_forward_hook and register_full_backward_hook.
_remove_hooks	Removes all registered hook handles to clean up the model.	Iterates through handle list and calls .remove().
Forward Hook	Captures activation statistics (mean, std, sparsity).	Monitors activation regimes (Normal, Saturated, Dead).
Backward Hook	Captures gradient statistics (norm, max, min).	Monitors gradient health (Healthy, Vanishing, Exploding).

Sources: neuraldbg.py:50-85, infrastructure_planning/CODEBASE_GUIDE.md

Data Flow and Event Detection

The hooks act as sensors that transform high-dimensional tensor data into low-dimensional semantic summaries. This process is essential for memory efficiency as it avoids storing full tensors.

Semantic Extraction Process

When a hook triggers during training, it performs the following logic:

1. Raw Observation: Accesses the input/output tensors of the layer.
2. Metric Calculation: Computes norms or health indicators (e.g., gradient norm < 1e-6 for vanishing).
3. State Classification: Compares metrics against defined thresholds for ActivationHealth or GradientHealth.
4. Event Generation: If a transition is detected (e.g., Healthy → Vanishing), a SemanticEvent is created and stored.

This sequence diagram shows how raw tensor data is processed into semantic events during a single pass.
Sources: neuraldbg.py:120-150, plans/reasoning_engine_analysis.md, infrastructure_planning/CODEBASE_GUIDE.md

Performance and Memory Safety

The injection system is built to be "Hub & Spokes" compliant, isolating the monitoring logic from the core framework to ensure reliability.

• Compiler Awareness: Hooks are attached at module boundaries to survive torch.compile optimizations.
• Non-Invasive: It does not pause execution or store full tensor history, maintaining a low memory footprint.
• Garbage Collection: The system uses handle management to ensure all PyTorch hooks are detached during __exit__.
• Complexity: Hook installation is $O(L)$ where $L$ is the number of layers, and processing overhead is $O(1)$ per layer per pass.

Sources: README.md, plans/reasoning_engine_analysis.md, ROADMAP.md

Limitations and Future Enhancements

The current implementation of the hook injection system focuses on a specific subset of training dynamics:

• Framework Limitation: Currently only supports PyTorch.
• Architecture Awareness: It does not yet specifically track skip connections (Residuals) or normalization layer internals (BatchNorm/LayerNorm running stats).
• Automatic Step Tracking: Future versions aim to increment step counts automatically via backward hooks to reduce temporal coupling issues.

Sources: plans/reasoning_engine_analysis.md, ROADMAP.md

The PyTorch Hooks Injection System serves as the foundational "sensor array" for NeuralDBG, enabling high-level causal analysis by providing the necessary semantic data without the overhead of traditional heavy-weight debuggers.