torch.compile Native Compatibility

The purpose of torch.compile Native Compatibility within NeuralDBG is to ensure that the causal inference engine can monitor and analyze neural network training dynamics even when PyTorch's Inductor compiler is active. torch.compile optimizes models by fusing kernels and altering the execution graph, which often breaks traditional debugging tools that rely on intercepting raw tensor operations or internal state.

NeuralDBG is designed to be "compiler-aware," meaning it operates at safe module boundaries where the execution graph is stable. This allows the engine to extract semantic events—such as gradient health transitions or activation regime shifts—without interfering with the compiler's optimization passes or requiring the storage of full tensors.
Sources: README.md, infrastructure_planning/CODEBASE_GUIDE.md

Architecture and Monitoring Strategy

NeuralDBG maintains compatibility with compiled models by avoiding deep tensor inspection in "hot paths." Instead, it utilizes standard PyTorch hooks attached to module boundaries. These hooks capture high-level statistical data during the forward and backward passes, which are then converted into semantic events.

Component Interaction Flow

The following diagram illustrates how NeuralDBG sits between the PyTorch Model (which may be compiled) and the Reasoning Engine to capture data without breaking the optimized execution graph.

The diagram shows that hooks are placed at the boundaries of the compiled artifacts to ensure semantic extraction persists after optimization.
Sources: README.md, plans/reasoning_engine_analysis.md:37-55

Compiler-Aware Event Extraction

To survive the transformation performed by torch.compile, NeuralDBG focuses on the extraction of high-level metadata rather than raw tensor data. This approach, referred to as "Compiler-Aware Monitoring," ensures that the engine does not force the compiler to "de-optimize" or fallback to eager mode due to unsupported tensor operations during the monitoring phase.

Monitoring Points

NeuralDBG identifies specific "safe operation points" at module boundaries. This ensures that:

• Hooks persist after the compilation process.
• Semantic events are correctly extracted from fused kernels.
• The performance overhead remains minimal as it avoids pausing execution.

Sources: ROADMAP.md:32-45, infrastructure_planning/CODEBASE_GUIDE.md

Data Flow for Compiled Models

The sequence below describes the interaction between a training loop using a compiled model and the NeuralDBG context manager.

The flow highlights that metadata is collected asynchronously via hooks during the pass, which is then utilized for post-mortem analysis.
Sources: infrastructure_planning/CODEBASE_GUIDE.md, README.md

Technical Implementation Details

The implementation of compatibility is primarily located in the NeuralDbg class within neuraldbg.py. Key methods are designed to be non-invasive to the model's computational graph.

Key Functions and Methods

• _install_hooks(model): Attaches PyTorch hooks to every layer in the model. Because these are standard hooks, they are generally respected by the compiler during graph capture.
• _classify_gradient_health(grad_norm): Categorizes gradient health using scalar thresholds. By reducing tensors to scalars immediately, it minimizes the memory footprint within the compiled loop.
• _detect_gradient_transition(): Detects changes between steps, providing the temporal data necessary for causal inference without requiring the compiler to store past tensor states.

Sources: infrastructure_planning/CODEBASE_GUIDE.md, plans/reasoning_engine_analysis.md:313-318

Performance Considerations

To maintain the performance gains of torch.compile, the engine adheres to a specific complexity profile:

• Forward/Backward Hook Complexity: $O(B \times F)$, where $B$ is batch size and $F$ is features, but primarily involves scalar reductions.
• Memory Overhead: $O(1)$ for raw tensor inspection per layer, as tensors are converted to scalar metadata immediately.

Sources: plans/reasoning_engine_analysis.md:275-285

Conclusion

The torch.compile Native Compatibility feature is essential for modern PyTorch workflows within NeuralDBG. By utilizing boundary hooks and semantic event extraction, the system provides a non-invasive way to debug complex training dynamics—such as vanishing gradients or dead neurons—while allowing the underlying model to benefit from compiler-level optimizations. This design prioritizes "Causal Compression" and "Metadata Extraction" over traditional, invasive debugging techniques.