fix: align _input_chunk dtype in FLCE grad_weight addmm under AMP

[justinhh4]

Summary

PR #1239 replaced grad_weight += torch.mm(...).float() with torch.addmm(..., out_dtype=torch.float32, out=grad_weight) in the FLCE backward, guarded only on grad_logits_chunk.dtype.

Unlike torch.mm, the addmm out_dtype=/out= overload does not participate in autocast operand casting. So under AMP (fp32 master params, bias=False), grad_logits is the autocast dtype (bf16/fp16) while _input_chunk stays fp32. addmm requires mat1 and mat2 to share a dtype, raising:

RuntimeError: mat1 and mat2 must have the same dtype, but got BFloat16 and Float

bias=True sidesteps the bug because logits + bias promotes grad_logits to fp32. As-is this regresses AMP training of bias=False LM heads (the common LLM case) on torch>=2.8.0.

Details

Align _input_chunk to the grad_logits dtype inside the addmm fast path — a no-op in the intended bf16-native case, exactly what autocast previously did to torch.mm. This preserves #1239's memory optimization and is numerically unchanged on the non-AMP path.

Scope: 1 file, +11/-3, branched directly on top of current main (which already includes #1239).

Testing Done

• test_amp: the 8 bias=False cases that crashed after perf: optimize grad_weight accumulation with addmm #1239 now pass — 16/16 green.

• test_fused_linear_cross_entropy.py: 137/137 pass.

• Numerically verified vs a pure-torch reference (loss / grad_input / grad_weight / grad_bias) across {bf16, fp16, no-AMP} x {bias, nobias}; the non-AMP mm().float() path is bit-for-bit unchanged.

• Hardware Type: B200

• run make test to ensure correctness

• run make checkstyle to ensure code style

• run make test-convergence to ensure convergence