ذخیره و بازیابی چک‌پوینت در آموزش مدل‌های یادگیری عمیق

توقف ناگهانی فرایند آموزش یکی از پرهزینه‌ترین اتفاقات در چرخه‌ی یادگیری عمیق است. قطع برق، اتمام زمان تخصیص منابع در کلاستر، یا ریبوت‌های ناگهانی سیستم می‌تواند ساعت‌ها محاسبه و هزینه‌ی پردازشی سنگین را بر باد دهد؛ به‌ویژه زمانی که به دلیل محدودیت دسترسی به زیرساخت‌ها و کارت‌های گرافیک ابری، دسترسی مداوم به سخت‌افزار بسیار گران و باارزش تمام می‌شود.

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

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

توقف ناگهانی فرایند آموزش (interrupt)، اجراها را از بین می‌برد؛ چک‌پوینت‌ها ادامه‌ی کار را ممکن می‌سازند. مدل، بهینه‌ساز (optimizer)، اسکژولر (scheduler)، تاریخچه‌ی تابع زیان (loss)، شمارنده‌ی گام‌ها و وضعیت RNG را به‌صورت اتمیک (atomic) ذخیره کنید تا هرگونه توقف ناگهانی، همیشه فایلی سالم و معتبر روی دیسک بر جای بگذارد.

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

• ثبت تمامیت وضعیت آموزش در قالب یک بسته‌ی داده‌ی (payload) واحد که بتوان آن را در فرایندی تازه بارگذاری کرد.
• پیاده‌سازی ذخیره‌ی اتمیک با الگوی «نوشتن در فایل موقت و سپس تغییر نام» تا در صورت کرش کردن برنامه، هیچ فایل نیمه‌کاره‌ای روی دیسک نماند.
• بازیابی وضعیت RNG برای Python، NumPy و PyTorch تا منحنی تابع زیان پس از ادامه‌ی اجرا، دقیقاً با خط مبنای (baseline) پیوسته همخوانی داشته باشد.
• ساختاردهی چک‌پوینت‌های تکه‌تکه‌شده (sharded) برای مدل‌هایی که دیگر در یک فایل جا نمی‌شوند، همراه با تکه‌های اعتبارسنجی‌شده با هش و یک نمایه‌ی JSON.

صورت مسئله

یک جاب آموزش مدل را برای 18 ساعت تنظیم کرده‌اید. اسکژولر سیستم، سقف هر تخصیص منابع را به 4 ساعت محدود کرده است؛ بنابراین این اجرا همین حالا هم زنجیره‌ای از اجراهای ادامه‌یافته به شمار می‌رود. سپس در ساعت 11، کلاستر ریبوت می‌شود، چون فردی در رده‌های بالاتر ارتقای کرنل (kernel) را تأیید کرده است. بدون چک‌پوینت، باید همه‌چیز را از ابتدا شروع کنید. بدون قابلیت ادامه‌ی اجرا (resume)، وضعیت بهینه‌ساز را هم که یادگیری‌اش 11 ساعت زمان برده از دست می‌دهید؛ یعنی حتی اگر وزن‌های مدل باقی بمانند، گشتاورهای AdamW از بین رفته‌اند و گام بعدی، مدل را به جهتی پرتاب می‌کند که مسیر کامل اجرا (trajectory) آموزش قبلاً از آن عبور کرده بود.

دست‌سازه‌ی (artifact) درست، فایلی واحد است که تمام ملزومات ادامه را در خود دارد: پارامترهای مدل، وضعیت بهینه‌ساز، وضعیت اسکژولر، تاریخچه‌ی تابع زیان برای رسم نمودارها، شمارنده‌های گام و دور (epoch) فعلی و بچِ درون دور، و وضعیت RNG برای تک‌تک منابع تولید تصادفی. بدون وضعیت RNG، منحنی تابع زیان پس از ادامه دادن، مسیر متفاوتی را طی می‌کند: همان مدل، همان داده‌ها، اما بُر خوردن متفاوت داده‌ها، ماسک دراپ‌اوت (dropout) متفاوت و عددی متفاوت روی داشبورد.

نوشتن اتمیک (atomic write)، نیمه‌ی دیگر این قرارداد است. نوشتن مستقیم روی نام فایل نهایی یعنی کرش کردن در میانه‌ی نوشتن، فایلی خراب به جا می‌گذارد و ادامه‌ی اجرا با داده‌های معیوب روبه‌رو می‌شود. نوشتن در یک فایل موقت در همان دایرکتوری و سپس تغییر نام آن، تضمین می‌کند که کرش حین نوشتن، فایل سالم قبلی را دست‌نخورده باقی می‌گذارد. عمل تغییر نام روی فایل‌سیستم‌های POSIX کاملاً اتمیک است.

مفاهیم پایه

پنج دسته‌ی وضعیت آموزش

ذخیره‌ی اتمیک

دو قاعده وجود دارد: اول اینکه فایل موقت حتماً در همان دایرکتوری فایل مقصد ساخته شود تا تغییر نام درون یک فایل‌سیستم یکسان انجام گیرد؛ تغییر نام بین دو درایو یا دستگاه مختلف اتمیک نیست. دوم اینکه نام فایل موقت در هر بار تلاش یکتا باشد تا دو نویسنده روی داده‌های یکدیگر بازنویسی نکنند.

چک‌پوینت‌های تکه‌تکه‌شده (Sharded)

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

این نمایه تعداد تکه‌ها، هش sha256 هر تکه و هش sha256 فایل متادیتا را ثبت می‌کند. لودر در صورت عدم تطابق هر هش، بلافاصله خطای صریح می‌دهد. این تکه‌ها می‌توانند روی دیسک‌های فیزیکی مختلف بنشینند؛ فایل متادیتا کوچک است و ابتدا خوانده می‌شود.

ادامه‌ی آموزش از میانه‌ی دور

اگر ادامه‌ی آموزش همیشه به ابتدای دور بعدی پرش کند، بسته به اندازه‌ی داده‌ها از چند دقیقه تا یک روز زمان تلف می‌شود. راه‌حل، نگه‌داشتن جفت (epoch, batch_in_epoch) به همراه وضعیت RNG است. بعد از بارگذاری، حلقه‌ی آموزش، مولد اعداد تصادفی را به اندازه‌ی بچ‌هایی که در دور جاری مصرف شده‌اند به جلو می‌برد (fast-forward) و دقیقاً از batch_in_epoch ادامه می‌دهد. کد این درس این کار را دقیقاً انجام می‌دهد؛ طوری که مسیر تابع زیان بعد از ادامه‌ی اجرا، با دقت کمتر از 1e-4 با خط مبنای بدون وقفه مطابقت دارد.

پیاده‌سازی

فایل code/main.py چهار جزء سازنده (primitive) و یک نمونه‌ی اجرایی ارائه می‌دهد.

گام 1: ثبت و بازیابی وضعیت RNG

تابع capture_rng_state یک دیکشنری شامل random.getstate پایتون، np.random.get_state نام‌پای و بایت‌های RNG سی‌پی‌یو و CUDA در PyTorch برمی‌گرداند. تک‌تک بخش‌ها به‌صورت اعداد ساده‌ی پایتون، تاپل‌ها و لیست‌ها ذخیره می‌شوند (آرایه‌ی کلیدی NumPy از طریق tolist() تبدیل می‌شود) تا لودر در گام 3 بتواند آن را بدون آن‌پیکل (unpickle) کردن اشیای دلخواه بخواند. تابع restore_rng_state این مسیر را برعکس می‌کند. تنسور CPU یک بافر بایت uint8 است که RNG پای‌تورچ نحوه‌ی مصرف آن را می‌داند.

گام 2: ذخیره‌ی اتمیک

تابع atomic_save بسته را در فایلی موقت داخل دایرکتوری مقصد می‌نویسد و سپس os.replace آن را با نام فایل نهایی جایگزین می‌کند. تابع atomic_write_json همین کار را برای نمایه‌ی شاردها انجام می‌دهد.

گام 3: چرخه‌ی کامل ذخیره و بازیابی چک‌پوینت

تابع save_checkpoint مدل، بهینه‌ساز، اسکژولر، وضعیت آموزش و RNG را در یک دیکشنری بسته‌بندی می‌کند. تابع load_checkpoint این فرایند را معکوس کرده و یک TrainState تحویل می‌دهد. فیلد اسکیما (schema) قلاب ارتقای کد است: با تغییر ساختار در آینده، نسخه‌ی آن بالا می‌رود و لودر مسیر متناسب را انتخاب می‌کند.

تابع load_checkpoint متد torch.load(..., weights_only=True) را فراخوانی می‌کند. فایل‌های .pt با ساختار pickle ذخیره می‌شوند و آن‌پیکل کردن فایلی غیرقابل‌اعتماد با weights_only=False، هر کدی را که درون فایل نام برده شده باشد اجرا خواهد کرد. لودرِ weights-only فقط تنسورها و ظرف‌های داده‌ی پایه‌ای را می‌پذیرد و باقی را رد می‌کند؛ به همین دلیل در گام 1 وضعیت RNG را در قالب لیست‌های ساده نگه داشتیم. بررسی‌های صحت داده‌ها به‌جای assert از ValueError استفاده می‌کنند، چون python -O گزاره‌های assert را حذف می‌کند. از نسخه PyTorch 2.6 یا بالاتر استفاده کنید؛ قبل از این نسخه، weights_only=True یک راه دور زدن شناخته‌شده داشت (CVE-2025-32434)، بنابراین تضمین امنیتی این درس فقط از نسخه 2.6 به بعد برقرار است.

گام 4: نسخه‌ی تکه‌تکه‌شده (Sharded)

تابع save_sharded_checkpoint کلیدهای پارامتر را با روش نوبت‌گردشی (round-robin) بین N تکه پخش می‌کند، هر تکه را با ذخیره‌ی اتمیک اختصاصی خودش می‌نویسد، یک فایل متادیتا شامل وضعیت بهینه‌ساز، اسکژولر و وضعیت آموزش ثبت می‌کند و نمایه‌ی JSON را همراه با هشدارهای sha256 تکه‌ها می‌نویسد. تابع load_sharded_checkpoint همه‌ی تکه‌ها را قبل از ادغام بررسی می‌کند و هر مسیر شاردی را که به خارج از دایرکتوری چک‌پوینت اشاره کند نمی‌پذیرد.

گام 5: آزمایش عملی ادامه‌ی اجرا

تابع run_resume_demo یک مدل کوچک را برای total_steps آموزش می‌دهد، در گام interrupt_at چک‌پوینت ذخیره می‌کند و سپس ادامه می‌دهد. فرایند دوم، چک‌پوینت را بازیابی کرده و بقیه‌ی گام‌ها را اجرا می‌کند. این تابع، بیشینه‌ی قدرمطلق تفاوت بین دو مسیر تابع زیان بعد از نقطه‌ی وقفه را برمی‌گرداند. با بازیابی دقیق RNG، این تفاوت صفر یا در حد نویز ممیز شناور است.

اجرایش کنید:

python3 code/main.py

هر دو آزمون تک‌فایلی و تکه‌تکه‌شده تضمین می‌کنند که بیشینه‌ی اختلاف کمتر از 1e-4 است. خلاصه‌ی نتایج در outputs/resume-demo.json ذخیره می‌شود.

کاربرد در محیط عملیاتی

استک‌های آموزشی در محیط عملیاتی (پروداکشن)، سیستم چک‌پوینت‌گیری (checkpointing) را به‌عنوان بخشی از trainer ارائه می‌دهند. ساختار کلی یکسان است: مدل + بهینه‌ساز + اسکژولر + شمارنده‌ها + RNG که به‌صورت اتمیک نوشته می‌شوند و بر اساس شماره‌ی گام نام‌گذاری می‌شوند تا آخرین نسخه به‌راحتی پیدا شود. لایه‌بندی‌های شاردشده، بارگذاری مدل‌های بزرگ را از طریق خواندن موازی ممکن می‌کنند؛ فایلی مثل index.json عامل اصلی این سازوکار است.

چهار الگویی که باید رعایت شوند:

• بارگذاری حتماً با weights_only=True. چک‌پوینتی که از درایو اشتراکی یا دانلودها می‌آید، یک ورودی غیرقابل‌اعتماد (untrusted input) است. لودر weights-only جلوی اجرای کدهای مخرب روی ماشینی را که فرایند آموزش را ادامه می‌دهد می‌گیرد.
• اسکیما به‌صورت رشته در بسته‌ی داده قرار گیرد. عملیات ارتقا بر اساس این فیلد branch می‌شوند. بدون آن نمی‌توانید فرمت ذخیره‌سازی را بدون از کار انداختن اجراهای قدیمی ارتقا دهید.
• محاسبه‌ی Sha256 برای تک‌تک تکه‌ها. دانلودی که به‌طور ناقص قطع شده و بی‌صدا باقی بماند، بدترین نوع باگ است؛ لودر باید فوراً متوقف شود، نه در مراحل پایانی و دیرتر.
• تنظیم منطقی بازه‌های چک‌پوینت‌گیری. ذخیره‌سازی را هر N گام یک‌بار و هر چند دقیقه بر اساس ساعت واقعی (هر کدام زودتر فرا رسید) انجام دهید. در غیر این صورت، گامی طولانی که دچار کرش می‌شود، کل بازه‌ی محاسباتی انجام‌شده را هدر می‌دهد.

تحویل نهایی

فایل outputs/skill-checkpoint-save-resume.md دستورالعملی آماده برای هر اسکریپت آموزش جدید است: ساختار بسته، نوشتن اتمیک، ثبت RNG و نمایه‌ی شاردها. این مهارت را به مخزن کد خود ببرید، save_checkpoint را در نقطه‌ی ذخیره‌سازی دوره‌ای و load_checkpoint را در زمان راه‌اندازی متصل کنید تا اجرای شما در برابر هرگونه توقف ناگهانی زنده بماند.

تمرین‌ها

1. شاردینگ نوبت‌گردشی را با شاردینگ بر اساس گروه‌های پارامتر (لایه‌های منتهی به .weight در برابر .bias) جایگزین کنید. هر یک از این ساختارها در چه شرایطی ارجحیت دارند؟
2. حلقه‌ی ذخیره را طوری توسعه دهید که فقط K چک‌پوینت آخر حفظ شوند و نسخه‌های قدیمی‌تر پاک شوند. زمانی که فضای دیسک محدود است، مقدار مناسب K چقدر است؟
3. فلگ --ckpt-every-seconds را اضافه کنید تا ذخیره‌سازی بر اساس یک بازه‌ی زمانی مشخص روی ساعت واقعی اجرا شود، نه فقط شمارنده‌ی گام‌ها.
4. مسیری برای بررسی چکسام بسازید که در ابتدای اجرا فعال شود، تک‌تک چک‌پوینت‌های موجود در دایرکتوری را اسکن کند و موارد آسیب‌دیده را گزارش دهد.
5. تابعی با نام migrate_v1_to_v2 پیاده‌سازی کنید که فیلد جدیدی به بسته اضافه کند و نسخه‌ی رشته‌ی اسکیما را ارتقا دهد. منطق بارگذاری را طوری بنویسید که هر دو نسخه را مدیریت کند.

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

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

• ساختار و رفتار تابع rename در استاندارد POSIX برای اطمینان از اتمیک بودن که os.replace به آن متکی است.
• مستندات PyTorch درباره‌ی torch.save و torch.load، شامل پارامتر map_location برای بازیابی در دستگاه‌های سخت‌افزاری دیگر و weights_only جهت بارگذاری فایل‌های غیرقابل‌اعتماد.
• درس 46 از فاز 19 که به موضوع انباشت گرادیان (gradient accumulation) می‌پردازد؛ داده‌های چک‌پوینت این درس در میان این گام‌ها سالم باقی می‌مانند.
• درس 48 از فاز 19 درباره‌ی رپرهای توزیع‌شده که ساختار دیکشنری وضعیت پارامترهای آن‌ها با این الگو هماهنگ است.
• مستندات تابع fsync در کرنل لینوکس برای درک تضمین‌های دوام داده در تغییر نام‌های اتمیک.