feat(ce,flce): decouple gradients computation for no_grad mode

[Tcc0403]

Summary

Add a flag HAS_GRADIENTS to cross entropy kernel. No more gradients computation if there's no need.

Testing Done

Cross Entropy forward with no_grad

Fused Linear Cross Entropy forward with no_grad

• Hardware Type:
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[shimizust]

Nice, thanks for adding this. The forward flce still is significantly slower than hf since we're still computing grad_input applying token scaling logic? Also do you know why fp32 accum is faster?

[Tcc0403]

@shimizust The slower forward pass is kinda expected because:

1. instead of one big matmul, we slice input rowwise and do multiple matmuls. It brings some kernel launch overhead. Plus, it's not gauranteed to ivnoke the most efficient kernels (tiling size, tail effect, etc... from kernel's prespective).
2. multiple cross entropy kernel launch overhead and similar issue also add up.

I just found that we can remove this line in eval mode as well, cutting another matmul for each interation should be significant.

Liger-Kernel/src/liger_kernel/ops/fused_linear_cross_entropy.py

Line 175 in 5c2a04d

grad_input[start_idx:end_idx] = grad_logits_chunk @ weight

and some grad tensors allocations too