raise error for nan hidden states

scripts/data_generation_offline.py

@@ -23,10 +23,11 @@
 
 import openai
 from datasets import load_from_disk
-from safetensors import safe_open
+from safetensors.torch import load_file
 from tqdm import tqdm
 
 from speculators.data_generation.offline import (
+    check_hidden_states,
     get_existing_hidden_state_indices,
     get_indices_to_process,
 )
@@ -191,23 +192,6 @@ def parse_args():
     return parser.parse_args()
 
 
-def check_safetensors_file(path: Path, tokens: list[int]):
-    with safe_open(path, "pt") as f:
-        t_ids = f.get_tensor("token_ids").tolist()
-        if t_ids != tokens:
-            raise ValueError(
-                f"Token ids in {path} don't match expected token ids {tokens}"
-            )
-
-        hs_slice = f.get_slice("hidden_states")
-        hs_shape = list(hs_slice.get_shape())
-        if len(tokens) != hs_shape[0]:
-            raise ValueError(
-                f"Sequence length of hidden states {hs_shape[0]} in {path}"
-                f" doesn't match num tokens {len(tokens)}"
-            )
-
-
 async def worker(  # noqa: C901
     client,
     model: str,
@@ -258,11 +242,15 @@ async def worker(  # noqa: C901
                     shutil.move, hidden_states_path, target_hidden_states_path
                 )
                 if validate_outputs:
-                    await asyncio.to_thread(
-                        check_safetensors_file,
-                        target_hidden_states_path,
-                        item["input_ids"],
-                    )
+
+                    def _load_and_check(
+                        path=target_hidden_states_path,
+                        tokens=item["input_ids"],
+                    ):
+                        loaded = load_file(path)
+                        check_hidden_states(loaded, tokens)
+
+                    await asyncio.to_thread(_load_and_check)
         except Exception as e:
             if fail_on_error:
                 logger.exception(

src/speculators/data_generation/offline.py

@@ -4,6 +4,21 @@
 logger = logging.getLogger(__name__)
 
 
+def check_hidden_states(data: dict, tokens: list[int]):
+    t_ids = data["token_ids"].tolist()
+    if t_ids != tokens:
+        raise ValueError(f"Token ids don't match expected token ids {tokens}")
+
+    hs = data["hidden_states"]
+    if hs.isnan().any():
+        raise ValueError("Hidden states contain NaN values")
+    if len(tokens) != hs.shape[0]:
+        raise ValueError(
+            f"Sequence length of hidden states {hs.shape[0]}"
+            f" doesn't match num tokens {len(tokens)}"
+        )
+
+
 def get_existing_hidden_state_indices(output_path: Path) -> list[int]:
     """Find existing `hs_i.safetensors` files (where i is the file index)"""
 

src/speculators/train/data.py

@@ -16,6 +16,7 @@
 from safetensors.torch import load_file
 from torch.utils.data import Dataset
 
+from speculators.data_generation.offline import check_hidden_states
 from speculators.data_generation.vllm_client import (
     DEFAULT_MAX_RETRIES,
     DEFAULT_REQUEST_TIMEOUT,
@@ -322,6 +323,10 @@ def _maybe_generate_hs(self, index: int) -> dict[str, torch.Tensor] | None:
             )
 
             loaded_hs = _maybe_load_hs_file(Path(hs_filepath))
+            if loaded_hs is None:
+                raise ValueError(f"Failed to load hidden states from {hs_filepath}")
+
+            check_hidden_states(loaded_hs, dataset_item["input_ids"].tolist())
 
             match self.on_generate:
                 case "cache":
@@ -330,7 +335,9 @@ def _maybe_generate_hs(self, index: int) -> dict[str, torch.Tensor] | None:
                     shutil.move(hs_filepath, target_path)
                 case "delete":
                     Path(hs_filepath).unlink()
-        except Exception as e:  # noqa: BLE001
+        except Exception as e:
+            if isinstance(e, ValueError) and "NaN" in str(e):
+                raise
             warnings.warn(
                 f"Failed to load/cache hidden states for sample {index}: {e}",
                 stacklevel=1,
@@ -365,7 +372,7 @@ def _get_raw_data(self, index):
             return loaded_hs
 
         # loaded_hs structure: {
-        #   "hidden_states": [seq_len, 4, hidden_size]
+        #   "hidden_states": [seq_len, num_layers, hidden_size]
         #   "token_ids": [seq_len]
         # }