KOEN
Read Original
#AI

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

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

Summary

Goodbye Multi-GPU Debug Bottlenecks: How Pinterest's ModuleLens Profiles PyTorch Modules in Seconds

Shifting from Model-First to Module-First architecture to empower both human engineers and autonomous AI agents.

Debugging large distributed ML models often requires re-running entire multi-GPU pipelines just to test a single submodule edit. ModuleLens solves this by capturing runtime weights and tensor I/O into isolated, reproducible modules that can be reloaded locally in seconds. Beyond drastically reducing developer iteration cycles, it serves as an execution substrate for AI agents like Claude Code to autonomously refactor code and generate custom GPU kernels.

  • 01Shift from legacy Model-First to an isolated Module-First engineering paradigm
  • 02Non-intrusive PyTorch forward hooks and import patching to capture weights, tensors, and init args without code changes
  • 03Versioned module storage using WandB Registry with binary payloads and JSON metadata
  • 04Support for Jupyter extensions, IPython CLI, and LLM Agent Skills (e.g., Claude Code, Cursor)
  • 05Integration with KernelAgent for LLM-driven custom Triton/CUDA kernel synthesis and verification

RECOMMENDATION

Highly recommended for ML Infrastructure teams facing slow iteration loops in multi-GPU distributed training, as well as teams looking to integrate AI agents into automated kernel synthesis and model optimization workflows.

The Problem

Modern large-scale recommendation systems rely on Distributed Model Parallelism (DMP) across multiple GPUs, making local single-GPU execution impossible. Modifying downstream sub-modules required re-running the entire multi-GPU pipeline to capture realistic input data, resulting in severely slow iteration feedback loops and wasted compute.

The Solution

Pinterest developed ModuleLens, shifting from a legacy 'Model-First' approach to an isolated 'Module-First' framework. Using non-intrusive PyTorch forward hooks and in-memory init tracking, ModuleLens captures module parameter weights and tensor I/O into a WandB Registry-backed ModuleRepo, enabling human engineers (CLI/Jupyter) and AI agents (Claude Code, Cursor) to reconstruct modules locally in seconds.

The Result

Engineers can isolate and profile targeted PyTorch modules locally within seconds, drastically accelerating iteration loops. Furthermore, ModuleLens powers internal tools like Training Analyzer and KernelAgent, providing an execution substrate for AI agents to autonomously synthesize, benchmark, and optimize GPU kernels.

Trade-off

Capturing module traces introduces runtime GPU overhead, which ModuleLens mitigates by running early-stop callbacks after model warmup. For massive tensors, it relies on shape-only logging, requiring synthetic tensor construction during restoration which may not reflect real input distributions.

03

Key Concepts

Concept · 01

Module-First Paradigm

An engineering approach that isolates individual sub-modules from large neural network architectures to enable independent profiling and execution without initializing the full model. This eliminates multi-GPU dependencies for local testing.

  • Used in ModuleLens to allow local module reconstruction and profiling in seconds without GPU clusters.
Concept · 02

ModuleTracer

A runtime tracer that non-intrusively captures parameter weights, I/O tensors, and initialization arguments from PyTorch modules using forward hooks and import patching. It minimizes GPU overhead using early-stop callbacks after warmup.

  • Leverages shape-only logging for large tensors and deduplicates shared input tensors using reference links.
Concept · 03

KernelAgent

An LLM-driven agent framework designed to synthesize high-performance GPU kernels (e.g., Triton/CUDA). It utilizes captured ModuleLens inputs and baseline outputs to benchmark and validate kernel speed and accuracy.

  • Leverages captured module data to test, profile, and validate autonomously generated GPU kernels.