Tính Cross Entropy Loss mà vẫn tiết kiệm bộ nhớ

Bài viết chỉ ra vì sao trong quá trình huấn luyện LLM với context dài và vocab lớn, LM head + cross entropy lại trở thành một trong những nơi tiêu tốn bộ nhớ nhiều nhất. Với context 128K, chỉ riêng một tensor logits đã có thể chạm ngưỡng gần 40GB, thậm chí còn lớn hơn cả trọng số của mô hình.

Bắt đầu từ sự cố OOM thực tế khi huấn luyện mô hình 16B với context 128K, bài viết suy ra lại forward/backward của cross entropy từ đầu, rồi cho thấy vì sao cách đơn giản là chia trục sequence thành các chunk vẫn không thể giảm peak memory (do autograd giữ graph của từng chunk cho tới khi backward). Sau đó, bài viết giải thích cách FLCE tính gradient của từng chunk ngay trong forward pass để các tensor lớn không còn bị giữ lại trong graph. Cuối bài còn phân tích tradeoff giữa bộ nhớ/độ trễ và đi qua từng bước triển khai kernel thực tế.