Fix orthogonal_ init for low-precision dtypes (bf16/fp16) (#47252)
* Fix orthogonal_ init for low-precision dtypes (bf16/fp16)
torch.nn.init.orthogonal_ relies on torch.linalg.qr, whose geqrf kernel
is only implemented for float32/float64 (on both CPU and CUDA). Since v5
materializes weights in the requested dtype before _init_weights runs,
any QR-based init on a bf16/fp16 tensor crashes with e.g.
NotImplementedError: "geqrf_cuda" not implemented for 'BFloat16'
(RWKV being one affected model).
Run the init in float32 and copy the result back for low-precision
floating dtypes, on any device.
Fixes #47225
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* Move orthogonal_ low-precision test into test_modeling_utils.py
Address review: drop the standalone test file and keep a single
condensed low-precision regression test in an existing file.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* Update tests/utils/test_modeling_utils.py
---------
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Matt <Rocketknight1@users.noreply.github.com>