fix bfloat16 as it's not supported by numpy directly

Author: ilyalasy

clt/activation_generation/generator.py

@@ -41,6 +41,8 @@
 from collections import defaultdict
 import psutil
 
+from clt.training.utils import to_bfloat16
+
 try:
     import GPUtil
 except ImportError:
@@ -764,13 +766,14 @@ def _write_chunk(
 
                                 # Convert to numpy
                                 with self._conditional_measure(f"chunk_{chunk_idx}_layer_{lid}_convert_numpy"):
-                                    inp_np = inp_perm.to(self.torch_dtype).numpy()
-                                    tgt_np = tgt_perm.to(self.torch_dtype).numpy()
-
-                                # Handle bfloat16 conversion
-                                if h5py_dtype_str == "uint16" and inp_np.dtype == np.dtype("bfloat16"):
-                                    inp_np = inp_np.view(np.uint16)
-                                    tgt_np = tgt_np.view(np.uint16)
+                                    # Handle bfloat16 conversion
+                                    if h5py_dtype_str == "uint16":
+                                        inp_np = to_bfloat16(inp_perm.to(torch.float32).numpy())
+                                        tgt_np = to_bfloat16(tgt_perm.to(torch.float32).numpy())
+
+                                    else:
+                                        inp_np = inp_perm.to(self.torch_dtype).numpy()
+                                        tgt_np = tgt_perm.to(self.torch_dtype).numpy()
 
                                 # Store prepared data
                                 layer_data[lid] = (inp_np, tgt_np)

clt/training/utils.py

@@ -1,5 +1,7 @@
 import datetime
 
+import numpy as np
+
 
 # Helper function to format elapsed time
 def _format_elapsed_time(seconds: float) -> str:
@@ -11,3 +13,7 @@ def _format_elapsed_time(seconds: float) -> str:
         return f"{td.days * 24 + hours:02d}:{minutes:02d}:{seconds:02d}"
     else:
         return f"{minutes:02d}:{seconds:02d}"
+
+
+def to_bfloat16(x: np.ndarray) -> np.ndarray:
+    return np.frombuffer(np.array(x, dtype=np.float32).tobytes()[::2], dtype=np.uint16)