KOEN
Read Original
#AI

ModuleLens 소개: 인간 및 AI 에이전트 엔지니어링을 위한 PyTorch 모듈 집중 정밀 분석

ModuleLens 소개: 인간 및 AI 에이전트 엔지니어링을 위한 PyTorch 모듈 집중 정밀 분석
01

Summary

멀티 GPU 모델 재실행은 그만! PyTorch 모듈만 쏙 떼어 로컬에서 초단위로 프로파일링하는 ModuleLens

핀터레스트가 구축한 모듈 중심 프레임워크로 엔지니어 생산성과 AI 에이전트 자율 최적화를 동시에 달성하는 방법

거대한 분산 학습 모델 개발 시 특정 하위 모듈 하나를 수정하기 위해 전체 파이프라인을 돌리는 것은 엄청난 시간과 자원 낭비입니다. ModuleLens는 PyTorch 런타임 상태를 모듈 단위로 캡처하여 분산 환경 없이도 로컬에서 단 수 초 만에 재현 및 분석할 수 있게 해줍니다. 이 프레임워크는 개발자의 디버깅 속도를 혁신할 뿐만 아니라 Claude Code, Cursor 같은 AI 에이전트가 자동 커널 생성 및 모듈 최적화를 수행하는 강력한 실행 기반이 됩니다.

  • 01기존 모델 중심(Model-First) 패러다임에서 모듈 중심(Module-First) 격리 프레임워크로의 전환
  • 02PyTorch Forward Hook과 __import__ 패칭을 이용해 코드 수정 없이 런타임 입출력 및 가중치 비침습적 캡처
  • 03WandB Registry 기반 ModuleRepo에 버전 관리된 메타데이터 및 torch.save 바이너리 저장
  • 04Jupyter CLI 및 Claude Code/Cursor 지원 AI Agent Skills를 통해 자율적 모듈 최적화 지원
  • 05KernelAgent를 활용한 자동 GPU 커널 생성 및 Accuracy/Latency 검증 체계 구축

RECOMMENDATION

대규모 분산 학습 환경에서 느린 피드백 루프로 고통받는 ML Infra 팀이나, AI 에이전트를 모듈 단위 코드 최적화 및 커널 생성에 도입하려는 엔지니어링 조직에 강력히 추천합니다.

The Problem

초대형 추천 및 광고 모델은 멀티 GPU 분산 학습 환경(DMP)을 필수적으로 요구하여 로컬 워크스페이스에서 전체 모델을 초기화하거나 실행하는 것이 불가능합니다. 특정 하위 모듈 코드를 수정할 때에도 전체 분산 파이프라인을 재실행해야 하므로 입력 데이터 확보가 어렵고 반복 실험 피드백 루프가 크게 지연되었습니다.

The Solution

핀터레스트는 기존 '모델 우선' 방식에서 벗어나 모듈을 독립적으로 격리하여 분석하는 'ModuleLens' 프레임워크를 개발했습니다. PyTorch Forward Hook과 메모리 내 초기화 추적을 사용해 모듈 입출력 텐서 및 가중치를 비침습적으로 캡처하고, WandB Registry 기반 저장소에 저장하여 개발자 CLI 및 AI 에이전트(Claude Code, Cursor 등)가 로컬에서 수 초 내 재구성할 수 있도록 구현했습니다.

The Result

엔지니어는 멀티 GPU 파이프라인 재실행 없이 수 초 만에 모듈을 로컬에서 복원 및 프로파일링하여 개발 속도를 획기적으로 개선했습니다. 또한 Training Analyzer, KernelAgent, Scaling Law ROI Tools와 연동되어 AI 에이전트가 로컬 모듈 단위로 고성능 GPU 커널을 자동 생성하고 정밀도를 검증하는 환경을 완성했습니다.

Trade-off

런타임 트레이싱에 따른 연산 오버헤드를 줄이기 위해 1회 스텝 캡처 후 조기 종료하는 방식을 취하며, 대형 텐서의 경우 형상 정보만 기록하는 shape-only 덤프를 적용하므로 실제 입력 데이터의 고유한 분포를 완전히 재현하는 데 제한이 있을 수 있습니다.

03

Key Concepts

Concept · 01

Module-First Paradigm

전체 모델을 초기화하지 않고 특정 하위 모듈만 격리하여 입출력 데이터와 가중치를 복원 및 실행하는 엔지니어링 접근 방식입니다. 분산 환경 의존성을 제거해 개발 및 프로파일링 루프를 수 초 단위로 단축시킵니다.

  • ModuleLens에서 멀티 GPU 초기화 없이 로컬 워크스페이스에서 수 초 내 모듈을 재구성하는 핵심 접근법으로 활용되었습니다.
Concept · 02

ModuleTracer

PyTorch 런타임 중 모듈의 입출력 텐서, 파라미터 가중치, 초기화 인자를 비침습적으로 캡처하는 추적기입니다. 1회성 스텝 캡처 및 조기 종료 기법으로 GPU 오버헤드를 최소화합니다.

  • PyTorch Forward Hook 및 in-memory __init__ tracking을 적용해 코드 변경 없이 모듈 상태를 캡처합니다.
Concept · 03

KernelAgent

LLM을 기반으로 고성능 GPU 커널(Triton/CUDA)을 자동 생성하는 AI 에이전트입니다. ModuleLens가 캡처한 실제 모듈 데이터를 바탕으로 생성된 커널의 정밀도와 속도를 검증합니다.

  • ModuleLens 캡처 데이터를 활용해 벤치마크, 프로파일링 및 연산 정확도 검증을 수행합니다.