Add FP32 accumulation option to fused scaled cross entropy

[yueyiming2009]

Summary

• Add an optional accum_dtype to fused scaled cross entropy so chunked weight gradients can accumulate in FP32 before casting back to the parameter dtype.
• Support the option in the cuTile implementation and PyTorch fallback while preserving the existing default behavior; reject it explicitly for the fused CuTe SM90 path.
• Add correctness coverage, a cutile-accum-fp32 benchmark provider, and usage documentation.

Details

BF16 grad_weight was previously rounded after every token chunk because the running accumulator inherited weight.dtype. With accum_dtype=torch.float32, supported CUDA systems use torch.addmm(..., out_dtype=torch.float32) to avoid that chunk-boundary precision loss. The tradeoff is an additional V x H FP32 buffer.

On an NVIDIA B200 at M=8192, H=4096, V=131072, FP32 accumulation added 0.7% full-pass latency and 2048 MB peak memory. A chunk-count sweep showed the gradient error improvement growing from 1.34x at 16 chunks to 7.31x at 1024 chunks.

Testing Done

• Hardware Type: NVIDIA B200
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence
• Local code review completed
• Ran test/transformers/test_fused_linear_scaled_cross_entropy.py with LIGER_KERNEL_IMPL=cutile on B200: 36 passed
• Ran targeted ruff check and ruff format --check on all changed Python files
• Benchmarked cutile versus cutile-accum-fp32 on B200 for speed, memory, and precision

🤖 Generated with GitHub Copilot CLI