Introducing ModuleLens: Zooming into PyTorch Modules for Human and Agentic Engineering

Lida Li, Sr. Staff Software Engineer, Product ML Infra Introduction 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: Local Machine Hardware Limits : Because production models require multi-GPU setups, initializing a full model locally on a single-gpu workstation is virtually impossible. Input Data Dependencies : When an engineer modifies code inside a specific downstream module, getting realistic input data is a major hurdle. Traditionally, obtaining that input requires initializing the entire multi-GPU pipeline from scratch and executing every preceding module in the network — wasting immense amounts of time and computation resources. Slow Iteration Feedback Loops : Engineers face prolonged iteration loops when attempting to verify sub-module changes, as every code edit requires re-executing or re-initializing heavy distributed trainer setups. 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. Architecture 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 ModuleTracer is designed to capture and record all essential data — including input and output tensors, model parameter weights, and initialization arguments. This comprehensive tracking …