Introducing ModuleLens: Zooming into PyTorch Modules for Human and Agentic Engineering
Modern recommendation systems face significant challenges in testing and debugging due to their large scale and distributed nature. Traditional "Model-First" approaches lead to hardware limits, input data dependencies, and slow iteration feedback loops. ModuleLens introduces a "Module-First" paradigm, enabling engineers to isolate, reconstruct, and profile individual modules locally. This significantly accelerates developer workflows and facilitates autonomous AI agent optimization.The ModuleLens architecture comprises ModuleTracer, ModuleRepo, and a Consumption Layer. ModuleTracer captures essential runtime data, including tensors, parameters, and initialization arguments, non-intrusively via PyTorch forward hooks. It efficiently handles large tensors and tracks module initialization in-memory.ModuleRepo centralizes extracted module states in a Weights & Biases (WandB) Registry, organizing them by collection indexing and versioned artifacts containing metadata and binary payloads. The Consumption Layer offers interactive Jupyter extensions and IPython CLI for human engineers, and Agent Skills for LLM agents.This enables debugging, exploration, and autonomous module analysis and optimization. ModuleLens has been adopted for various internal tools, including Training Analyzer for bottleneck identification, KernelAgent for GPU kernel generation and validation, and Scaling Law ROI Tools for assessing infrastructure costs.Future plans involve stabilizing the core system, enhancing Agent Skills for advanced optimization like Triton/CUDA kernel synthesis, and enabling autonomous training and serving optimization. ModuleLens aims to boost human engineering velocity and provide an execution substrate for model auto-research and recursive self-improvement.By offering reproducible access to isolated modules and performance signals, ModuleLens empowers agents to experiment, modify components, benchmark, and retain successful changes. The project acknowledges numerous contributors from Product ML Infra, ATG, HomeFeed, and Ads ranking teams for their leadership, design insights, testing, adoption, and code development.