کاهش مصرف رم در پای‌تورچ با الگوریتم کوله‌پشتی و روش Hirschberg

هنگام ساخت گراف محاسباتی در PyTorch، سیستم باید تصمیم بگیرد کدام عملیات ذخیره شوند و کدام‌یک دوباره محاسبه (recompute) گردند تا هم سرعت بالا برود و هم در مصرف حافظه صرفه‌جویی شود. این فرآیند از دیدگاه بهینه‌سازی به مسئله کلاسیک «کوله‌پشتی صفر و یک» (0/1 Knapsack) شباهت دارد.

پس‌زمینه و مسئله کوله‌پشتی در PyTorch

در مسئله کوله‌پشتی، شما ظرفیت محدودی (بودجه حافظه رم) دارید و مجموعه‌ای از آیتم‌ها (عملیات مدل) که هرکدام وزن (حافظه مورد نیاز) و ارزشی (صرفه‌جویی در زمان اجرا) دارند. هدف این است که با انتخاب بهینه آیتم‌ها، ارزش کل را حداکثر کنید بدون اینکه از بودجه بگذرید.

پیش‌فرض PyTorch استفاده از روش برنامه‌نویسی پویا یعنی dp_knapsack است. این الگوریتم یک جدول دوبعدی کامل با ابعاد (تعداد آیتم‌ها، حداکثر وزن) می‌سازد. در مدل‌های بزرگ با تعداد لایه‌ها و عملیات بالا، ساخت این جدول حافظه رم زیادی مصرف می‌کند و گاهی به ارور Out of Memory منجر می‌شود.

راه حل جدید: ترکیب پنجره لغزان و ترفند Hirschberg

برای حل این مشکل، رویکرد جدیدی به نام dp_knapsack_sliding_hirschberg ارائه شده است که دو ایده کلیدی را ترکیب می‌کند:

• ترفند پنجره لغزان (Sliding Window): به جای نگهداری کل جدول دوبعدی، تنها به ردیف قبلی و ردیف فعلی نیاز داریم.
• ترفند Hirschberg: یک رویکرد تقسیم و حل (Divide and Conquer) شبیه به الگوریتم Quicksort که نیاز به مرحله بازگشتی (Backtracking) برای پیدا کردن آیتم‌های انتخابی را از بین می‌برد.

با این روش، اندازه جدول از (تعداد آیتم‌ها × حداکثر وزن) به (۲ × حداکثر وزن) کاهش می‌یابد. به عنوان مثال، در سیستمی با ۶۴ گیگابایت رم، الگوریتم پیش‌فرض روی ۱۰۰ آیتم کرش می‌کند، در حالی که نسخه جدید به‌راحتی از پس ۲۰۰0 آیتم برمی‌آید و حدود ۳۷ درصد نیز سرعت اجرای بهتری دارد.

نحوه استفاده در کد

اگر از پکیج‌های پیش‌فرض PyTorch استفاده می‌کنید، نیازی به تغییر نیست زیرا dp_knapsack همچنان پیش‌فرض است. اما برای فعال‌سازی الگوریتم جدید می‌توانید این کد را اجرا کنید:

import torch
import torch._functorch.config as fconfig

fconfig.activation_memory_budget_solver = fconfig.dp_knapsack_sliding_hirschberg

دیدگاه برنامه‌نویسان در هکرنیوز

در بحث‌های هکرنیوز، توسعه‌دهندگان به نکات جالبی اشاره کردند:

• برخی کاربران پیشنهاد کردند برای نمونه‌های متوسط تا بزرگ می‌توان از الگوریتم‌های ترکیبی پیشرفته‌تر (مانند کدهای Pisinger) استفاده کرد که بسیار سریع‌تر هستند.
• همچنین بحث شد که اگر سرعت اهمیت بیشتری از دقت کامل داشته باشد، گزینه‌هایی مثل greedy_knapsack یا در صورت استفاده از SciPy استفاده از ilp_knapsack گزینه‌های جایگزین مناسبی هستند.

این تغییر ساده در کانفیگ پای‌تورچ برای مهندسانی که با مدل‌های سنگین و محدودیت حافظه رم در زمان کامپایل گراف مواجهند، اهمیت بالایی دارد.