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_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():