‏AutoSP: آموزش مدل‌های زبانی با کانتکست‌های بسیار بلند، بدون دردسر

آموزش مدل‌های زبانی بزرگ (LLM) با کانتکست‌های خیلی بلند یه چالش جدی مهندسیه. وقتی تعداد توکن‌ها از ۱۰۰ هزار رد می‌شه، حتی با تکنیک‌های رایجی مثل ZeRO یا FSDP هم خطای out-of-memory (OOM) گریبان‌گیر می‌شه. یه راه‌حل معمول، Sequence Parallelism (SP) هست؛ یعنی توکن‌های ورودی رو بین چند GPU تقسیم کنیم تا آموزش روی کانتکست‌های بلندتر ممکن بشه.

مشکل اینه که پیاده‌سازی SP به‌صورت دستی فوق‌العاده سخته. باید توی کتابخونه‌هایی مثل DeepSpeed یا HuggingFace تغییرات عمیقی بدیم، توکن‌ها و اکتیویشن‌های میانی رو تقسیم کنیم، عملیات ارتباطی بین GPUها رو جا بندازیم، و همه این‌ها رو هم برای forward pass و هم backward pass انجام بدیم. نتیجه اینه که محققان وقت زیادی رو صرف مهندسی سیستم می‌کنن، نه پژوهش واقعی.

AutoSP برای حل این مشکل طراحی شده. این ابزار یه کامپایلر اتوماتیکه که کد معمولی تک-GPU ترنسفورمر رو به طور خودکار به کد چند-GPU با Sequence Parallelism تبدیل می‌کنه. AutoSP درون DeepCompile (اکوسیستم کامپایلر DeepSpeed) پیاده‌سازی شده و با ZeRO نیز سازگاره. کافیه چند تنظیم ساده توی کانفیگ DeepSpeed اضافه بشه:

config = {
    "compile": {
        "deepcompile": True,
        "passes": ["autosp"]
    },
    "sequence_parallel_size": 4,
    "zero_optimization": {
        "stage": 1,
    },
}
model, _, _ = deepspeed.initialize(config=config, model=model)
model.compile(compile_kwargs={"dynamic": True})

بعد از این، کاربر فقط باید از تابع کمکی prepare_auto_sp_inputs برای تگ‌گذاری سبک روی توکن‌های ورودی، attention mask و position id استفاده کنه. بقیه کار رو کامپایلر AutoSP انجام می‌ده و به‌طور خودکار Sequence Parallelism رو فعال می‌کنه.

استراتژی SP که AutoSP پیاده‌سازی می‌کنه DeepSpeed-Ulysses هست. دلیل انتخاب این استراتژی اینه که سربار ارتباطی‌اش با افزایش تعداد GPU ثابت می‌مونه؛ چیزی که در شبکه‌های NVLink یا fat-tree خیلی مهمه. البته یه محدودیت هم داره: مقیاس SP نمی‌تونه از تعداد attention head های مدل بیشتر بشه (مثلاً در مدل‌های ۷-۸ میلیارد پارامتری، حداکثر ۳۲ تاست).

AutoSP علاوه بر SP، یه استراتژی سفارشی Activation Checkpointing (AC) هم داره به اسم Sequence-aware AC یا SAC. در حالت عادی PyTorch 2.0 یه روش AC مبتنی بر max-flow min-cut ارائه می‌ده که برای آموزش کانتکست‌های بلند خیلی محافظه‌کارانه‌ست. SAC این مشکل رو حل می‌کنه و با توجه به دینامیک خاص FLOP در کانتکست‌های طولانی، اکتیویشن‌های ارزون‌تر رو آزاد می‌کنه و در backward pass دوباره محاسبه‌شون می‌کنه. این کار کمی throughput رو کاهش می‌ده، ولی بدون اون، آموزش روی کانتکست‌های بلند اصلاً ممکن نیست.

نتایج بنچمارک روی مدل‌های مختلف Llama 3.1 روی نود ۸ تایی A100-80GB نشون می‌ده که AutoSP در مقایسه با بیس‌لاین‌های دستی مثل RingFlashAttention، DeepSpeed-Ulysses و ZeRO-3، هم حداکثر طول کانتکست قابل‌آموزش رو به شکل قابل‌توجهی بالا می‌بره و هم سربار زمانی ناچیزیه.

البته AutoSP دو محدودیت مهم داره. اول اینکه مدل باید به‌صورت یه artifact کامپایلی واحد کامپایل بشه؛ کامپایل جداگانه‌ی توابع مختلف و چسبوندن‌شون به هم کار نمی‌کنه. دوم اینکه هیچ graph break ای توی مدل قابل‌تحمل نیست، چون تحلیل و انتشار اطلاعات رو پیچیده می‌کنه.

نکات کلیدی:

• AutoSP کد آموزش معمولی تک-GPU ترنسفورمر رو به‌طور خودکار به کد چند-GPU با Sequence Parallelism تبدیل می‌کنه
• فعال‌سازی فقط با تغییر چند خط در کانفیگ DeepSpeed انجام می‌شه
• از استراتژی DeepSpeed-Ulysses استفاده می‌کنه که سربار ارتباطی ثابتی داره
• استراتژی SAC سفارشی برای Activation Checkpointing در کانتکست‌های بلند طراحی شده
• با ZeRO Stage 1 به‌صورت کامل سازگاره
• محدودیت: مدل باید به‌صورت یه واحد کامپایلی واحد کامپایل بشه و graph break مجاز نیست
• بنچمارک روی Llama 3.1 روی A100 نشون‌دهنده‌ی افزایش قابل‌توجه حداکثر طول کانتکست با سربار زمانی کم هست