چک‌پوینت تکه‌تکه‌شده و ادامه‌ی اتمیک در آموزش توزیع‌شده

در آموزش مدل‌های زبانی حجیم روی چندین کارت گرافیک، وقوع خطا و متوقف شدن کار بخشی اجتناب‌ناپذیر از چرخه‌ی آموزش است. اگر از روش‌های سنتی برای ذخیره‌ی چک‌پوینت (Checkpoint) استفاده کنید، تمام پارامترها و وضعیت بهینه‌ساز باید روی یک گره متمرکز شوند و سپس روی دیسک قرار بگیرند. این کار علاوه‌بر هدر دادن ساعت‌ها از زمان باارزش محاسباتی، کل کلاستر را معطل درگاه شبکه‌ی یک کارت گرافیک نگه می‌دارد.

راهکار استاندارد برای پشت سر گذاشتن این گلوگاه، استفاده از چک‌پوینت‌های تکه‌تکه‌شده (Sharded Checkpoints) و الگوی نوشتن اتمیک (Atomic Write) است. با این شیوه، هر کارت یا رتبه (Rank) شارد داده‌های خود را به‌صورت کاملاً موازی در دیسک می‌نویسد و یک فایل مانیفست (Manifest) وظیفه‌ی ردیابی و مدیریت مالکیت این داده‌ها را بر عهده می‌گیرد.

در این درس یاد می‌گیرید چطور سیستم چک‌پوینت توزیع‌شده را از صفر بسازید، با نوشتن اتمیک در فایل‌های موقت جلوی خراب شدن داده‌ها در اثر کرش‌های ناگهانی را بگیرید و کدهایی مقاوم در برابر خطاهای رایج مثل تغییر تعداد گره‌ها یا ناقص ماندن فایل‌ها پیاده کنید.

در فرایند آموزش یک مدل 70 میلیارد پارامتری (70B)، خرابی گره‌ها هر چند ساعت یک‌بار کار را متوقف می‌کند. فرمت چک‌پوینت (checkpoint) تعیین می‌کند که 30 دقیقه از کار را از دست می‌دهید یا 30 ساعت را. چک‌پوینت تکه‌تکه‌شده (sharded checkpoint)، شارد مربوط به هر رتبه (rank) را به‌صورت موازی می‌نویسد و مالکیت هر بخش را در یک مانیفست (manifest) ثبت می‌کند. هنگام ادامه‌ی کار (resume)، شاردِ هر رتبه از فایل اختصاصی خودش بارگذاری می‌شود، وضعیت را روی همان تعداد فرایند (world size) بازسازی می‌کند و بهینه‌ساز طوری گام برمی‌دارد که انگار هیچ اتفاقی نیفتاده است. نوشتن اتمیک (atomic write) مانع از آن می‌شود که یک چک‌پوینت نیمه‌کاره اجرای بعدی را خراب کند.

اهداف یادگیری

• ذخیره‌ی یک چک‌پوینت چندرتبه‌ای در قالب فایل‌های شارد به‌ازای هر رتبه، به‌همراه مانیفستی که مشخص می‌کند کدام رتبه مالک کدام بخش است.
• استفاده از الگوی نوشتن اتمیک (نوشتن در یک مسیر موقت و سپس تغییر نام آن) تا قطعی یا کرش در میانه‌ی نوشتن، هرگز چک‌پوینت نیمه‌کاره تولید نکند.
• ادامه‌ی کار از روی مانیفست و اعتبارسنجی انطباق بایت‌به‌بایت داده‌ها، هم برای پارامترهای fp16 و هم برای وضعیت بهینه‌ساز ZeRO در تک‌تک رتبه‌ها.
• محافظت از اسکیما (schema)ی مانیفست در برابر سه حالت شکست (failure modes): تغییر world size، عدم تطابق تعداد شاردها، و نوشتن ناقص.

صورت مسئله

یک چک‌پوینت معمولی، همه‌ی پارامترها و وضعیت بهینه‌ساز را در رتبه‌ی 0 جمع‌آوری (gather) می‌کند و سپس همه‌چیز را در یک تک‌فایل می‌نویسد. برای یک مدل 70B، این یعنی عبور 1.1 ترابایت داده‌ی وضعیت از درگاه شبکه‌ی یک رتبه‌ی واحد. این نوشتن، بقیه‌ی رتبه‌ها را معطل می‌کند؛ چون باید منتظر پایان مرحله‌ی جمع‌آوری بمانند و بیکار می‌نشینند. در این حالت، پهنای باند ورودی/خروجی (IO) برابر با کندترین اتصال شبکه‌ی یک GPU است، نه مجموع پهنای باند کل کلاستر. روی یک کلاستر واقعی، مرحله‌ی «جمع‌آوری و سپس نوشتن» می‌تواند بیش از کل یک ساعت آموزش قبلی طول بکشد؛ این یعنی خروجی کل کار به کمتر از یک چک‌پوینت در هر روز کاری کاهش می‌یابد.

چک‌پوینت‌های تکه‌تکه‌شده این الگو را دگرگون می‌کنند: هر رتبه شارد خود را به‌صورت موازی در فایل اختصاصی خودش می‌نویسد. مانیفست ثبت می‌کند که کدام رتبه مالک کدام شارد بوده است تا فرایند resume بتواند هر شارد را دقیقاً سر جای اولش برگرداند. مجموع پهنای باند نوشتن، همگام با مقیاس کلاستر افزایش می‌یابد. یک چک‌پوینت 1 ترابایتی که نوشتن آن از طریق یک رتبه 4 ساعت زمان می‌برد، با 64 رتبه تنها 4 دقیقه طول می‌کشد. علاوه‌بر این، مانیفست قراردادی روشن برای شناسایی موارد ناسازگار در زمان بازیابی ارائه می‌دهد: تغییر world size و موارد نوشتن ناقص قابل تشخیص هستند، و مسیر بارگذاری داده می‌تواند به‌جای استفاده‌ی خاموش از داده‌های نامعتبر، با توقف با صدای رسا (fail loud) خطای خود را اعلام کند.

مفاهیم کلیدی

اسکیمای مانیفست

{
  "world_size": 4,
  "step": 1234,
  "wall_clock_seconds": 4521,
  "shards": [
    {"rank": 0, "path": "rank0.bin", "sha256": "...", "param_shard_offset": 0, "param_shard_numel": 65536},
    {"rank": 1, "path": "rank1.bin", "sha256": "...", "param_shard_offset": 65536, "param_shard_numel": 65536}
  ],
  "schema_version": 1
}

سه فیلد در این اسکیما نقش حیاتی دارند: world_size باعث می‌شود بازیابی روی تعداد متفاوتی از فرایندها به‌جای ایجاد خرابی خاموش، بلافاصله با توقف با صدای رسا روبه‌رو شود. فیلد sha256 به‌ازای هر شارد، موارد نوشتن ناقص یا خراب‌شده را تشخیص می‌دهد. فیلدهای param_shard_offset و param_shard_numel برای هر شارد به لودر اجازه می‌دهند تنسور (tensor) مسطح پارامترها را در موقعیت درست خود بازسازی کند.

نوشتن اتمیک

الگوی استاندارد به این صورت است: هر شارد را در فایلی با پسوند <name>.tmp می‌نویسید، مانیفست را در manifest.json.tmp ذخیره می‌کنید، روی تک‌تک آن‌ها fsync می‌زنید، نام فایل‌های شارد را تغییر می‌دهید و در آخرین مرحله نام فایل مانیفست را عوض می‌کنید. یک عملیات تغییر نام در استاندارد POSIX روی یک سیستم‌فایل یکسان، اتمیک است: یعنی یا فایل جدید به‌طور کامل وجود دارد یا فایل قبلی. مجموعه‌ای از تغییر نام‌ها یک تراکنش واحد به‌شمار نمی‌روند؛ بنابراین تغییر نام مانیفست همان نقطه‌ی قطعی‌شدن (commit) است و لودر هش همه‌ی شاردها را با مانیفست مقایسه و بررسی می‌کند. هر چک‌پوینت را داخل دایرکتوری گام اختصاصی خودش ذخیره کنید؛ در این صورت، اگر پیش از تغییر نام مانیفست کرشی رخ دهد، چک‌پوینت قبلی همچنان معتبر و فعال باقی می‌ماند. استفاده‌ی مجدد از یک دایرکتوری مشترک این تضمین را از بین می‌برد؛ چرا که تغییر نام شاردها پیش از رسیدن مانیفست جدید، شاردهای قدیمی را بازنویسی می‌کند. بدون نوشتن اتمیک، یک کرش ناگهانی می‌تواند شاردی ناقص به‌جا بگذارد در حالی که مانیفست موجود به آن اشاره می‌کند؛ در نتیجه، در زمان بازیابی، وضعیت بهینه‌ساز تخریب می‌شود.

سه حالت شکستی که اسکیما باید در برابرشان ایمن باشد

هر یک از این راهکارها مانع بارگذاری داده‌های نامعتبر در همان ابتدا می‌شوند؛ در غیر این صورت با خطای خاموشی روبه‌رو می‌شوید که شاید 100 گام بعد، وقتی تابع زیان (loss) مقدارش NaN شد، خودش را نشان دهد.

چرا فایل‌های مجزا به‌ازای هر رتبه، نه یک تک‌فایل بزرگ؟

نوشتن هم‌زمان در یک فایل با استفاده از O_APPEND در POSIX برای نوشتن‌های بایت‌محور کار می‌کند؛ اما در عمل، آفست‌های داخل یک شارد اسپن‌هایی (span) در مقیاس چند مگابایت را پوشش می‌دهند و قفل‌گذاری روی فایل سربار غالبی ایجاد می‌کند. فایل‌های مجزا برای هر رتبه هیچ رقابتی با هم ندارند و اگر سیستم‌فایل زیرساختی موازی باشد (مثل Lustre یا GPFS)، از تکنیک توزیع بلوک‌ها (striping) بیشترین بهره را می‌برند. به همین دلیل تمام استک‌های عملیاتی (DeepSpeed، FSDP و NeMo) از فایل‌های مجزا به‌ازای هر رتبه استفاده می‌کنند.

پیاده‌سازی از صفر

فایل code/main.py این بخش‌ها را پیاده‌سازی می‌کند:

• دیتاکلاس ShardManifest با اسکیمای معرفی‌شده در بالا به‌همراه توابع to_json/from_json.
• تابع save_sharded(state_dict_per_rank, dir, step) که وضعیت باینری هر رتبه را با استفاده از الگوی اتمیکِ «نوشتن موقت و سپس تغییر نام» در فایل اختصاصی خودش می‌نویسد و سپس مانیفست را ذخیره می‌کند.
• تابع load_sharded(dir, expected_world_size) که مانیفست را می‌خواند، sha256 هر شارد را بررسی می‌کند و دیکشنری وضعیت پارامترها (state dict) را به‌ازای هر رتبه برمی‌گرداند.
• یک آزمون رفت‌وبرگشت (round-trip): ساخت وضعیت به‌ازای هر رتبه، ذخیره، بارگذاری، و بررسی انطباق دقیق بایت‌به‌بایت.

اجرای کد:

python3 code/main.py

خروجی: 4 فایل شارد به‌همراه فایل مانیفست ذخیره می‌شوند و سپس با اعتبارسنجی انطباق بایت‌به‌بایت دوباره بارگذاری می‌گردند.

الگوهای کاربردی در محیط‌های عملیاتی

سه الگو سطح پایداری چک‌پوینت را برای استفاده در محیط عملیاتی (پروداکشن) تضمین می‌کنند.

نوشتن غیرهم‌زمان (Async write). استک‌های عملیاتی نوشتن چک‌پوینت را روی ترد (thread) یا فرایند (process) مجزایی ارسال می‌کنند تا آموزش بدون وقفه ادامه یابد. مانع هماهنگی (barrier) در چک‌پوینت بعدی قرار دارد: تا زمانی که ذخیره‌ی قبلی کامل نشده باشد، نباید ذخیره‌ی بعدی شروع شود. فلگ async_io در DeepSpeed دقیقاً همین کار را انجام می‌دهد. در این درس نوشتن را هم‌زمان نگه داشتیم تا مراحل به‌وضوح دیده شوند.

اول دیسک محلی سریع، سپس آپلود غیرهم‌زمان. ابتدا روی NVMe محلی (بسیار سریع) بنویسید و سپس به‌شکل غیرهم‌زمان داده‌ها را روی S3 یا GCS آپلود کنید. این الگوی دومرحله‌ای، سرعت چک‌پوینت‌گیری را روی کلاستر برای ادامه‌ی سریع کار بالا نگه می‌دارد و هم‌زمان نسخه‌ای پایدار را نیز خارج از کلاستر آرشیو می‌کند. مانیفست اصلی مسیر محلی را نگه می‌دارد؛ مانیفست آپلود مسیر مخزن راه‌دور (remote) را ثبت می‌کند.

چرخش چک‌پوینت‌ها (Rotation). اجراهای عملیاتی معمولاً آخرین K چک‌پوینت (معمولاً 3 تا 5 مورد) را نگه می‌دارند و قدیمی‌ترین‌ها را پاک می‌کنند. بدون چرخش، دیسک در میانه‌ی فرایند پر می‌شود و چک‌پوینت بعدی شکست می‌خورد. با چرخش، ذخیره‌سازی بعدی ابتدا قدیمی‌ترین فایل را حذف می‌کند و فضای دیسک را آزاد می‌سازد.

استفاده از ابزارهای واقعی

الگوهای موجود در محیط‌های عملیاتی:

• چک‌پوینت‌گیری در DeepSpeed. دستور deepspeed.save_checkpoint(tag=step) فایل‌های هر رتبه و یک فایل latest را ذخیره می‌کند که به تگ فعال اشاره دارد.
• چک‌پوینت‌گیری با FSDP در PyTorch. ماژول torch.distributed.checkpoint وضعیت شاردشده را به‌کمک یک Planner ذخیره می‌کند؛ این پلنر چیدمان داده‌ها را برای هر رتبه تعیین می‌کند.
• کتابخانه‌ی NeMo. ابزارهای DeepSpeed و FSDP را در رابط کاربری یکپارچه‌ی save_to_checkpoint قرار می‌دهد و متادیتای موردنیاز را اضافه می‌کند.

خروجی این درس

در درس 81، یک چک‌پوینت شاردشده از اجرای سرتاسری DDP+ZeRO ذخیره می‌کنیم و آن را روی همان world size دوباره بارگذاری می‌کنیم تا مطمئن شویم قرارداد بازیابی (resume) بدون نقص کار می‌کند.

تمرین‌ها

1. نوشتن غیرهم‌زمان را اضافه کنید: ذخیره‌سازی را در یک ترد مجزا آغاز کنید و بگذارید آموزش ادامه یابد. چک‌پوینت بعدی را تا زمان اتمام چک‌پوینت قبلی متوقف نگه دارید.
2. قابلیت چرخش last_5_steps را اضافه کنید: 5 چک‌پوینت آخر را نگه دارید و پیش از نوشتن چک‌پوینت جدید، قدیمی‌ترین را پاک کنید.
3. یک مسیر بررسی سریع با CRC اضافه کنید تا بازخوانی‌های داخل حلقه سریع‌تر انجام شوند (جایی که در فرایند چرخش، بدون نیاز به محاسبه‌ی کامل sha256، یک چک‌پوینت به‌عنوان چک‌پوینت فعالِ جدید تعیین می‌شود).
4. امکان بارگذاری در ابعاد متفاوت (cross-world-size) را پیاده کنید: تعادل مجدد شاردها از N=4 به N=8 را با خواندن مانیفست، ادغام و سپس شاردبندی مجدد انجام دهید.
5. قابلیت آپلود به یک شبیه‌ساز S3 (یک دایرکتوری دوم) را اضافه کنید و مانیفست آپلود را بنویسید. از منطق سیاست ذخیره‌سازی دومرحله‌ای دفاع کنید.

اصطلاحات کلیدی

برای مطالعه‌ی بیشتر

• چک‌پوینت‌گیری در DeepSpeed
• مستندات torch.distributed.checkpoint در PyTorch
• اتمیک بودن Rename در POSIX
• فاز 19، درس 78 - وضعیت ZeRO که این سیستم چک‌پوینت برای ذخیره‌ی آن طراحی شده است
• فاز 19، درس 81 - پیاده‌سازی سرتاسری که وضعیت ذخیره‌شده را در چرخه‌ی رفت‌وبرگشت آزمایش می‌کند