‏Monarch پای‌تورچ روی GPUهای AMD

تیم پای‌تورچ Monarch رو به GPUهای AMD Instinct با پلتفرم ROCm آورده. Monarch یه رانتایم توزیع‌شدهٔ «تک‌کنترلره»؛ یعنی می‌تونی کل یه کلاستر GPU رو با یه برنامهٔ پایتون واحد مدیریت کنی. تا الان این مدل فقط روی محیط CUDA (انویدیا) کار می‌کرد و حالا به سخت‌افزار AMD هم گسترش پیدا کرده.

مشکلی که این کار حل می‌کنه قابلیت اطمینان توی مقیاس بزرگه. وقتی داری یه مدل رو روی صدها یا هزاران GPU آموزش می‌دی، خرابی سخت‌افزار یه اتفاق عادیه نه استثنا. روش سنتی اینه که مرتب از کل وضعیت مدل checkpoint بگیری و بعد از هر خرابی، کل کار از آخرین checkpoint دوباره شروع شه—که هم کند و هم پرهزینه‌ست.

به گفتهٔ نویسنده‌ها، Monarch به جای این کار از یه مدل اکتور-محور و «درخت نظارت» (supervision tree) استفاده می‌کنه که خرابی‌ها رو ایزوله می‌کنه. وقتی یه نود خراب می‌شه، بقیهٔ نودهای سالم آموزش رو ادامه می‌دن و نودِ خراب به صورت محلی ری‌استارت می‌شه و دوباره به جمع می‌پیونده—بدون اینکه کل کار متوقف شه.

توی این معماری، Monarch نقش ارکستریتور رو داره، TorchFT مدیریت خطا رو در سطح هر قدم انجام می‌ده و TorchTitan هم موتور آموزشه. نکتهٔ جالب اینه که وقتی یه نود می‌خواد برگرده، به جای بارگذاری یه checkpoint سراسری، وضعیتش رو مستقیم از یه نود سالم دیگه (peer) کپی می‌کنه—پس زمان تلف‌شده خیلی کمتره.

پورت کردن به ROCm کار مهندسی زیادی برد. تیم با ابزار hipify_torch کد پل ‪C++‬ رو از CUDA به HIP تبدیل کرد و به کتابخونهٔ RCCL (که API اون مثل NCCL هست) لینک زد. واسه اینکه مجبور نشن کد Rust رو شاخه‌شاخه کنن، یه ماژول سازگاری به اسم rocm_compat ساختن که نام‌های HIP رو زیر اسم‌های CUDA دوباره صادر می‌کنه:

pub type cudaError_t = hipError_t;
pub use hipSetDevice as cudaSetDevice;

با این ترفند بقیهٔ کد Rust بدون تغییر و مستقل از پلتفرم می‌مونه.

این مجموعه با موفقیت روی کلاسترهای AMD تست شده؛ مثلاً آموزش یه مدل Llama 3 8B روی ۱۲۸ تا GPU از نوع MI300 که هر ۱۸۰ ثانیه بهش خرابی تزریق می‌شد، بدون ری‌استارت کامل ادامه پیدا کرد و منحنی loss تقریباً مثل حالت بدون خرابی هم‌گرا شد. یه آزمایش بزرگ‌تر هم روی ۲۵۶ تا GPU از نوع MI355 روی Kubernetes پایدار موند.

نکات کلیدی:

• Monarch یه رانتایم تک‌کنترلر واسه مدیریت کل کلاستر GPU با یه برنامهٔ پایتونه
• حالا علاوه بر CUDA روی ROCm و GPUهای AMD هم کار می‌کنه
• به جای ری‌استارت کامل، وضعیت نودِ خراب مستقیم از یه نود سالم کپی می‌شه
• روی کلاسترهای MI300 و MI355 با تزریق خرابی مداوم، آموزش بدون توقف ادامه پیدا کرد