Cơ chế tuần tự hóa và khôi phục trạng thái bộ tối ưu hóa (Optimizer state_dict Checkpointing)

Mô tả bài toán

Trong các bài toán huấn luyện quy mô lớn (LLM Pretraining), quá trình huấn luyện có thể gặp sự cố phần cứng hoặc phải dừng lại định kỳ. Nếu chỉ lưu trữ trọng số mô hình mà quên lưu trạng thái của Optimizer (step, các vector đà m và phương sai v), quá trình huấn luyện tiếp tục sẽ bị hỏng hoàn toàn do mất toàn bộ thống kê lịch sử gradient.

Hãy xây dựng lớp ResumableAdamW:

• __init__(self, params: dict[str, np.ndarray], lr: float = 1e-3, beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-8, weight_decay: float = 0.01)
• step(self, grads: dict[str, np.ndarray]) -> None: cập nhật tham số theo thuật toán AdamW đầy đủ.
• state_dict(self) -> dict:
• Trả về dictionary chứa toàn bộ trạng thái nội bộ:

    {
        "step": self.step_count,
        "hyperparameters": {
            "lr": self.lr, "beta1": self.beta1, "beta2": self.beta2,
            "eps": self.eps, "weight_decay": self.weight_decay
        },
        "state": {
            key: {"m": np.copy(self.m[key]), "v": np.copy(self.v[key])}
            for key in self.params
        }
    }

• load_state_dict(self, state_dict: dict) -> None:
• Khôi phục chính xác step_count, các siêu tham số, và các mảng buffer m, v.
• Nếu cấu trúc state_dict bị thiếu khóa hoặc kích thước mảng không khớp với params, raise ValueError.

Ràng buộc

• Thời gian chạy tối đa: 6000ms.
• Giới hạn bộ nhớ: 512MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.