Bộ Điều Chỉnh Tốc Độ Học Cosine Annealing Kèm Linear Warmup

Mục tiêu kiến thức

• Hiểu tại sao chiến lược Cosine Annealing kèm Linear Warmup là lựa chọn mặc định của GPT, LLaMA, Chinchilla:
• Giai đoạn Warmup (t ≤ T_warmup): Tăng tuyến tính từ 0 lên lr_max để tránh gradient khổng lồ ở các epoch đầu phá hủy các trọng số khởi tạo ngẫu nhiên.

lr(t) = lr_max · tT_warmup

• Giai đoạn Cosine Decay (T_warmup < t ≤ T_max): Giảm dần theo đường cong hàm cosin từ lr_max về lr_min:

lr(t) = lr_min + 12 (lr_max - lr_min) ( 1 + cos( π t - T_warmupT_max - T_warmup ) )

• Sau T_max (t > T_max): Giữ nguyên lr_min.

Yêu cầu

Xây dựng lớp CosineWarmupLRScheduler:

class CosineWarmupLRScheduler:
    def __init__(self, max_lr: float, min_lr: float, warmup_steps: int, total_steps: int):
        pass

    def get_lr(self, step: int) -> float:
        pass

• step: Bước huấn luyện hiện tại (nguyên ≥ 0).

Ràng buộc

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