ModuleLens: Accelerating ML Development and Autonomous Agent Engineering at Pinterest

The Challenge of Large-Scale ML Model Development

Modern recommendation systems, with their increasingly complex and large PyTorch model architectures, often exceed the memory capacity of a single GPU, necessitating distributed training. This scale introduces significant bottlenecks for engineers: inability to initialize full models locally, difficulty in obtaining realistic input data for specific modules, and slow iteration feedback loops due to heavy distributed setups. These issues hinder rapid experimentation and debugging. ModuleLens was developed to address these challenges by providing an isolated, module-first engineering framework.

Module-First Paradigm: Architecture and Components

ModuleLens pivots to a 'Module-First' approach, allowing engineers to isolate and reconstruct any given PyTorch module within seconds, without requiring a multi-GPU model initialization. This not only accelerates human developer workflows but also creates a structured environment for autonomous AI agents to analyze, refactor, and optimize PyTorch code. The system is designed with three core components:

• ModuleTracer: Captures runtime data including input/output tensors, model parameter weights, and initialization arguments, enabling offline module reconstruction. It uses non-intrusive PyTorch forward hooks, in-memory init tracking, and resource-efficient tracing loops. It also optimizes handling of large tensors by logging shapes only and deduplicating shared tensors.
• ModuleRepo: A central storage built on Weights & Biases (WandB) Registry, persisting extracted module states. It organizes modules into versioned artifacts with JSON metadata and binary state payloads (input tensors, output tensors, parameter state_dict). Collections are indexed by an MD5 hash of model and module names, supporting efficient search.
• Consumption Layer: Exposes ModuleLens capabilities through interactive Jupyter Extensions and IPython CLI for human engineers, and provides Agent Skills for LLM agents to autonomously discover, rebuild, benchmark, and profile modules.

Key Use Cases and Future Directions

• Training Analyzer: Visualizes module-level metrics (latency, FLOPs, memory) to identify training bottlenecks.
• KernelAgent: An LLM-driven agent that generates and validates high-performance GPU kernels using captured ModuleLens data for benchmarking.
• Scaling Law ROI Tools: Assesses infrastructure costs against model scaling gains by tracking FLOPS and memory usage.
• Autonomous Research and RSI: Future plans include enabling AI agents to formulate experiments, modify model components, benchmark results, and retain successful changes in a continuous self-improvement loop for models and optimization strategies.

The architecture highlights how capturing and making granular runtime data accessible can empower both human engineers and AI agents, leading to faster iteration, automated optimization, and the potential for recursive self-improvement in complex ML systems. This pattern of isolating components for easier development and testing is applicable across various domains, not just ML.