Introducing ModuleLens: Isolated PyTorch Module Debugging and Agentic Optimization
A module-first framework enabling local debugging and autonomous AI agent optimization for large-scale models.
1The Multi-GPU Bottleneck
Modern recommendation systems require distributed model parallelism, making local debugging on single-GPU workstations impossible. Engineers face significant hurdles when modifying downstream modules, as traditional workflows demand initializing the entire multi-GPU pipeline. This process wastes immense computation resources and creates slow feedback loops, hindering the development velocity of complex machine learning architectures.
2Module-First Engineering Paradigm
ModuleLens shifts the development paradigm from model-first to module-first, allowing engineers to isolate specific components. By reconstructing and profiling modules locally within seconds, teams avoid full multi-GPU initialization. This modular foundation provides a structured execution environment, enabling both human developers and autonomous AI agents to analyze and refactor PyTorch code independently.
3ModuleTracer Runtime Capture
ModuleTracer captures essential runtime data, including input tensors, parameter weights, and initialization arguments. It utilizes non-intrusive PyTorch forward hooks and dynamic Python import patching to ensure seamless data logging. To maintain efficiency, the system employs resource-efficient tracing loops that wait for model warmup and capture data across a single training step.
4Optimized Tensor Handling
To manage large-scale model memory, ModuleTracer implements shape-only tracing, logging only tensor dimensions for synthetic input reconstruction. If an input tensor is shared across multiple modules, the system retains a single copy and uses reference links. These mechanisms significantly reduce GPU compute overhead while ensuring accurate module state reconstruction during offline execution.
5Centralized Module Repository
Extracted module states are persisted within a queryable backend built on the Weights & Biases Registry. Collections are indexed using MD5 hashes of model and module names, supporting attribute-based search. Each entry contains a JSON metadata file for execution context and a binary state payload serialized via torch.save for precise module restoration.
6Dual Consumption Interfaces
ModuleLens exposes data through two distinct surfaces: an interactive Jupyter extension for human engineers and specialized agent skills for LLMs. These interfaces allow agents like Claude Code or Cursor to autonomously discover, benchmark, and profile modules. This capability enables automated testing and validation without requiring manual intervention from the engineering team.
7Internal Tooling Applications
The system powers internal tools such as the Training Analyzer, which visualizes latency and memory usage across the module tree. Additionally, KernelAgent leverages captured data to generate and validate high-performance GPU kernels. Scaling Law ROI tools also utilize ModuleLens to track FLOPS and activation memory, helping teams assess infrastructure costs against model scaling gains.
8Future Agentic Optimization
The project aims to foster adoption across machine learning teams while expanding agent skills for automated Triton and CUDA kernel synthesis. By providing a reproducible execution substrate, ModuleLens enables recursive self-improvement. Agents can formulate experiments and retain successful changes, laying the groundwork for systems that continuously optimize both model quality and infrastructure strategies.