Fix orthogonal_ init for low-precision dtypes (bf16/fp16)

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
 

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