PyTorch Monarch chạy trên AMD ROCm

Giới thiệu

PyTorch Monarch là runtime để điều phối distributed training từ một chương trình Python duy nhất. Bài viết này thông báo Monarch đã được port sang AMD Instinct GPU qua ROCm, thay vì chỉ sống trong hệ CUDA.

Bối cảnh khá rõ: training LLM nhiều tỷ parameter trên hàng trăm hoặc hàng nghìn GPU không chỉ cần scale tốt, mà còn cần chịu lỗi tốt. Một GPU lỗi bộ nhớ, network partition hoặc node crash có thể làm phí nhiều ngày compute nếu job phải restart từ checkpoint.

Tính năng chính

Monarch kết hợp vài khối hạ tầng:

• Python API cho single-program distributed execution.
• Actor runtime và process mesh để quản lý cụm GPU.
• Rust runtime dựa trên Tokio.
• Tích hợp RDMA, RCCL/NCCL, SLURM, Kubernetes và SkyPilot.
• Supervision tree để cô lập lỗi ở mức actor hoặc replica.
• Tensor sharding và orchestration cho training, evaluation hoặc reinforcement learning trong một script.

Điểm quan trọng là fault tolerance được tách khỏi parallelism strategy. Khi một phần training replica hỏng, lỗi có thể được xử lý ở tầng gần nhất thay vì làm toàn bộ cluster dừng lại.

Port sang ROCm

Nhóm PyTorch mô tả ba phần chính khi đưa Monarch sang AMD GPU:

• Collective communications: dùng hipify_torch để chuyển bridge C++ từ CUDA sang HIP và link với RCCL, lớp tương đương NCCL trong hệ ROCm.
• GPU memory management: build system tự phát hiện platform và route CUDA driver API sang HIP equivalent.
• RDMA integration: giữ đường libibverbs cho RDMA, chỉ đổi binding phía GPU từ CUDA sang HIP.

Một chi tiết kỹ thuật hay là họ không fork Rust code theo hai nhánh CUDA/ROCm. Sau khi bindgen sinh ra type HIP, các module compatibility re-export symbol HIP dưới tên CUDA như cudaError_t hoặc cudaSetDevice. Nhờ vậy phần còn lại của Rust runtime ít phải biết đang chạy trên NVIDIA hay AMD.

Nhóm tác giả nói toàn bộ 1.171 test pass và các thay đổi đã được upstream qua PR #2393 và #2891 trong repository Monarch.

Fault recovery hoạt động ra sao

Case study trong bài kết hợp Monarch, TorchFT và TorchTitan:

• Monarch quản lý process và cluster orchestration.
• TorchFT xử lý quorum, step-level fault tolerance và skip node hỏng.
• TorchTitan chạy forward, backward, optimizer, checkpoint và metrics.

Kịch bản ví dụ có bốn replica group, mỗi replica có tám GPU process. Khi một GPU process trong replica 0 crash, supervisor của Monarch bắt lỗi và restart replica đó tại chỗ. Các replica còn lại vẫn tiếp tục training. Sau đó Lighthouse chọn một replica khỏe làm donor để chuyển checkpoint ngang hàng, rồi quorum mới được tạo và training tiếp tục.

Ý nghĩa với Dev

Nếu bạn chỉ train nhỏ trên một node, Monarch có thể hơi xa nhu cầu hằng ngày. Nhưng với team đang tính chuyện AMD GPU, neocloud, hoặc cluster nhiều node, bài này là tín hiệu tốt: hệ sinh thái ROCm đang có thêm runtime production-grade thay vì chỉ có benchmark model.

Điểm nên kiểm tra khi áp dụng:

• Khả năng chạy với scheduler hiện có: SLURM, Kubernetes hoặc SkyPilot.
• Chi phí checkpoint và thời gian mất khi restart job.
• Compatibility của RCCL/RDMA trên phần cứng cụ thể.
• Observability khi actor crash, restart và rejoin quorum.

Nói ngắn gọn: scaling training chưa đủ; job càng lớn thì recovery càng là feature chính.