It eliminates the hardware bottlenecks of debugging massive distributed models by enabling local, isolated module execution. This drastically reduces iteration times and provides a structured environment for AI agents to autonomously optimize PyTorch code and GPU kernels.
Lida Li, Sr. Staff Software Engineer, Product ML Infra
Modern recommendation systems, including those used for ads ranking and content discovery, have grown substantially in scale and complexity. This growth stems from both increasingly large model architectures. The resulting memory footprint often exceeds the capacity of a single GPU, making it necessary to shard production models across multiple GPUs using distributed training techniques such as Distributed Model Parallelism (DMP).
This evolution brings challenges to the everyday engineer work on testing and debugging the model code. The bottleneck includes:
To eliminate these bottlenecks, we build ModuleLens, which pivots from a legacy “Model-First” paradigm to an isolated “Module-First” engineering framework. The following table contrasts these two operational paradigms across key engineering metrics:
Under the module-first paradigm, engineers can isolate any given module, avoid multi-GPU model initialization entirely, and reconstruct and profile it locally within seconds. In addition to accelerating developer workflows, this modular foundation offers a fast, structured execution environment that empowers autonomous AI agents to analyze, refactor, and optimize PyTorch code independently.
The architecture of ModuleLens is organized into three primary components: ModuleTracer for capturing runtime data, Module Repo for central storage, and a Consumption Layer serving human interfaces and AI-based tooling. The following diagram depicts this structure.
ModuleTracer is designed to capture and record all essential data — including input and output tensors, model parameter weights, and initialization arguments. This comprehensive tracking allows any target PyTorch module to be reconstructed seamlessly within the consumption layer, enabling offline execution with captured inputs and validation against expected outputs.
To carry out trace capture, ModuleTracer relies on several key mechanisms:
Extracted module state is persisted within a queryable backend built on top of the Weights & Biases (WandB) Registry. All captured modules are stored into one registry and organized around Collection and Versioned Artifacts:
Collection Indexing: Target collections are indexed using an MD5 hash of the model and module names: hash(model_name, module_name). Additionally, top-level trace executions log overall run history to special collections structured as registered_model_xxx, indexed directly by model name. The registry supports both attribute-based and tag-based search filtering.
Versioned Artifact Payload: Every captured Module produces a distinct version entry containing two standardized artifacts:
ModuleLens exposes its capabilities through two tailored surfaces to consume the traced data:
Human Engineers: Interactive Jupyter Extension and IPython CLI for hands-on debugging, exploration, and manual profiling.
AI Agents: Agent Skills that allow LLM agents (e.g., Claude Code, Cursor) to autonomously discover, rebuild, benchmark, and profile modules without human intervention.
The consume layer provides the following capabilities:
The screenshot below illustrates the command-line interface available to users.
Ever since its deployment across our Ads and Core models, ModuleLens has also been adopted for various internal tools built by our engineering teams.
Evaluates module-level metrics across different architectures. Through hierarchical visual representations of latency, FLOPs, and memory usage across the module tree, engineers can rapidly identify training bottlenecks.
An LLM-driven agent designed to generate high-performance GPU kernels. It leverages captured ModuleLens data to benchmark, profile, test, and validate the accuracy and speed of generated kernels.
Evaluates the return on investment (ROI) for model scaling. By utilizing ModuleLens to track training and serving FLOPS, activation and parameter memory usage, and execution throughput, it helps assess infrastructure costs against model scaling gains.
Our immediate priority is stabilizing the core system and fostering adoption across machine learning teams. Concurrently, we are enriching our Agent Skills repository to furnish AI coding agents with advanced optimization capabilities, including automated Triton/CUDA kernel synthesis as well as autonomous training and serving optimization workflows. By building reliable, execution-friendly interfaces for both developers and autonomous agents, ModuleLens optimizes human engineering velocity while laying the groundwork for scalable, fully automated performance optimization.
Beyond optimization, ModuleLens can serve as an execution substrate for model autoresearch and RSI (recursive self-improvement). By giving agents reproducible access to isolated modules, captured inputs and outputs, and trustworthy performance signals, it enables them to formulate experiments, modify model components, benchmark results, and retain successful changes in a tight feedback loop. Over time, this foundation can support systems that continuously improve both model quality and the tools and strategies used to optimize the model.
This project benefited from the valuable contributions of many individuals. We express our gratitude to Archer Liu, leader of Product ML Infra (PMLI), for his leadership and support, Chen Yang (ATG) for insightful design discussions and feedback; Abhinav Naikawadi (HomeFeed) and Paulo Soares (Ads ranking) for leading testing, adoption, and advocacy efforts; and Jiasong Sun, Xinyi Zhang, Diane Li and Hongze Zhao (Product ML Infra) for their contributions to code development and review.
Introducing ModuleLens: Zooming into PyTorch Modules for Human and Agentic Engineering was originally published in Pinterest Engineering Blog on Medium, where people are continuing the conversation by highlighting and responding to this story.
Continue reading on the original blog to support the author
Read full articleThis approach shows how to collapse complex multi-stage recommendation pipelines into a single generative task. It proves that distilled 4B-parameter LLMs can provide production-grade performance for high-scale personalization, significantly improving user intent understanding and engagement.
Engineers often struggle to balance model complexity with real-time latency at scale. This 3-tower approach demonstrates a practical way to integrate rich cross-features and multi-task learning into high-throughput ranking systems without exceeding strict compute budgets.
Scaling embedding retrieval to billions of vectors is prohibitively expensive in RAM. Pinterest's approach demonstrates how quantization and SSD-based indexing (SPANN) can slash infrastructure costs by 30-40% while maintaining high recall and low latency for real-time discovery.
Serving VLMs at scale requires overcoming massive compute and memory bottlenecks not found in text-only LLMs. This architecture demonstrates how to optimize multimodal inference using disaggregated serving and specialized hardware to achieve better TCO and performance for visual discovery.
