From 2343d4989211590ccf37db21757eda1d5fe09b8c Mon Sep 17 00:00:00 2001 From: jbernloehr Date: Fri, 10 Jul 2026 13:41:17 +0200 Subject: [PATCH 1/2] 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 --- src/transformers/initialization.py | 7 ++++ tests/utils/test_initialization.py | 54 ++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+) create mode 100644 tests/utils/test_initialization.py diff --git a/src/transformers/initialization.py b/src/transformers/initialization.py index b0ebb053086b..74e0e07320a8 100644 --- a/src/transformers/initialization.py +++ b/src/transformers/initialization.py @@ -144,6 +144,13 @@ def orthogonal_( generator: torch.Generator | None = None, ) -> torch.Tensor: if not getattr(tensor, "_is_hf_initialized", False): + # The QR decomposition (`geqrf`) used by `torch.nn.init.orthogonal_` is only implemented for + # float32/float64, so for lower-precision dtypes we run the init in float32 and copy back + if tensor.is_floating_point() and torch.finfo(tensor.dtype).bits < 32: + fp32_tensor = torch.empty_like(tensor, dtype=torch.float32) + TORCH_INIT_FUNCTIONS["orthogonal_"](fp32_tensor, gain=gain, generator=generator) + with torch.no_grad(): + return tensor.copy_(fp32_tensor) return TORCH_INIT_FUNCTIONS["orthogonal_"](tensor, gain=gain, generator=generator) return tensor diff --git a/tests/utils/test_initialization.py b/tests/utils/test_initialization.py new file mode 100644 index 000000000000..9c715736c197 --- /dev/null +++ b/tests/utils/test_initialization.py @@ -0,0 +1,54 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +from transformers import is_torch_available +from transformers.testing_utils import require_torch + + +if is_torch_available(): + import torch + + from transformers import initialization as init + + +@require_torch +class TestInitialization(unittest.TestCase): + def test_orthogonal_low_precision(self): + # The QR decomposition (`geqrf`) used by `torch.nn.init.orthogonal_` is not implemented for + # low-precision dtypes, in which case the init runs in float32 and is copied back (see #47225) + gain = 2.0 + for dtype in (torch.bfloat16, torch.float16): + with self.subTest(dtype=dtype): + tensor = torch.empty(16, 16, dtype=dtype) + init.orthogonal_(tensor, gain=gain) + self.assertEqual(tensor.dtype, dtype) + self.assertTrue(torch.isfinite(tensor).all()) + # Rows should be orthogonal with norm equal to `gain`, up to low-precision rounding + product = (tensor.float() @ tensor.float().T) / gain**2 + torch.testing.assert_close(product, torch.eye(16), atol=5e-2, rtol=0) + + def test_orthogonal_full_precision_matches_torch(self): + tensor = torch.empty(16, 16) + init.orthogonal_(tensor, generator=torch.Generator().manual_seed(0)) + expected = torch.empty(16, 16) + torch.nn.init.orthogonal_(expected, generator=torch.Generator().manual_seed(0)) + torch.testing.assert_close(tensor, expected, atol=0, rtol=0) + + def test_orthogonal_respects_hf_initialized_flag(self): + tensor = torch.zeros(16, 16, dtype=torch.bfloat16) + tensor._is_hf_initialized = True + init.orthogonal_(tensor) + self.assertTrue((tensor == 0).all()) From d4701f1fc81a751d658854276969360fef4537c2 Mon Sep 17 00:00:00 2001 From: jbernloehr Date: Wed, 15 Jul 2026 15:22:56 +0200 Subject: [PATCH 2/2] 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 --- tests/utils/test_initialization.py | 54 ------------------------------ tests/utils/test_modeling_utils.py | 15 +++++++++ 2 files changed, 15 insertions(+), 54 deletions(-) delete mode 100644 tests/utils/test_initialization.py diff --git a/tests/utils/test_initialization.py b/tests/utils/test_initialization.py deleted file mode 100644 index 9c715736c197..000000000000 --- a/tests/utils/test_initialization.py +++ /dev/null @@ -1,54 +0,0 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import unittest - -from transformers import is_torch_available -from transformers.testing_utils import require_torch - - -if is_torch_available(): - import torch - - from transformers import initialization as init - - -@require_torch -class TestInitialization(unittest.TestCase): - def test_orthogonal_low_precision(self): - # The QR decomposition (`geqrf`) used by `torch.nn.init.orthogonal_` is not implemented for - # low-precision dtypes, in which case the init runs in float32 and is copied back (see #47225) - gain = 2.0 - for dtype in (torch.bfloat16, torch.float16): - with self.subTest(dtype=dtype): - tensor = torch.empty(16, 16, dtype=dtype) - init.orthogonal_(tensor, gain=gain) - self.assertEqual(tensor.dtype, dtype) - self.assertTrue(torch.isfinite(tensor).all()) - # Rows should be orthogonal with norm equal to `gain`, up to low-precision rounding - product = (tensor.float() @ tensor.float().T) / gain**2 - torch.testing.assert_close(product, torch.eye(16), atol=5e-2, rtol=0) - - def test_orthogonal_full_precision_matches_torch(self): - tensor = torch.empty(16, 16) - init.orthogonal_(tensor, generator=torch.Generator().manual_seed(0)) - expected = torch.empty(16, 16) - torch.nn.init.orthogonal_(expected, generator=torch.Generator().manual_seed(0)) - torch.testing.assert_close(tensor, expected, atol=0, rtol=0) - - def test_orthogonal_respects_hf_initialized_flag(self): - tensor = torch.zeros(16, 16, dtype=torch.bfloat16) - tensor._is_hf_initialized = True - init.orthogonal_(tensor) - self.assertTrue((tensor == 0).all()) diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py index 638e17a22c8d..6e167555dc2e 100644 --- a/tests/utils/test_modeling_utils.py +++ b/tests/utils/test_modeling_utils.py @@ -437,6 +437,21 @@ def test_get_total_byte_count_does_not_require_process_group(self): self.assertIn(torch.device("cpu"), total_byte_count) self.assertGreater(total_byte_count[torch.device("cpu")], 0) + def test_orthogonal_init_low_precision(self): + # `torch.nn.init.orthogonal_` uses a QR decomposition (`geqrf`) that is only implemented for + # float32/float64, so `init.orthogonal_` runs it in float32 and copies back for low-precision + # dtypes. Without the fallback this raises `NotImplementedError: "geqrf_..." not implemented`. + gain = 2.0 + for dtype in (torch.bfloat16, torch.float16): + with self.subTest(dtype=dtype): + tensor = torch.empty(16, 16, dtype=dtype) + init.orthogonal_(tensor, gain=gain) + self.assertEqual(tensor.dtype, dtype) + self.assertTrue(torch.isfinite(tensor).all()) + # Rows should be orthogonal with norm `gain`, up to low-precision rounding + product = (tensor.float() @ tensor.float().T) / gain**2 + torch.testing.assert_close(product, torch.eye(16), atol=5e-2, rtol=0) + def test_hub_retry(self): @hub_retry(max_attempts=2) def test_func():