Cộng Dồn Gradient Không Bộ Đệm Bằng ufunc.at Trong Huấn Luyện AI

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

• Hiểu rõ cơ chế đệm (buffering) trong phép gán mảng nâng cao arr[indices] += values của NumPy và hiện tượng mất mát dữ liệu khi có chỉ số trùng lặp.
• Sử dụng phương thức không bộ đệm np.add.at để thực hiện phép toán Scatter-Add chuẩn xác trong học sâu.
• Cập nhật gradient cho ma trận nhúng từ (Embedding Matrix) trong mô hình ngôn ngữ lớn.

Mô tả bài toán

Trong quá trình lan truyền ngược (Backpropagation), một câu văn bản đầu vào có thể chứa nhiều từ giống nhau (ví dụ: ["học", "lập", "trình", "và", "học", "AI"] chứa 2 từ "học" cùng có mã vựng token_id = 42). Nếu lập trình viên thực hiện phép gán thông thường: grad_table[token_ids] += grads NumPy sẽ nạp giá trị cũ vào bộ đệm một lần duy nhất, dẫn đến việc gradient của từ "học" thứ hai ghi đè lên từ "học" thứ nhất thay vì cộng dồn cả hai!

Hãy viết hàm: scatter_accumulate_embeddings(vocab_size: int, embed_dim: int, token_ids: np.ndarray, token_grads: np.ndarray) -> np.ndarray

Yêu cầu:

• Kiểm tra tham số:
• Nếu vocab_size <= 0 hoặc embed_dim <= 0: ném ngoại lệ ValueError("vocab_size va embed_dim phai > 0").
• Nếu len(token_ids) != len(token_grads): ném ngoại lệ ValueError("Kich thuoc token_ids va token_grads phai bang nhau").
• Nếu có bất kỳ chỉ số nào trong token_ids nằm ngoài phạm vi [0, vocab_size - 1]: ném ngoại lệ IndexError("token_id vuot qua vocab_size").
• Khởi tạo và tích lũy:
• Khởi tạo ma trận gradient kích thước (vocab_size, embed_dim) toàn số 0 với cùng dtype của token_grads.
• Sử dụng np.add.at để cộng dồn toàn bộ các gradient vào đúng vị trí dòng trong ma trận, bảo toàn 100% gradient của các token trùng lặp.
• Kết quả trả về: Ma trận gradient hoàn chỉnh.

Ràng buộc

• 1 ≤ vocab_size ≤ 50000.
• 1 ≤ embed_dim ≤ 1024.
• Thời gian chạy tối đa: 1000ms.
• Giới hạn bộ nhớ: 256MB.