Fused Cross Entropy Loss

[winglian]

pytorch/pytorch#124480 served as the impetus for getting this integrated.

Adapted the code from @mgmalek's https://github.com/mgmalek/efficient_cross_entropy/blob/main/modules.py so that it properly handles causal/next-token prediction.

[YouJiacheng]

mgmalek's implementation might have numerical issues when using half precision dtype.
The gradient w.r.t. the weight is accumulated through chunks with a half precision accumulator.
Some users have reported loss curve mismatch on Twitter.

[winglian]

mgmalek's implementation might have numerical issues when using half precision dtype.
The gradient w.r.t. the weight is accumulated through chunks with a half precision accumulator.
Some users have reported loss curve mismatch on Twitter.

haha, yeah, that was probably me reporting the loss curve differences out on twitter. Do you have any insights on the best way to fix the kernels? standard CEL uses bfloat16 as well iirc.

[YouJiacheng]

The problem is the lm_head weight (proj_weight), not CEL. By default bf16 matmul use fp32 accumulator (along the GEMM K axis). Since matmul is performed tile by tile in M&N axis (MxK@KxN matmul), registers are used as accumulator, without the need of a fp32 global memory buffer.

However, here we perform chunking along the GEMM K axis (DxT@TxV matmul)…

A simple fix can be using fp32 grad_proj_weight, but that will incur a significant memory overhead.

Can you check the loss curve with fp32 grad_proj_weight?

[winglian]

@YouJiacheng

gave your recommendation a try, but changing grad_proj_weight to fp32 didn't seem to help much.

[YouJiacheng]

Ah Oh, now we need more investigation…

[YouJiacheng]

It seems that the fused version has a different loss & grad_norm from the first step? That looks strange…