Hàm Mất Mát Cân Bằng Tải Phụ Trợ (MoE Auxiliary Load Balancing Loss)

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

• Hiện tượng Sụp đổ Chuyên gia (Expert Collapse / Routing Collapse):

Nếu không có cơ chế điều tiết, bộ định tuyến MoE sẽ có xu hướng gửi 95% số token cho 1 hoặc 2 chuyên gia quen thuộc vì chúng được cập nhật thường xuyên hơn. Các chuyên gia còn lại sẽ "chết lâm sàng", biến mô hình MoE trở lại thành mô hình Dense nhỏ bé.

• Hàm mất mát phụ trợ (Auxiliary Loss - Fedus et al., 2021):

Cho N tokens trong batch và E chuyên gia: L_aux = α · E ∑_i=1^E f_i · P_i Trong đó:

• f_i = Số token được gán cho expert iN · K: Tỷ lệ tải thực tế của chuyên gia i.
• P_i = 1N ∑_t=1^N Softmax(router_logits)_i: Xác suất trung bình của chuyên gia i trên toàn bộ batch.
• α: Hệ số phạt (thường α = 0.01).
• Giá trị nhỏ nhất đạt được khi f_i = 1/E và P_i = 1/E (tải phân bổ đồng đều tuyệt đối).

Yêu cầu

Viết hàm:

def compute_moe_load_balancing_loss(
    router_logits: torch.Tensor,
    topk_indices: torch.Tensor,
    num_experts: int,
    alpha: float = 0.01
) -> torch.Tensor:
    pass

• router_logits: Tensor shape (total_tokens, num_experts).
• topk_indices: Tensor shape (total_tokens, top_k).
• Trả về giá trị scalar loss float.

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ả.