Synthyra/ESM2-8M
4581
1from __future__ import annotations
2
3import torch
4import torch._inductor.config as inductor_config
5import torch._dynamo as dynamo
6
7# Enable TensorFloat32 tensor cores for float32 matmul (Ampere+ GPUs)
8# Provides significant speedup with minimal precision loss
9torch.set_float32_matmul_precision('high')
10
11# Enable TF32 for matrix multiplications and cuDNN operations
12torch.backends.cuda.matmul.allow_tf32 = True
13torch.backends.cudnn.allow_tf32 = True
14
15# Enable cuDNN autotuner - finds fastest algorithms for your hardware
16# Best when input sizes are consistent; may slow down first iterations
17torch.backends.cudnn.benchmark = True
18
19# Deterministic operations off for speed (set True if reproducibility needed)
20torch.backends.cudnn.deterministic = False
21inductor_config.max_autotune_gemm_backends = "ATEN,CUTLASS,FBGEMM"
22
23dynamo.config.capture_scalar_outputs = True
24torch._dynamo.config.recompile_limit = 16
25
26import io
27import os
28import queue
29import sqlite3
30import struct
31import threading
32import time
33
34import networkx as nx
35import numpy as np
36import torch
37from tqdm.auto import tqdm
38from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Tuple
39from torch.utils.data import DataLoader
40from torch.utils.data import Dataset as TorchDataset
41from transformers import PreTrainedTokenizerBase
42
43
44# SQLite stores tensors as compact blobs. Keep this header format compatible
45# with Protify readers that share the same dtype/version codes.
46_COMPACT_VERSION = 0x01
47_DTYPE_TO_CODE = {torch.float16: 0, torch.bfloat16: 1, torch.float32: 2}
48_CODE_TO_DTYPE = {0: torch.float16, 1: torch.bfloat16, 2: torch.float32}
49_CODE_TO_NP_DTYPE = {0: np.float16, 1: np.float16, 2: np.float32}
50
51
52def tensor_to_embedding_blob(tensor: torch.Tensor) -> bytes:
53 """Serialize a tensor to compact binary format for SQLite blob storage.
54
55 Format: [version:1][dtype_code:1][ndim:4][shape:4*ndim][raw_bytes]
56 bfloat16 tensors are stored as float16 bytes (numpy lacks bfloat16)
57 but tagged with dtype_code=1 so they can be cast back on read.
58 Falls back to torch.save for unsupported dtypes.
59 """
60 t = tensor.cpu()
61 if t.dtype not in _DTYPE_TO_CODE:
62 buffer = io.BytesIO()
63 torch.save(t, buffer)
64 return buffer.getvalue()
65 dtype_code = _DTYPE_TO_CODE[t.dtype]
66
67 if t.dtype == torch.bfloat16:
68 raw = t.half().numpy().tobytes()
69 else:
70 raw = t.numpy().tobytes()
71
72 shape = t.shape
73 header = struct.pack(f'<BBi{len(shape)}i', _COMPACT_VERSION, dtype_code, len(shape), *shape)
74 return header + raw
75
76
77def _compact_header(dtype: torch.dtype, shape: tuple) -> bytes:
78 """Build just the compact header for a given dtype and shape."""
79 dtype_code = _DTYPE_TO_CODE[dtype]
80 return struct.pack(f'<BBi{len(shape)}i', _COMPACT_VERSION, dtype_code, len(shape), *shape)
81
82
83def batch_tensor_to_blobs(batch: torch.Tensor) -> List[bytes]:
84 """Serialize a batch of same-shape tensors to compact blobs (fast path for vectors).
85
86 Builds the header once and slices raw bytes per row. Much faster than
87 per-row tensor_to_embedding_blob calls for uniform-shape batches.
88 """
89 assert batch.ndim >= 2, f"Expected batch with >= 2 dims, got {batch.ndim}"
90 t = batch.cpu()
91 store_dtype = t.dtype
92 if t.dtype not in _DTYPE_TO_CODE:
93 return [tensor_to_embedding_blob(t[i]) for i in range(t.shape[0])]
94
95 if t.dtype == torch.bfloat16:
96 arr = t.half().numpy()
97 store_dtype = torch.bfloat16
98 else:
99 arr = t.numpy()
100
101 row_shape = tuple(t.shape[1:])
102 header = _compact_header(store_dtype, row_shape)
103 raw = arr.tobytes()
104 stride = len(raw) // t.shape[0]
105 return [header + raw[i * stride:(i + 1) * stride] for i in range(t.shape[0])]
106
107
108def embedding_blob_to_tensor(blob: bytes, fallback_shape: Optional[Tuple[int, ...]] = None) -> torch.Tensor:
109 """Deserialize a blob back to a tensor. Auto-detects compact vs legacy formats."""
110 if len(blob) >= 6 and blob[0] == _COMPACT_VERSION:
111 dtype_code = blob[1]
112 ndim = struct.unpack_from('<i', blob, 2)[0]
113 shape = struct.unpack_from(f'<{ndim}i', blob, 6)
114 data_offset = 6 + 4 * ndim
115 np_dtype = _CODE_TO_NP_DTYPE[dtype_code]
116 arr = np.frombuffer(blob, dtype=np_dtype, offset=data_offset).copy().reshape(shape)
117 t = torch.from_numpy(arr)
118 target_dtype = _CODE_TO_DTYPE[dtype_code]
119 if target_dtype != t.dtype:
120 t = t.to(target_dtype)
121 return t
122
123 # Older `.pth`-style blobs were written with torch.save.
124 try:
125 buffer = io.BytesIO(blob)
126 return torch.load(buffer, map_location='cpu', weights_only=True)
127 except Exception:
128 pass
129
130 # Oldest SQLite rows stored raw float32 bytes and need a caller-supplied shape.
131 assert fallback_shape is not None, "Cannot deserialize blob: unknown format and no fallback_shape provided."
132 arr = np.frombuffer(blob, dtype=np.float32).copy().reshape(fallback_shape)
133 return torch.from_numpy(arr)
134
135
136def select_hidden_state_embeddings(
137 last_hidden_state: torch.Tensor,
138 hidden_states: Optional[Tuple[torch.Tensor, ...]],
139 hidden_state_index: int = -1,
140 store_all_hidden_states: bool = False,
141) -> torch.Tensor:
142 assert isinstance(hidden_state_index, int), "hidden_state_index must be an integer."
143 if store_all_hidden_states:
144 assert hidden_states is not None, "store_all_hidden_states requires output_hidden_states=True."
145 assert len(hidden_states) > 0, "Model returned no hidden states."
146 return torch.stack(tuple(hidden_states), dim=1)
147 if hidden_state_index == -1:
148 return last_hidden_state
149 assert hidden_states is not None, "hidden_state_index selection requires output_hidden_states=True."
150 return hidden_states[hidden_state_index]
151
152
153def _trim_full_embedding(embedding: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
154 mask = attention_mask.bool()
155 if embedding.ndim == 2:
156 return embedding[mask].reshape(-1, embedding.shape[-1])
157 if embedding.ndim == 3:
158 return embedding[:, mask, :].reshape(embedding.shape[0], -1, embedding.shape[-1])
159 raise AssertionError(f"Expected full embedding tensor with 2 or 3 dims, got {embedding.ndim}.")
160
161
162def pool_embeddings(
163 embeddings: Dict[str, torch.Tensor],
164 pooling_types: List[str] = ['mean'],
165 hidden_state_index: int = -1,
166) -> Dict[str, torch.Tensor]:
167 pooler = Pooler(pooling_types)
168 pooled: Dict[str, torch.Tensor] = {}
169 for sequence, embedding in embeddings.items():
170 assert isinstance(sequence, str), "Expected embedding dictionary keys to be sequences (str)."
171 assert isinstance(embedding, torch.Tensor), "Expected embedding dictionary values to be tensors."
172 if embedding.ndim == 1:
173 pooled[sequence] = embedding.cpu()
174 continue
175 if embedding.ndim == 3:
176 embedding = embedding[hidden_state_index]
177 assert embedding.ndim == 2, f"Expected token-wise embedding with 2 dims, got {embedding.ndim}."
178 pooled[sequence] = pooler(embedding.unsqueeze(0)).squeeze(0).cpu()
179 return pooled
180
181
182def load_pooled_embeddings_from_pth(
183 save_path: str,
184 pooling_types: List[str] = ['mean'],
185 hidden_state_index: int = -1,
186) -> Dict[str, torch.Tensor]:
187 assert os.path.exists(save_path), f"Embedding file does not exist: {save_path}"
188 payload = torch.load(save_path, map_location="cpu", weights_only=True)
189 assert isinstance(payload, dict), "Expected .pth embeddings file to contain a dictionary."
190 return pool_embeddings(payload, pooling_types=pooling_types, hidden_state_index=hidden_state_index)
191
192
193def load_pooled_embeddings_from_db(
194 db_path: str,
195 sequences: Optional[List[str]] = None,
196 pooling_types: List[str] = ['mean'],
197 hidden_state_index: int = -1,
198) -> Dict[str, torch.Tensor]:
199 assert os.path.exists(db_path), f"Embedding database does not exist: {db_path}"
200 loaded: Dict[str, torch.Tensor] = {}
201 with sqlite3.connect(db_path, timeout=30) as conn:
202 cursor = conn.cursor()
203 if sequences is None:
204 cursor.execute("SELECT sequence, embedding FROM embeddings")
205 else:
206 if len(sequences) == 0:
207 return loaded
208 placeholders = ",".join(["?"] * len(sequences))
209 cursor.execute(
210 f"SELECT sequence, embedding FROM embeddings WHERE sequence IN ({placeholders})",
211 tuple(sequences),
212 )
213 for sequence, embedding_bytes in cursor.fetchall():
214 loaded[sequence] = embedding_blob_to_tensor(embedding_bytes)
215 return pool_embeddings(loaded, pooling_types=pooling_types, hidden_state_index=hidden_state_index)
216
217
218def maybe_compile(model: torch.nn.Module, dynamic: bool = False) -> torch.nn.Module:
219 """Compile model with torch.compile if possible.
220
221 Skips compilation when dynamic=True (padding='longest') because
222 flex attention's create_block_mask is incompatible with dynamic shapes
223 under torch.compile, causing CUDA illegal memory access.
224 """
225 if dynamic:
226 print("Skipping torch.compile (dynamic shapes + flex attention incompatible)")
227 return model
228 try:
229 model = torch.compile(model)
230 print("Model compiled")
231 except Exception as e:
232 print(f"Skipping torch.compile: {e}")
233 return model
234
235
236def build_collator(
237 tokenizer: PreTrainedTokenizerBase,
238 padding: str = 'max_length',
239 max_length: int = 512,
240) -> Callable[[List[str]], Dict[str, torch.Tensor]]:
241 def _collate_fn(sequences: List[str]) -> Dict[str, torch.Tensor]:
242 kwargs: Dict[str, Any] = dict(
243 return_tensors="pt", padding=padding, truncation=True, max_length=max_length,
244 )
245 if padding != 'max_length':
246 kwargs['pad_to_multiple_of'] = 8
247 return tokenizer(sequences, **kwargs)
248 return _collate_fn
249
250
251def _make_embedding_progress(
252 dataloader: DataLoader,
253 padding: str,
254 n_warmup: int = 3,
255 n_calibration: int = 5,
256) -> Iterator[Tuple[int, Any]]:
257 """Progress-bar wrapper for embedding loops. Drop-in replacement for enumerate(dataloader).
258
259 When padding='max_length', all batches have uniform cost so plain tqdm works.
260 When padding='longest' (sorted longest-first), batch times vary dramatically.
261 In that case: yield warmup batches first (compiler warmup + OOM check on longest
262 sequences), then time mid-length calibration batches to estimate total ETA.
263
264 Keep in sync with protify/embedder.py and core/atlas/precomputed.py.
265 """
266 total = len(dataloader)
267 if padding == 'max_length' or total <= n_warmup + n_calibration:
268 for i, batch in tqdm(enumerate(dataloader), total=total, desc='Embedding batches'):
269 yield i, batch
270 return
271
272 dl_iter = iter(dataloader)
273
274 # Warm up on the longest batches first; sorted inputs make these the OOM-risk
275 # and compile-stabilization cases.
276 warmup_bar = tqdm(range(n_warmup), desc='Warmup (longest batches)', leave=False)
277 for i in warmup_bar:
278 batch = next(dl_iter)
279 yield i, batch
280 warmup_bar.close()
281
282 # Move toward mid-length batches for ETA calibration, yielding every real
283 # batch on the way so no sequences are skipped.
284 mid_start = total // 2
285 intermediate_bar = tqdm(
286 range(n_warmup, mid_start), desc='Embedding batches', leave=False,
287 )
288 for i in intermediate_bar:
289 batch = next(dl_iter)
290 yield i, batch
291 intermediate_bar.close()
292
293 # Mid-length batches give a better remaining-time estimate than the longest
294 # warmup batches.
295 calibration_times: List[float] = []
296 cal_bar = tqdm(range(n_calibration), desc='Calibrating ETA', leave=False)
297 for j in cal_bar:
298 t0 = time.perf_counter()
299 batch = next(dl_iter)
300 yield mid_start + j, batch
301 calibration_times.append(time.perf_counter() - t0)
302 cal_bar.close()
303
304 avg_time = sum(calibration_times) / len(calibration_times)
305 remaining_start = mid_start + n_calibration
306 remaining_count = total - remaining_start
307 estimated_total_seconds = avg_time * remaining_count
308
309 # Finish the tail with the calibrated ETA shown in the progress bar.
310 main_bar = tqdm(
311 range(remaining_count),
312 desc='Embedding batches',
313 bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]',
314 )
315 main_bar.set_postfix_str(f'ETA ~{estimated_total_seconds:.0f}s (calibrated)')
316 for k in main_bar:
317 batch = next(dl_iter)
318 yield remaining_start + k, batch
319 main_bar.close()
320
321
322class _SQLWriter:
323 """Context manager for async SQL embedding writes. Matches core/embed/storage.SQLEmbeddingWriter."""
324
325 def __init__(self, conn: sqlite3.Connection, queue_maxsize: int = 4) -> None:
326 self._conn = conn
327 self._queue: queue.Queue = queue.Queue(maxsize=queue_maxsize)
328 self._thread: Optional[threading.Thread] = None
329
330 def __enter__(self) -> "_SQLWriter":
331 self._thread = threading.Thread(target=self._writer_loop, daemon=True)
332 self._thread.start()
333 return self
334
335 def write_batch(self, rows: List[Tuple[str, bytes]]) -> None:
336 self._queue.put(rows)
337
338 def _writer_loop(self) -> None:
339 cursor = self._conn.cursor()
340 while True:
341 item = self._queue.get()
342 if item is None:
343 break
344 cursor.executemany("INSERT OR REPLACE INTO embeddings VALUES (?, ?)", item)
345 if self._queue.qsize() == 0:
346 self._conn.commit()
347 self._conn.commit()
348
349 def __exit__(self, *exc) -> None:
350 if self._thread is not None:
351 self._queue.put(None)
352 self._thread.join()
353 self._thread = None
354
355
356class Pooler:
357 def __init__(self, pooling_types: List[str]) -> None:
358 self.pooling_types = pooling_types
359 self.pooling_options: Dict[str, Callable] = {
360 'mean': self.mean_pooling,
361 'max': self.max_pooling,
362 'norm': self.norm_pooling,
363 'median': self.median_pooling,
364 'std': self.std_pooling,
365 'var': self.var_pooling,
366 'cls': self.cls_pooling,
367 'parti': self._pool_parti,
368 }
369
370 def _create_pooled_matrices_across_layers(self, attentions: torch.Tensor) -> torch.Tensor:
371 assert isinstance(attentions, torch.Tensor)
372 maxed_attentions = torch.max(attentions, dim=1)[0]
373 return maxed_attentions
374
375 def _page_rank(self, attention_matrix: np.ndarray, personalization: Optional[dict] = None, nstart: Optional[dict] = None, prune_type: str = "top_k_outdegree") -> Dict[int, float]:
376 G = self._convert_to_graph(attention_matrix)
377 if G.number_of_nodes() != attention_matrix.shape[0]:
378 raise Exception(
379 f"The number of nodes in the graph should be equal to the number of tokens in sequence! You have {G.number_of_nodes()} nodes for {attention_matrix.shape[0]} tokens.")
380 if G.number_of_edges() == 0:
381 raise Exception(f"You don't seem to have any attention edges left in the graph.")
382
383 return nx.pagerank(G, alpha=0.85, tol=1e-06, weight='weight', personalization=personalization, nstart=nstart, max_iter=100)
384
385 def _convert_to_graph(self, matrix: np.ndarray) -> nx.DiGraph:
386 G = nx.from_numpy_array(matrix, create_using=nx.DiGraph)
387 return G
388
389 def _calculate_importance_weights(self, dict_importance: Dict[int, float], attention_mask: Optional[torch.Tensor] = None) -> np.ndarray:
390 if attention_mask is not None:
391 for k in list(dict_importance.keys()):
392 if attention_mask[k] == 0:
393 del dict_importance[k]
394
395 total = sum(dict_importance.values())
396 return np.array([v / total for _, v in dict_importance.items()])
397
398 def _pool_parti(self, emb: torch.Tensor, attentions: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
399 maxed_attentions = self._create_pooled_matrices_across_layers(attentions).numpy()
400 emb_pooled = []
401 for e, a, mask in zip(emb, maxed_attentions, attention_mask):
402 dict_importance = self._page_rank(a)
403 importance_weights = self._calculate_importance_weights(dict_importance, mask)
404 num_tokens = int(mask.sum().item())
405 emb_pooled.append(np.average(e[:num_tokens], weights=importance_weights, axis=0))
406 pooled = torch.tensor(np.array(emb_pooled))
407 return pooled
408
409 def mean_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
410 if attention_mask is None:
411 return emb.mean(dim=1)
412 else:
413 attention_mask = attention_mask.unsqueeze(-1)
414 return (emb * attention_mask).sum(dim=1) / attention_mask.sum(dim=1)
415
416 def max_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
417 if attention_mask is None:
418 return emb.max(dim=1).values
419 else:
420 mask = attention_mask.unsqueeze(-1).bool()
421 return emb.masked_fill(~mask, float('-inf')).max(dim=1).values
422
423 def norm_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
424 if attention_mask is None:
425 return emb.norm(dim=1, p=2)
426 else:
427 attention_mask = attention_mask.unsqueeze(-1)
428 return (emb * attention_mask).norm(dim=1, p=2)
429
430 def median_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
431 if attention_mask is None:
432 return emb.median(dim=1).values
433 else:
434 mask = attention_mask.unsqueeze(-1).bool()
435 return emb.masked_fill(~mask, float('nan')).nanmedian(dim=1).values
436
437 def std_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
438 if attention_mask is None:
439 return emb.std(dim=1)
440 else:
441 var = self.var_pooling(emb, attention_mask, **kwargs)
442 return torch.sqrt(var)
443
444 def var_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
445 if attention_mask is None:
446 return emb.var(dim=1)
447 else:
448 attention_mask = attention_mask.unsqueeze(-1)
449 mean = (emb * attention_mask).sum(dim=1) / attention_mask.sum(dim=1)
450 mean = mean.unsqueeze(1)
451 squared_diff = (emb - mean) ** 2
452 var = (squared_diff * attention_mask).sum(dim=1) / attention_mask.sum(dim=1)
453 return var
454
455 def cls_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
456 return emb[:, 0, :]
457
458 def __call__(
459 self,
460 emb: torch.Tensor,
461 attention_mask: Optional[torch.Tensor] = None,
462 attentions: Optional[torch.Tensor] = None
463 ) -> torch.Tensor:
464 if attention_mask is not None:
465 assert attention_mask.sum(dim=-1).min() > 0, (
466 "Pooler received samples with all-zero attention masks. "
467 "This causes NaN from division by zero. Filter empty inputs before pooling."
468 )
469 final_emb: List[torch.Tensor] = []
470 for pooling_type in self.pooling_types:
471 final_emb.append(self.pooling_options[pooling_type](emb=emb, attention_mask=attention_mask, attentions=attentions))
472 return torch.cat(final_emb, dim=-1)
473
474
475class ProteinDataset(TorchDataset):
476 """Simple dataset for protein sequences."""
477 def __init__(self, sequences: List[str]) -> None:
478 self.sequences = sequences
479
480 def __len__(self) -> int:
481 return len(self.sequences)
482
483 def __getitem__(self, idx: int) -> str:
484 return self.sequences[idx]
485
486
487def parse_fasta(fasta_path: str) -> List[str]:
488 assert os.path.exists(fasta_path), f"FASTA file does not exist: {fasta_path}"
489 sequences = []
490 current_seq = []
491 with open(fasta_path, 'r') as f:
492 for line in f:
493 line = line.strip()
494 if not line:
495 continue
496 if line.startswith('>'):
497 if current_seq:
498 sequences.append(''.join(current_seq))
499 current_seq = []
500 else:
501 current_seq.append(line)
502 if current_seq:
503 sequences.append(''.join(current_seq))
504 return sequences
505
506
507class EmbeddingMixin:
508 def _embed(
509 self,
510 input_ids: torch.Tensor,
511 attention_mask: Optional[torch.Tensor] = None,
512 hidden_state_index: int = -1,
513 store_all_hidden_states: bool = False,
514 ) -> torch.Tensor:
515 raise NotImplementedError
516
517 @property
518 def device(self) -> torch.device:
519 """Get the device of the model."""
520 return next(self.parameters()).device
521
522 def _read_sequences_from_db(self, db_path: str) -> Set[str]:
523 """Read sequences from SQLite database."""
524 with sqlite3.connect(db_path, timeout=30) as conn:
525 c = conn.cursor()
526 c.execute("SELECT sequence FROM embeddings")
527 return {row[0] for row in c.fetchall()}
528
529 def _ensure_embeddings_table(self, conn: sqlite3.Connection) -> None:
530 cursor = conn.cursor()
531 cursor.execute(
532 "CREATE TABLE IF NOT EXISTS embeddings ("
533 "sequence TEXT PRIMARY KEY, "
534 "embedding BLOB NOT NULL"
535 ")"
536 )
537 conn.commit()
538
539 def load_embeddings_from_pth(self, save_path: str) -> Dict[str, torch.Tensor]:
540 assert os.path.exists(save_path), f"Embedding file does not exist: {save_path}"
541 payload = torch.load(save_path, map_location="cpu", weights_only=True)
542 assert isinstance(payload, dict), "Expected .pth embeddings file to contain a dictionary."
543 for sequence, tensor in payload.items():
544 assert isinstance(sequence, str), "Expected embedding dictionary keys to be sequences (str)."
545 assert isinstance(tensor, torch.Tensor), "Expected embedding dictionary values to be tensors."
546 return payload
547
548 def load_embeddings_from_db(self, db_path: str, sequences: Optional[List[str]] = None) -> Dict[str, torch.Tensor]:
549 assert os.path.exists(db_path), f"Embedding database does not exist: {db_path}"
550 loaded: Dict[str, torch.Tensor] = {}
551 with sqlite3.connect(db_path, timeout=30) as conn:
552 self._ensure_embeddings_table(conn)
553 cursor = conn.cursor()
554 if sequences is None:
555 cursor.execute("SELECT sequence, embedding FROM embeddings")
556 else:
557 if len(sequences) == 0:
558 return loaded
559 placeholders = ",".join(["?"] * len(sequences))
560 cursor.execute(
561 f"SELECT sequence, embedding FROM embeddings WHERE sequence IN ({placeholders})",
562 tuple(sequences),
563 )
564
565 rows = cursor.fetchall()
566 for row in rows:
567 sequence = row[0]
568 embedding_bytes = row[1]
569 loaded[sequence] = embedding_blob_to_tensor(embedding_bytes)
570 return loaded
571
572 def pool_embeddings(
573 self,
574 embeddings: Dict[str, torch.Tensor],
575 pooling_types: List[str] = ['mean'],
576 hidden_state_index: int = -1,
577 ) -> Dict[str, torch.Tensor]:
578 return pool_embeddings(embeddings, pooling_types=pooling_types, hidden_state_index=hidden_state_index)
579
580 def load_pooled_embeddings_from_pth(
581 self,
582 save_path: str,
583 pooling_types: List[str] = ['mean'],
584 hidden_state_index: int = -1,
585 ) -> Dict[str, torch.Tensor]:
586 return load_pooled_embeddings_from_pth(
587 save_path,
588 pooling_types=pooling_types,
589 hidden_state_index=hidden_state_index,
590 )
591
592 def load_pooled_embeddings_from_db(
593 self,
594 db_path: str,
595 sequences: Optional[List[str]] = None,
596 pooling_types: List[str] = ['mean'],
597 hidden_state_index: int = -1,
598 ) -> Dict[str, torch.Tensor]:
599 return load_pooled_embeddings_from_db(
600 db_path,
601 sequences=sequences,
602 pooling_types=pooling_types,
603 hidden_state_index=hidden_state_index,
604 )
605
606 def embed_dataset(
607 self,
608 sequences: Optional[List[str]] = None,
609 tokenizer: Optional[PreTrainedTokenizerBase] = None,
610 batch_size: int = 2,
611 max_len: int = 512,
612 truncate: bool = True,
613 full_embeddings: bool = False,
614 embed_dtype: torch.dtype = torch.float32,
615 pooling_types: List[str] = ['mean'],
616 num_workers: int = 0,
617 sql: bool = False,
618 save: bool = True,
619 sql_db_path: str = 'embeddings.db',
620 save_path: str = 'embeddings.pth',
621 fasta_path: Optional[str] = None,
622 padding: str = 'max_length',
623 hidden_state_index: int = -1,
624 store_all_hidden_states: bool = False,
625 **kwargs,
626 ) -> Optional[Dict[str, torch.Tensor]]:
627 """
628 Embed a dataset of protein sequences.
629
630 Supports two modes:
631 - Tokenizer mode (ESM2/ESM++): provide `tokenizer` or use `self.tokenizer`.
632 - Sequence mode (E1): pass `tokenizer=None`, `_embed(sequences, return_attention_mask=True, **kwargs)` is used.
633
634 Sequences can be supplied as a list via `sequences`, parsed from a FASTA file via
635 `fasta_path`, or both (the two sources are combined). At least one must be provided.
636 """
637 if fasta_path is not None:
638 fasta_sequences = parse_fasta(fasta_path)
639 sequences = list(sequences or []) + fasta_sequences
640 assert sequences is not None and len(sequences) > 0, \
641 "Must provide at least one sequence via `sequences` or `fasta_path`."
642 assert isinstance(hidden_state_index, int), "hidden_state_index must be an integer."
643 assert full_embeddings or not store_all_hidden_states, \
644 "store_all_hidden_states=True requires full_embeddings=True."
645 sequences = list(set([seq[:max_len] if truncate else seq for seq in sequences]))
646 sequences = sorted(sequences, key=len, reverse=True)
647 pooler = Pooler(pooling_types) if not full_embeddings else None
648 if tokenizer is None and self.config.model_type != "E1":
649 tokenizer = self.tokenizer
650 tokenizer_mode = tokenizer is not None
651
652 # Resolve padding and compilation
653 dynamic = padding == 'longest'
654 compiled_model = maybe_compile(self, dynamic=dynamic)
655
656 if tokenizer_mode:
657 collate_fn = build_collator(tokenizer, padding=padding, max_length=max_len)
658 device = self.device
659 else:
660 collate_fn = None
661 device = None
662
663 def get_embeddings(residue_embeddings: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
664 assert isinstance(residue_embeddings, torch.Tensor)
665 if full_embeddings or residue_embeddings.ndim == 2:
666 return residue_embeddings
667 return pooler(residue_embeddings, attention_mask)
668
669 def iter_batches(to_embed: List[str]):
670 if tokenizer_mode:
671 assert collate_fn is not None
672 assert device is not None
673 dataset = ProteinDataset(to_embed)
674 dataloader = DataLoader(
675 dataset,
676 batch_size=batch_size,
677 num_workers=num_workers,
678 prefetch_factor=2 if num_workers > 0 else None,
679 collate_fn=collate_fn,
680 shuffle=False,
681 pin_memory=True,
682 )
683 for i, batch in _make_embedding_progress(dataloader, padding):
684 seqs = to_embed[i * batch_size:(i + 1) * batch_size]
685 input_ids = batch['input_ids'].to(device)
686 attention_mask = batch['attention_mask'].to(device)
687 residue_embeddings = compiled_model._embed(
688 input_ids,
689 attention_mask,
690 hidden_state_index=hidden_state_index,
691 store_all_hidden_states=store_all_hidden_states,
692 )
693 yield seqs, residue_embeddings, attention_mask
694 else:
695 for batch_start in tqdm(range(0, len(to_embed), batch_size), desc='Embedding batches'):
696 seqs = to_embed[batch_start:batch_start + batch_size]
697 batch_output = compiled_model._embed(
698 seqs,
699 return_attention_mask=True,
700 hidden_state_index=hidden_state_index,
701 store_all_hidden_states=store_all_hidden_states,
702 **kwargs,
703 )
704 assert isinstance(batch_output, tuple), "Sequence mode _embed must return (last_hidden_state, attention_mask)."
705 assert len(batch_output) == 2, "Sequence mode _embed must return exactly two values."
706 residue_embeddings, attention_mask = batch_output
707 assert isinstance(attention_mask, torch.Tensor), "Sequence mode _embed must return attention_mask as a torch.Tensor."
708 yield seqs, residue_embeddings, attention_mask
709
710 if sql:
711 # Resume safely: skip sequences already present in the SQLite table.
712 conn = sqlite3.connect(sql_db_path, timeout=30, check_same_thread=False)
713 conn.execute('PRAGMA journal_mode=WAL')
714 conn.execute('PRAGMA busy_timeout=30000')
715 conn.execute('PRAGMA synchronous=OFF')
716 conn.execute('PRAGMA cache_size=-64000')
717 self._ensure_embeddings_table(conn)
718 already_embedded = self._read_sequences_from_db(sql_db_path)
719 to_embed = [seq for seq in sequences if seq not in already_embedded]
720 print(f"Found {len(already_embedded)} already embedded sequences in {sql_db_path}")
721 print(f"Embedding {len(to_embed)} new sequences")
722 if len(to_embed) > 0:
723 # Embed batches synchronously; serialize/write them on the SQL writer thread.
724 with _SQLWriter(conn) as writer:
725 with torch.inference_mode():
726 for seqs, residue_embeddings, attention_mask in iter_batches(to_embed):
727 embeddings = get_embeddings(residue_embeddings, attention_mask).to(embed_dtype)
728 if full_embeddings:
729 batch_rows = []
730 for seq, emb, mask in zip(seqs, embeddings, attention_mask):
731 batch_rows.append((seq, tensor_to_embedding_blob(_trim_full_embedding(emb, mask))))
732 else:
733 blobs = batch_tensor_to_blobs(embeddings)
734 batch_rows = list(zip(seqs, blobs))
735 writer.write_batch(batch_rows)
736 conn.close()
737 return None
738
739 embeddings_dict = {}
740 if os.path.exists(save_path):
741 embeddings_dict = self.load_embeddings_from_pth(save_path)
742 to_embed = [seq for seq in sequences if seq not in embeddings_dict]
743 print(f"Found {len(embeddings_dict)} already embedded sequences in {save_path}")
744 print(f"Embedding {len(to_embed)} new sequences")
745 else:
746 to_embed = sequences
747 print(f"Embedding {len(to_embed)} new sequences")
748
749 if len(to_embed) > 0:
750 with torch.inference_mode():
751 for seqs, residue_embeddings, attention_mask in iter_batches(to_embed):
752 embeddings = get_embeddings(residue_embeddings, attention_mask).to(embed_dtype)
753 for seq, emb, mask in zip(seqs, embeddings, attention_mask):
754 if full_embeddings:
755 emb = _trim_full_embedding(emb, mask)
756 embeddings_dict[seq] = emb.cpu()
757
758 if save:
759 torch.save(embeddings_dict, save_path)
760
761 return embeddings_dict
762
763
764if __name__ == "__main__":
765 # Manual smoke test for pooling shape behavior.
766 pooler = Pooler(pooling_types=['max', 'parti'])
767 batch_size = 8
768 seq_len = 64
769 hidden_size = 128
770 num_layers = 12
771 emb = torch.randn(batch_size, seq_len, hidden_size)
772 attentions = torch.randn(batch_size, num_layers, seq_len, seq_len)
773 attention_mask = torch.ones(batch_size, seq_len)
774 y = pooler(emb=emb, attention_mask=attention_mask, attentions=attentions)
775 print(y.shape)
776
777"""Shared attention infrastructure for all FastPLMs models.
778
779Contains: AttentionBackend enum, backend resolution, mask creation,
780flex attention helpers, flash kernel detection/dispatch, and pad/unpad utilities.
781"""
782from enum import Enum
783from typing import Dict, List, Optional, Tuple
784
785import torch
786import torch.nn as nn
787from torch.nn import functional as F
788from einops import rearrange
789
790try:
791 from torch.nn.attention.flex_attention import create_block_mask, flex_attention, BlockMask
792except ImportError:
793 create_block_mask = None
794 flex_attention = None
795 BlockMask = None
796
797_compiled_flex_attention = None
798
799
800def _get_flex_attention_fn():
801 """Return flex_attention callable: compiled (fused kernel) by default, or eager when debug flag is set."""
802 global _compiled_flex_attention
803 if flex_attention is None:
804 return None
805 flex_mod = torch.nn.attention.flex_attention
806 if getattr(flex_mod, "_FLEX_ATTENTION_DISABLE_COMPILE_DEBUG", False):
807 return flex_attention
808 if _compiled_flex_attention is None:
809 _compiled_flex_attention = torch.compile(
810 flex_attention,
811 dynamic=False,
812 )
813 return _compiled_flex_attention
814
815
816# HuggingFace `kernels` exposes slightly different APIs for Flash Attention 2
817# and 3. Detect the loaded variant once so every caller uses the same dispatch.
818def _infer_kernels_flash_variant(kernel) -> Optional[str]:
819 if hasattr(kernel, "fwd") and hasattr(kernel, "varlen_fwd"):
820 return "flash_attn2"
821 if hasattr(kernel, "flash_attn_func") and hasattr(kernel, "flash_attn_varlen_func"):
822 return "flash_attn3"
823 return None
824
825
826def _try_get_kernels_flash():
827 try:
828 from kernels import get_kernel
829 except ImportError:
830 return None, None
831
832 flash_kernel = None
833 flash_kernel_variant = None
834 try:
835 flash_kernel = get_kernel("kernels-community/flash-attn3")
836 flash_kernel_variant = _infer_kernels_flash_variant(flash_kernel)
837 assert flash_kernel_variant is not None, "Loaded flash-attn3 kernel does not expose a supported API."
838 except Exception:
839 try:
840 flash_kernel = get_kernel("kernels-community/flash-attn2")
841 flash_kernel_variant = _infer_kernels_flash_variant(flash_kernel)
842 assert flash_kernel_variant is not None, "Loaded flash-attn2 kernel does not expose a supported API."
843 except Exception:
844 flash_kernel = None
845 flash_kernel_variant = None
846 return flash_kernel, flash_kernel_variant
847
848
849_FLASH_KERNELS_LOADED = False
850FLASH_KERNEL = None
851FLASH_KERNEL_VARIANT = None
852
853
854def _ensure_flash_kernels_loaded():
855 global _FLASH_KERNELS_LOADED, FLASH_KERNEL, FLASH_KERNEL_VARIANT
856 if _FLASH_KERNELS_LOADED:
857 return
858 _FLASH_KERNELS_LOADED = True
859 FLASH_KERNEL, FLASH_KERNEL_VARIANT = _try_get_kernels_flash()
860
861
862def _kernels_flash_forward(
863 query_states: torch.Tensor,
864 key_states: torch.Tensor,
865 value_states: torch.Tensor,
866 causal: bool = False,
867 softmax_scale: Optional[float] = None,
868) -> torch.Tensor:
869 """Flash-attention forward, optionally overriding the softmax scale.
870
871 When `softmax_scale is None`, the flash kernel applies its default
872 `1 / sqrt(head_dim)`. Pass `softmax_scale=1.0` if the caller has already
873 pre-scaled Q (the convention used by ESM2, DPLM, DPLM2, E1, ESMFold).
874 Failing to override when Q is pre-scaled applies the scale twice. On
875 DPLM-150M, that produced pooled-embedding cosine around -0.12 and argmax
876 agreement around 0.27 vs SDPA.
877 """
878 assert FLASH_KERNEL is not None, "Kernel Flash Attention is not available in this environment."
879 if FLASH_KERNEL_VARIANT == "flash_attn2":
880 return FLASH_KERNEL.fwd(
881 q=query_states, k=key_states, v=value_states,
882 softmax_scale=softmax_scale, is_causal=causal,
883 )[0]
884 if FLASH_KERNEL_VARIANT == "flash_attn3":
885 try:
886 output = FLASH_KERNEL.flash_attn_func(
887 q=query_states, k=key_states, v=value_states,
888 softmax_scale=softmax_scale, causal=causal,
889 )
890 except TypeError:
891 output = FLASH_KERNEL.flash_attn_func(
892 query_states, key_states, value_states,
893 0.0, softmax_scale, causal,
894 )
895 if isinstance(output, tuple):
896 return output[0]
897 return output
898 raise AssertionError(f"Unsupported kernels flash attention variant: {FLASH_KERNEL_VARIANT}")
899
900
901def _kernels_flash_varlen_forward(
902 query_states: torch.Tensor,
903 key_states: torch.Tensor,
904 value_states: torch.Tensor,
905 cu_seqlens_q: torch.Tensor,
906 cu_seqlens_k: torch.Tensor,
907 max_seqlen_in_batch_q: int,
908 max_seqlen_in_batch_k: int,
909 causal: bool = False,
910 softmax_scale: Optional[float] = None,
911) -> torch.Tensor:
912 """Varlen flash-attention forward, optionally overriding the softmax scale.
913
914 See `_kernels_flash_forward` docstring for why `softmax_scale=1.0` must be
915 passed when Q has been pre-scaled by the caller.
916 """
917 assert FLASH_KERNEL is not None, "Kernel Flash Attention is not available in this environment."
918 if FLASH_KERNEL_VARIANT == "flash_attn2":
919 return FLASH_KERNEL.varlen_fwd(
920 q=query_states, k=key_states, v=value_states,
921 cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k,
922 max_seqlen_q=max_seqlen_in_batch_q, max_seqlen_k=max_seqlen_in_batch_k,
923 softmax_scale=softmax_scale, is_causal=causal,
924 )[0]
925 if FLASH_KERNEL_VARIANT == "flash_attn3":
926 try:
927 output = FLASH_KERNEL.flash_attn_varlen_func(
928 q=query_states, k=key_states, v=value_states,
929 cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k,
930 max_seqlen_q=max_seqlen_in_batch_q, max_seqlen_k=max_seqlen_in_batch_k,
931 softmax_scale=softmax_scale, causal=causal,
932 )
933 except TypeError:
934 output = FLASH_KERNEL.flash_attn_varlen_func(
935 query_states, key_states, value_states,
936 cu_seqlens_q, cu_seqlens_k,
937 max_seqlen_in_batch_q, max_seqlen_in_batch_k,
938 0.0, softmax_scale, causal,
939 )
940 if isinstance(output, tuple):
941 return output[0]
942 return output
943 raise AssertionError(f"Unsupported kernels flash attention variant: {FLASH_KERNEL_VARIANT}")
944
945
946# Varlen flash attention runs only on real tokens. These helpers remove padding
947# before the kernel call and restore the original padded batch shape afterward.
948class IndexFirstAxis(torch.autograd.Function):
949 @staticmethod
950 def forward(ctx, input, indices) -> torch.Tensor:
951 ctx.save_for_backward(indices)
952 assert input.ndim >= 2
953 ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
954 second_dim = other_shape.numel()
955 return torch.gather(
956 rearrange(input, "b ... -> b (...)"), 0, indices.unsqueeze(1).expand(-1, second_dim)
957 ).reshape(-1, *other_shape)
958
959 @staticmethod
960 def backward(ctx, grad_output) -> Tuple[torch.Tensor, None]:
961 (indices,) = ctx.saved_tensors
962 assert grad_output.ndim >= 2
963 other_shape = grad_output.shape[1:]
964 grad_output = rearrange(grad_output, "b ... -> b (...)")
965 grad_input = torch.zeros(
966 [ctx.first_axis_dim, grad_output.shape[1]], device=grad_output.device, dtype=grad_output.dtype
967 )
968 grad_input.scatter_(0, indices.unsqueeze(1).expand(-1, grad_output.shape[1]), grad_output)
969 return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
970
971
972class IndexPutFirstAxis(torch.autograd.Function):
973 @staticmethod
974 def forward(ctx, values, indices, first_axis_dim) -> torch.Tensor:
975 ctx.save_for_backward(indices)
976 assert indices.ndim == 1
977 assert values.ndim >= 2
978 output = torch.zeros(first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype)
979 output[indices] = values
980 return output
981
982 @staticmethod
983 def backward(ctx, grad_output) -> Tuple[torch.Tensor, None, None]:
984 (indices,) = ctx.saved_tensors
985 return grad_output[indices], None, None
986
987
988index_first_axis = IndexFirstAxis.apply
989index_put_first_axis = IndexPutFirstAxis.apply
990
991
992def pad_input(hidden_states: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int) -> torch.Tensor:
993 output = index_put_first_axis(hidden_states, indices, batch * seqlen)
994 return rearrange(output, "(b s) ... -> b s ...", b=batch)
995
996
997def _unpad_input(
998 query_layer: torch.Tensor,
999 key_layer: torch.Tensor,
1000 value_layer: torch.Tensor,
1001 attention_mask_2d: torch.Tensor,
1002) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Tuple[torch.Tensor, torch.Tensor], Tuple[int, int]]:
1003 batch_size, seq_len, num_heads, head_dim = query_layer.shape
1004 seqlens = attention_mask_2d.sum(dim=1).int()
1005 cu_seqlens = F.pad(seqlens.cumsum(0, dtype=torch.int32), (1, 0))
1006 max_seqlen = int(seqlens.max().item())
1007 indices = attention_mask_2d.flatten().nonzero(as_tuple=False).flatten()
1008 query_layer = index_first_axis(query_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices)
1009 key_layer = index_first_axis(key_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices)
1010 value_layer = index_first_axis(value_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices)
1011 return query_layer, key_layer, value_layer, indices, (cu_seqlens, cu_seqlens), (max_seqlen, max_seqlen)
1012
1013
1014def kernels_flash_attention_func(
1015 query_states: torch.Tensor,
1016 key_states: torch.Tensor,
1017 value_states: torch.Tensor,
1018 attention_mask_2d: Optional[torch.Tensor] = None,
1019 causal: bool = False,
1020 softmax_scale: Optional[float] = None,
1021) -> torch.Tensor:
1022 """Public flash-attention entry point with optional padding handling.
1023
1024 `softmax_scale`:
1025 None -> kernel applies its default `1 / sqrt(head_dim)`.
1026 float -> kernel uses the given scale (pass 1.0 when Q is pre-scaled
1027 by the caller).
1028
1029 Caller contract: if a model family pre-scales Q by `1/sqrt(head_dim)`
1030 before calling this function (ESM2, DPLM, DPLM2, E1, and ESMFold do), pass
1031 `softmax_scale=1.0`. Otherwise the flash kernel applies its default scale
1032 again, yielding an effective `1/head_dim` scale that drifts across layers.
1033 """
1034 assert FLASH_KERNEL is not None, "Kernel Flash Attention is not available in this environment."
1035 if not causal and attention_mask_2d is not None:
1036 batch_size, q_len = query_states.shape[:2]
1037 (
1038 query_states, key_states, value_states,
1039 indices_q, (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k),
1040 ) = _unpad_input(query_states, key_states, value_states, attention_mask_2d)
1041 attn_output_unpad = _kernels_flash_varlen_forward(
1042 query_states=query_states, key_states=key_states, value_states=value_states,
1043 cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k,
1044 max_seqlen_in_batch_q=max_seqlen_q, max_seqlen_in_batch_k=max_seqlen_k,
1045 softmax_scale=softmax_scale,
1046 )
1047 return pad_input(attn_output_unpad, indices_q, batch_size, q_len)
1048 else:
1049 return _kernels_flash_forward(
1050 query_states=query_states, key_states=key_states, value_states=value_states,
1051 causal=causal, softmax_scale=softmax_scale,
1052 )
1053
1054
1055# User-facing backend strings resolve to this enum before attention dispatch.
1056class AttentionBackend(Enum):
1057 AUTO = "auto"
1058 KERNELS_FLASH = "kernels_flash"
1059 FLEX = "flex"
1060 SDPA = "sdpa"
1061
1062
1063VALID_ATTENTION_BACKENDS = tuple(b.value for b in AttentionBackend)
1064
1065
1066_BACKEND_CONFIRMED = False
1067
1068
1069def resolve_attention_backend(requested_backend: str) -> AttentionBackend:
1070 global _BACKEND_CONFIRMED
1071 assert requested_backend in VALID_ATTENTION_BACKENDS, (
1072 f"Unsupported attention backend: {requested_backend}. Expected one of {VALID_ATTENTION_BACKENDS}."
1073 )
1074 if requested_backend in (AttentionBackend.AUTO.value, AttentionBackend.KERNELS_FLASH.value):
1075 _ensure_flash_kernels_loaded()
1076 if requested_backend == AttentionBackend.AUTO.value:
1077 if FLASH_KERNEL is not None:
1078 resolved = AttentionBackend.KERNELS_FLASH
1079 elif flex_attention is not None:
1080 resolved = AttentionBackend.FLEX
1081 else:
1082 resolved = AttentionBackend.SDPA
1083 elif requested_backend == AttentionBackend.KERNELS_FLASH.value:
1084 assert FLASH_KERNEL is not None, "Kernels Flash Attention is not available in this environment."
1085 resolved = AttentionBackend.KERNELS_FLASH
1086 elif requested_backend == AttentionBackend.FLEX.value:
1087 assert flex_attention is not None, "Flex Attention is not available in this environment."
1088 resolved = AttentionBackend.FLEX
1089 elif requested_backend == AttentionBackend.SDPA.value:
1090 resolved = AttentionBackend.SDPA
1091 else:
1092 raise AssertionError(f"Unsupported attention backend: {requested_backend}")
1093 if not _BACKEND_CONFIRMED:
1094 print(f"Attention backend: config='{requested_backend}' -> resolved='{resolved.value}'")
1095 _BACKEND_CONFIRMED = True
1096 return resolved
1097
1098
1099@torch.compiler.disable
1100def get_attention_mask(
1101 effective_backend: AttentionBackend,
1102 batch_size: int,
1103 seq_len: int,
1104 device: torch.device,
1105 attention_mask: Optional[torch.Tensor] = None,
1106) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[BlockMask]]:
1107 """Build padding masks once for all encoder layers.
1108
1109 Returns (attention_mask_2d, attention_mask_4d, flex_block_mask).
1110 """
1111 if attention_mask is None:
1112 return None, None, None
1113
1114 attention_mask_2d = attention_mask.bool()
1115
1116 if effective_backend == AttentionBackend.KERNELS_FLASH:
1117 return attention_mask_2d, None, None
1118
1119 if effective_backend == AttentionBackend.FLEX:
1120 assert create_block_mask is not None, "Flex attention backend requested but torch.create_block_mask is unavailable."
1121 valid_lens = attention_mask_2d.sum(dim=-1)
1122
1123 def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
1124 return (q_idx < valid_lens[batch_idx]) & (kv_idx < valid_lens[batch_idx])
1125
1126 flex_block_mask = create_block_mask(mask_mod, batch_size, 1, seq_len, seq_len, device=device)
1127 return attention_mask_2d, None, flex_block_mask
1128
1129 # SDPA/manual masks only keys. Padding queries still attend to real keys, so
1130 # their outputs stay finite instead of softmaxing over all -inf scores.
1131 attention_mask_4d = attention_mask_2d[:, None, None, :]
1132 return attention_mask_2d, attention_mask_4d, None
1133
1134
1135def bool_to_additive_mask(
1136 bool_mask: torch.Tensor,
1137 dtype: torch.dtype,
1138) -> torch.Tensor:
1139 """Convert a bool mask (True = valid) to a float additive mask (0.0 valid, -inf invalid).
1140
1141 Why this exists: calling `bool_mask.masked_fill(bool_mask.logical_not(), float('-inf'))`
1142 directly on a bool tensor returns a bool tensor because `-inf` casts to `True`.
1143 That silently drops the mask. Always allocate a float tensor first, then fill it.
1144 This helper is the sanctioned way to build an SDPA additive mask from a bool validity mask.
1145 """
1146 assert bool_mask.dtype == torch.bool, (
1147 f"bool_to_additive_mask requires a bool tensor, got dtype={bool_mask.dtype}"
1148 )
1149 additive = torch.zeros_like(bool_mask, dtype=dtype)
1150 additive.masked_fill_(bool_mask.logical_not(), float("-inf"))
1151 return additive
1152
1153import typing as T
1154from dataclasses import dataclass, fields
1155
1156import torch
1157import torch.nn as nn
1158import torch.nn.functional as F
1159
1160
1161@dataclass
1162class TTTConfig:
1163 lr: float = 4e-4
1164 steps: int = 30
1165 ags: int = 16
1166 batch_size: int = 2
1167 mask_ratio: float = 0.15
1168 crop_size: int = 1024
1169 bert_leave_prob: float = 0.1
1170 bert_replace_prob: float = 0.1
1171 optimizer: str = "sgd"
1172 momentum: float = 0.0
1173 weight_decay: float = 0.0
1174 seed: int | None = 0
1175 lora_rank: int = 8
1176 lora_alpha: float = 32.0
1177 lora_target_replace_module: str | None = None
1178 lora_target_modules: tuple[str, ...] | None = None
1179 initial_state_reset: bool = True
1180 automatic_best_state_reset: bool = False
1181 eval_each_step: bool = False
1182 gradient_clip: bool = False
1183 gradient_clip_max_norm: float = 1.0
1184
1185 @classmethod
1186 def from_kwargs(cls, **kwargs: T.Any) -> "TTTConfig":
1187 valid_names = {field.name for field in fields(cls)}
1188 unknown_names = set(kwargs) - valid_names
1189 assert len(unknown_names) == 0, f"Unknown TTTConfig fields: {sorted(unknown_names)}"
1190 return cls(**kwargs)
1191
1192 def merged(self, overrides: T.Mapping[str, T.Any] | "TTTConfig" | None) -> "TTTConfig":
1193 if overrides is None:
1194 return self
1195 if isinstance(overrides, TTTConfig):
1196 return overrides
1197 values = {field.name: self.__dict__[field.name] for field in fields(self)}
1198 for name, value in overrides.items():
1199 assert name in values, f"Unknown TTTConfig field: {name}"
1200 values[name] = value