ModuleLensのご紹介:人間とエージェントによるエンジ... ノート

ModuleLensのご紹介:人間とエージェントによるエンジニアリングのためのPyTorchモジュールへのズームイン

現代のレコメンデーションシステムは、その大規模かつ分散された性質から、テストとデバッグにおいて重大な課題に直面しています。従来の「モデルファースト」アプローチは、ハードウェアの制約、入力データの依存関係、そして遅いイテレーションフィードバックループにつながります。ModuleLensは「モジュールファースト」パラダイムを導入し、エンジニアが個々のモジュールをローカルで分離、再構築、プロファイリングできるようにします。これにより、開発者のワークフローが大幅に加速され、自律的なAIエージェントの最適化が促進されます。ModuleLensアーキテクチャは、ModuleTracer、ModuleRepo、およびConsumption Layerで構成されます。ModuleTracerは、PyTorchのフォワードフックを介して、テンソル、パラメータ、初期化引数を含む重要なランタイムデータを非侵襲的にキャプチャします。これは、大規模なテンソルを効率的に処理し、モジュールの初期化をインメモリで追跡します。ModuleRepoは、抽出されたモジュール状態をWeights & Biases (WandB) Registryに一元化し、コレクションインデックスとメタデータおよびバイナリペイロードを含むバージョン管理されたアーティファクトによって整理します。Consumption Layerは、人間のエンジニア向けにインタラクティブなJupyter拡張機能とIPython CLIを、LLMエージェント向けにAgent Skillsを提供します。これにより、デバッグ、探索、および自律的なモジュール分析と最適化が可能になります。ModuleLensは、ボトルネック特定のためのTraining Analyzer、GPUカーネル生成と検証のためのKernelAgent、インフラコスト評価のためのScaling Law ROI Toolsなど、さまざまな社内ツールに採用されています。将来の計画には、コアシステムの安定化、Triton/CUDAカーネル合成などの高度な最適化のためのAgent Skillsの強化、および自律的なトレーニングとサービングの最適化の実現が含まれます。ModuleLensは、人間のエンジニアリング速度を向上させ、モデルの自動リサーチと再帰的な自己改善のための実行基盤を提供することを目指しています。分離されたモジュールとパフォーマンスシグナルへの再現可能なアクセスを提供することにより、ModuleLensはエージェントが実験し、コンポーネントを変更し、ベンチマークを実行し、成功した変更を保持することを可能にします。このプロジェクトは、Product ML Infra、ATG、HomeFeed、およびAds rankingチームの多数の貢献者によるリーダーシップ、設計に関する洞察、テスト、採用、およびコード開発に感謝しています。
CdXz5zHNQW_NEug4e7mVW.png