| import logging |
| from pathlib import Path |
|
|
| import torch |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| def check_hidden_states(data: dict, tokens: list[int]): |
| required = {"token_ids", "hidden_states"} |
| missing = required - data.keys() |
| if missing: |
| raise ValueError(f"Hidden-state payload is missing keys: {missing}") |
|
|
| 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 not isinstance(hs, torch.Tensor): |
| raise ValueError(f"Hidden states must be a tensor, got {type(hs).__name__}") |
| 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)}" |
| ) |
|
|
| nan_count = 0 |
| inf_count = 0 |
| affected_layers: set[int] = set() |
| rows_per_chunk = 256 |
| for start in range(0, hs.shape[0], rows_per_chunk): |
| |
| chunk = hs[start : start + rows_per_chunk] |
| finite = torch.isfinite(chunk) |
| if finite.all(): |
| continue |
| nan_count += int(torch.isnan(chunk).sum().item()) |
| inf_count += int(torch.isinf(chunk).sum().item()) |
| if hs.ndim >= 3: |
| bad_layers = (~finite).flatten(start_dim=2).any(dim=(0, 2)) |
| affected_layers.update( |
| bad_layers.nonzero(as_tuple=False).flatten().tolist() |
| ) |
|
|
| if nan_count or inf_count: |
| details = ( |
| f"shape={tuple(hs.shape)}, dtype={hs.dtype}, " |
| f"nan_count={nan_count}, inf_count={inf_count}" |
| ) |
| if affected_layers: |
| details += f", affected layer slots={sorted(affected_layers)}" |
| raise ValueError(f"Hidden states contain non-finite values ({details})") |
|
|
|
|
| def get_existing_hidden_state_indices(output_path: Path) -> list[int]: |
| """Find existing `hs_i.safetensors` files (where i is the file index)""" |
|
|
| existing_file_indices_set: set[int] = set() |
|
|
| if not output_path.exists(): |
| return [] |
|
|
| for file_path in output_path.iterdir(): |
| if file_path.name.startswith("hs_") and file_path.name.endswith(".safetensors"): |
| index_str = file_path.stem[3:] |
| try: |
| file_index = int(index_str) |
| existing_file_indices_set.add(file_index) |
| except ValueError: |
| continue |
|
|
| return sorted(existing_file_indices_set) |
|
|
|
|
| def get_indices_to_process( |
| num_samples: int, |
| max_samples: int | None, |
| existing: list[int], |
| world_size: int, |
| rank: int, |
| ) -> list[int]: |
| """Determines which indices should be processed. If max_samples is None |
| returns all dataset indices not in existing. Otherwise gets the first |
| `max_samples - len(existing)` samples not already in existing. |
| |
| Args: |
| num_samples: Total size of preprocessed dataset |
| max_samples: (Optional) limit for number of samples to process |
| existing: list of ids that have already been processed |
| world_size: Number of nodes to generate on |
| rank: The rank of the local node |
| |
| Returns: |
| list of dataset indices to process |
| """ |
|
|
| target = min(max_samples, num_samples) if max_samples is not None else num_samples |
|
|
| if target <= 0: |
| return [] |
|
|
| chunk_size = target // world_size |
| remainder = target % world_size |
| |
| |
| start = rank * chunk_size + min(rank, remainder) |
| end = start + chunk_size + (1 if rank < remainder else 0) |
|
|
| existing_s = set(existing) |
| to_process = [i for i in range(start, end) if i not in existing_s] |
|
|
| if not to_process: |
| logger.info("All samples for this rank already processed!") |
| return [] |
|
|
| if len(existing_s & set(range(start, end))) > 0: |
| logger.info( |
| f"Found {len(existing_s & set(range(start, end)))} existing samples" |
| f" for rank {rank}." |
| ) |
|
|
| return to_process |
|
|
