bugfix

Author: ilyalasy

clt/activation_generation/generator.py

@@ -41,7 +41,7 @@
 from collections import defaultdict
 import psutil
 
-from clt.training.utils import to_bfloat16
+from clt.training.utils import torch_bfloat16_to_numpy_uint16
 
 try:
     import GPUtil
@@ -768,9 +768,8 @@ def _write_chunk(
                                 with self._conditional_measure(f"chunk_{chunk_idx}_layer_{lid}_convert_numpy"):
                                     # 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())
-
+                                        inp_np = torch_bfloat16_to_numpy_uint16(inp_perm)
+                                        tgt_np = torch_bfloat16_to_numpy_uint16(tgt_perm)
                                     else:
                                         inp_np = inp_perm.to(self.torch_dtype).numpy()
                                         tgt_np = tgt_perm.to(self.torch_dtype).numpy()

clt/training/utils.py

@@ -1,6 +1,7 @@
 import datetime
 
 import numpy as np
+import torch
 
 
 # Helper function to format elapsed time
@@ -15,5 +16,5 @@ def _format_elapsed_time(seconds: float) -> str:
         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)
+def torch_bfloat16_to_numpy_uint16(x: torch.Tensor) -> np.ndarray:
+    return np.frombuffer(x.float().numpy().tobytes(), dtype=np.uint16)[1::2].reshape(x.shape)