| """Frozen feature extraction wrappers for A1 baseline models.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import asdict, dataclass |
| import json |
| import os |
| from pathlib import Path |
| import re |
| from typing import Any |
| import warnings |
|
|
| import numpy as np |
| import pandas as pd |
|
|
|
|
| def slugify_model_id(model_id: str) -> str: |
| cleaned = re.sub(r"[^a-zA-Z0-9._-]+", "_", model_id.strip()) |
| return cleaned.strip("_") or "unknown_model" |
|
|
|
|
| def _is_llada_model_id(model_id: str) -> bool: |
| return "llada" in model_id.lower() |
|
|
|
|
| def _apply_llada_compat_patches(model_id: str, local_files_only: bool) -> None: |
| """Apply compatibility patches for LLaDA remote-code models.""" |
| try: |
| import transformers.modeling_utils as modeling_utils |
|
|
| if not hasattr(modeling_utils.PreTrainedModel, "all_tied_weights_keys"): |
| modeling_utils.PreTrainedModel.all_tied_weights_keys = {} |
| elif not isinstance(modeling_utils.PreTrainedModel.all_tied_weights_keys, dict): |
| modeling_utils.PreTrainedModel.all_tied_weights_keys = {} |
| except Exception: |
| |
| pass |
|
|
| try: |
| from transformers import AutoConfig |
| from transformers.dynamic_module_utils import get_class_from_dynamic_module |
|
|
| config = AutoConfig.from_pretrained( |
| model_id, |
| trust_remote_code=True, |
| local_files_only=local_files_only, |
| ) |
| auto_map = getattr(config, "auto_map", None) or {} |
| class_ref = auto_map.get("AutoModelForCausalLM") |
| if not class_ref: |
| return |
|
|
| model_cls = get_class_from_dynamic_module( |
| class_ref, |
| model_id, |
| local_files_only=local_files_only, |
| ) |
|
|
| original_tie = getattr(model_cls, "tie_weights", None) |
| if original_tie is None or getattr(original_tie, "_a1_llada_safe_wrapped", False): |
| return |
|
|
| def _safe_tie_weights(self: Any, *args: Any, **kwargs: Any) -> Any: |
| kwargs.pop("missing_keys", None) |
| kwargs.pop("recompute_mapping", None) |
| try: |
| return original_tie(self, *args, **kwargs) |
| except TypeError as exc: |
| if "unexpected keyword argument" in str(exc): |
| return original_tie(self) |
| raise |
|
|
| _safe_tie_weights._a1_llada_safe_wrapped = True |
| setattr(model_cls, "tie_weights", _safe_tie_weights) |
| except Exception: |
| |
| pass |
|
|
|
|
| def _normalize_all_tied_weights_keys(model: Any) -> None: |
| """Normalize missing/incompatible all_tied_weights_keys on loaded models.""" |
| try: |
| tied = getattr(model, "all_tied_weights_keys", None) |
| if tied is None: |
| model.all_tied_weights_keys = {} |
| elif callable(tied): |
| value = tied() |
| model.all_tied_weights_keys = value if isinstance(value, dict) else {} |
| elif not isinstance(tied, dict): |
| model.all_tied_weights_keys = {} |
| except Exception: |
| pass |
|
|
|
|
| def _load_causal_lm( |
| model_id: str, |
| model_dtype: Any, |
| local_files_only: bool, |
| ) -> Any: |
| from transformers import AutoModelForCausalLM |
|
|
| is_llada = _is_llada_model_id(model_id) |
| if is_llada: |
| _apply_llada_compat_patches(model_id=model_id, local_files_only=local_files_only) |
|
|
| load_kwargs: dict[str, Any] = { |
| "output_hidden_states": True, |
| "dtype": model_dtype, |
| "local_files_only": local_files_only, |
| } |
| if is_llada: |
| load_kwargs["trust_remote_code"] = True |
|
|
| try: |
| model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs) |
| except AttributeError as exc: |
| if not (is_llada and "all_tied_weights_keys" in str(exc)): |
| raise |
|
|
| _apply_llada_compat_patches(model_id=model_id, local_files_only=local_files_only) |
| try: |
| model = AutoModelForCausalLM.from_pretrained( |
| model_id, |
| low_cpu_mem_usage=False, |
| **load_kwargs, |
| ) |
| except TypeError: |
| model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs) |
|
|
| _normalize_all_tied_weights_keys(model) |
| if hasattr(model, "config") and not hasattr(model.config, "use_cache"): |
| try: |
| model.config.use_cache = False |
| except Exception: |
| pass |
|
|
| return model |
|
|
|
|
| @dataclass(frozen=True) |
| class FeatureExtractionRecord: |
| """Summary row for cached run-level feature extraction.""" |
|
|
| model_id: str |
| model_slug: str |
| run: int |
| n_words: int |
| n_layers: int |
| hidden_dim: int |
| unmatched_words: int |
| max_words_per_chunk: int |
| dry_run: bool |
| device: str |
| features_npz_path: str |
| metadata_json_path: str |
|
|
|
|
| def _parse_model_max_length(model: Any, tokenizer: Any) -> int: |
| candidate_values: list[int] = [] |
|
|
| max_position_embeddings = getattr(getattr(model, "config", None), "max_position_embeddings", None) |
| if isinstance(max_position_embeddings, int) and max_position_embeddings > 0: |
| candidate_values.append(int(max_position_embeddings)) |
|
|
| tokenizer_max = getattr(tokenizer, "model_max_length", None) |
| if isinstance(tokenizer_max, int) and 0 < tokenizer_max < 100000: |
| candidate_values.append(int(tokenizer_max)) |
|
|
| if candidate_values: |
| return int(min(candidate_values)) |
|
|
| return 4096 |
|
|
|
|
| def _build_text_and_word_spans(words: list[str]) -> tuple[str, list[tuple[int, int]]]: |
| safe_words = [str(word) for word in words] |
|
|
| spans: list[tuple[int, int]] = [] |
| cursor = 0 |
| chunks: list[str] = [] |
|
|
| for idx, word in enumerate(safe_words): |
| start = cursor |
| end = start + len(word) |
| spans.append((start, end)) |
| chunks.append(word) |
|
|
| cursor = end |
| if idx < len(safe_words) - 1: |
| chunks.append(" ") |
| cursor += 1 |
|
|
| return "".join(chunks), spans |
|
|
|
|
| def _map_tokens_to_words( |
| token_offsets: list[tuple[int, int]], |
| word_spans: list[tuple[int, int]], |
| ) -> tuple[list[list[int]], int]: |
| token_to_word: list[list[int]] = [[] for _ in range(len(word_spans))] |
|
|
| valid_token_centers: list[tuple[int, float]] = [] |
|
|
| for token_index, (token_start, token_end) in enumerate(token_offsets): |
| if token_end <= token_start: |
| continue |
|
|
| valid_token_centers.append((token_index, (token_start + token_end) * 0.5)) |
|
|
| for word_index, (word_start, word_end) in enumerate(word_spans): |
| overlaps = token_end > word_start and token_start < word_end |
| if overlaps: |
| token_to_word[word_index].append(token_index) |
| break |
|
|
| unmatched_words = 0 |
| if valid_token_centers: |
| centers = np.array([center for _, center in valid_token_centers], dtype=np.float64) |
| indices = [idx for idx, _ in valid_token_centers] |
|
|
| for word_index, word_tokens in enumerate(token_to_word): |
| if word_tokens: |
| continue |
|
|
| unmatched_words += 1 |
| word_start, word_end = word_spans[word_index] |
| word_center = (word_start + word_end) * 0.5 |
| nearest_idx = int(np.argmin(np.abs(centers - word_center))) |
| token_to_word[word_index] = [indices[nearest_idx]] |
| else: |
| unmatched_words = len(token_to_word) |
|
|
| return token_to_word, unmatched_words |
|
|
|
|
| def _extract_chunk_features( |
| words_chunk: list[str], |
| model: Any, |
| tokenizer: Any, |
| device: str, |
| selected_layers: list[int], |
| ) -> tuple[dict[int, np.ndarray], dict[str, Any]]: |
| import torch |
|
|
| if not getattr(tokenizer, "is_fast", False): |
| raise RuntimeError( |
| "Fast tokenizer with offset mapping is required for word-level aggregation." |
| ) |
|
|
| text, word_spans = _build_text_and_word_spans(words_chunk) |
| max_length = _parse_model_max_length(model=model, tokenizer=tokenizer) |
|
|
| encoded = tokenizer( |
| text, |
| return_tensors="pt", |
| return_offsets_mapping=True, |
| truncation=True, |
| max_length=max_length, |
| add_special_tokens=True, |
| return_overflowing_tokens=True, |
| ) |
|
|
| input_ids = encoded["input_ids"] |
| if input_ids.shape[0] != 1: |
| raise RuntimeError( |
| "Tokenizer overflow produced multiple windows. " |
| "Decrease --max-words-per-chunk." |
| ) |
|
|
| offset_mapping = encoded.pop("offset_mapping")[0].cpu().numpy().tolist() |
|
|
| model_inputs: dict[str, Any] = {} |
| for key, value in encoded.items(): |
| if key in {"overflow_to_sample_mapping", "num_truncated_tokens"}: |
| continue |
| model_inputs[key] = value.to(device) |
|
|
| with torch.no_grad(): |
| try: |
| outputs = model(**model_inputs, output_hidden_states=True, use_cache=False) |
| except TypeError as exc: |
| if "unexpected keyword argument" not in str(exc) or "use_cache" not in str(exc): |
| raise |
| outputs = model(**model_inputs, output_hidden_states=True) |
|
|
| hidden_states = outputs.hidden_states |
| if hidden_states is None: |
| raise RuntimeError("Model did not return hidden states") |
|
|
| token_to_word, unmatched_words = _map_tokens_to_words( |
| token_offsets=[(int(start), int(end)) for start, end in offset_mapping], |
| word_spans=word_spans, |
| ) |
|
|
| per_layer_features: dict[int, np.ndarray] = {} |
| hidden_dim = int(hidden_states[selected_layers[0]].shape[-1]) |
|
|
| for layer_idx in selected_layers: |
| layer_tokens = hidden_states[layer_idx][0].detach().float().cpu().numpy() |
| layer_word = np.zeros((len(words_chunk), hidden_dim), dtype=np.float32) |
|
|
| for word_index, token_indices in enumerate(token_to_word): |
| valid = [idx for idx in token_indices if 0 <= idx < layer_tokens.shape[0]] |
| if not valid: |
| continue |
| layer_word[word_index] = np.mean(layer_tokens[valid], axis=0, dtype=np.float32) |
|
|
| per_layer_features[layer_idx] = layer_word |
|
|
| diagnostics = { |
| "n_words": int(len(words_chunk)), |
| "n_tokens": int(len(offset_mapping)), |
| "unmatched_words": int(unmatched_words), |
| } |
| return per_layer_features, diagnostics |
|
|
|
|
| def _extract_real_features_for_run( |
| words: list[str], |
| model: Any, |
| tokenizer: Any, |
| device: str, |
| layer_indices: list[int] | None, |
| max_words_per_chunk: int, |
| ) -> tuple[dict[int, np.ndarray], dict[str, Any]]: |
| if max_words_per_chunk <= 0: |
| raise ValueError("max_words_per_chunk must be positive") |
|
|
| if not words: |
| raise ValueError("Cannot extract features from an empty word list") |
|
|
| n_all_layers = int(getattr(model.config, "num_hidden_layers", 0)) + 1 |
| selected_layers = layer_indices if layer_indices is not None else list(range(n_all_layers)) |
|
|
| for layer_idx in selected_layers: |
| if layer_idx < 0 or layer_idx >= n_all_layers: |
| raise ValueError(f"Layer index {layer_idx} out of range [0, {n_all_layers - 1}]") |
|
|
| chunk_outputs: dict[int, list[np.ndarray]] = {layer_idx: [] for layer_idx in selected_layers} |
| total_unmatched_words = 0 |
| total_tokens = 0 |
|
|
| start = 0 |
| while start < len(words): |
| stop = min(start + max_words_per_chunk, len(words)) |
| chunk_words = words[start:stop] |
|
|
| per_layer_chunk, chunk_diag = _extract_chunk_features( |
| words_chunk=chunk_words, |
| model=model, |
| tokenizer=tokenizer, |
| device=device, |
| selected_layers=selected_layers, |
| ) |
|
|
| total_unmatched_words += int(chunk_diag["unmatched_words"]) |
| total_tokens += int(chunk_diag["n_tokens"]) |
|
|
| for layer_idx in selected_layers: |
| chunk_outputs[layer_idx].append(per_layer_chunk[layer_idx]) |
|
|
| start = stop |
|
|
| outputs: dict[int, np.ndarray] = { |
| layer_idx: np.concatenate(chunks, axis=0).astype(np.float32) |
| for layer_idx, chunks in chunk_outputs.items() |
| } |
|
|
| hidden_dim = int(outputs[selected_layers[0]].shape[1]) |
| diagnostics = { |
| "n_words": int(len(words)), |
| "n_layers": int(len(selected_layers)), |
| "hidden_dim": hidden_dim, |
| "unmatched_words": int(total_unmatched_words), |
| "n_tokens_total": int(total_tokens), |
| "selected_layers": selected_layers, |
| } |
| return outputs, diagnostics |
|
|
|
|
| def _extract_dry_run_features_for_run( |
| words: list[str], |
| model_id: str, |
| run: int, |
| dry_run_n_layers: int, |
| dry_run_hidden_dim: int, |
| ) -> tuple[dict[int, np.ndarray], dict[str, Any]]: |
| if dry_run_n_layers <= 0: |
| raise ValueError("dry_run_n_layers must be positive") |
| if dry_run_hidden_dim <= 0: |
| raise ValueError("dry_run_hidden_dim must be positive") |
|
|
| n_words = len(words) |
| seed = abs(hash((model_id, int(run), n_words))) % (2**32) |
| rng = np.random.default_rng(seed) |
|
|
| outputs: dict[int, np.ndarray] = {} |
| for layer_idx in range(dry_run_n_layers): |
| features = rng.standard_normal(size=(n_words, dry_run_hidden_dim)).astype(np.float32) |
| outputs[layer_idx] = features |
|
|
| diagnostics = { |
| "n_words": int(n_words), |
| "n_layers": int(dry_run_n_layers), |
| "hidden_dim": int(dry_run_hidden_dim), |
| "unmatched_words": 0, |
| "n_tokens_total": int(n_words), |
| "selected_layers": list(range(dry_run_n_layers)), |
| } |
| return outputs, diagnostics |
|
|
|
|
| def extract_and_cache_run_level_features( |
| run_events_df: pd.DataFrame, |
| model_ids: list[str], |
| output_dir: Path, |
| layer_indices: list[int] | None, |
| max_words_per_chunk: int, |
| dry_run: bool, |
| dry_run_n_layers: int, |
| dry_run_hidden_dim: int, |
| device: str, |
| local_files_only: bool, |
| overwrite: bool, |
| num_workers: int = 1, |
| ) -> tuple[pd.DataFrame, dict[str, Any]]: |
| """Extract and cache run-level word features for each model. |
| |
| When ``num_workers > 1`` and multiple CUDA devices are visible, the work is |
| sharded across one process per GPU (each process pinned via |
| ``CUDA_VISIBLE_DEVICES``). Runs are partitioned round-robin across workers; |
| each worker still iterates the full ``model_ids`` list internally. |
| """ |
| if run_events_df.empty: |
| raise ValueError("run_events_df is empty; cannot extract features") |
|
|
| required_columns = {"run", "word_index", "word", "onset_s", "offset_s"} |
| missing = required_columns.difference(run_events_df.columns) |
| if missing: |
| raise ValueError(f"run_events_df missing required columns: {sorted(missing)}") |
|
|
| output_dir = output_dir.resolve() |
| output_dir.mkdir(parents=True, exist_ok=True) |
|
|
| runs = sorted({int(run) for run in run_events_df["run"].tolist()}) |
|
|
| if ( |
| num_workers > 1 |
| and not dry_run |
| and len(runs) > 1 |
| and _multi_gpu_available(device=device, requested_workers=num_workers) |
| ): |
| return _dispatch_multi_gpu_feature_extraction( |
| run_events_df=run_events_df, |
| model_ids=model_ids, |
| output_dir=output_dir, |
| layer_indices=layer_indices, |
| max_words_per_chunk=max_words_per_chunk, |
| local_files_only=local_files_only, |
| overwrite=overwrite, |
| num_workers=num_workers, |
| runs=runs, |
| ) |
|
|
| summary_rows: list[FeatureExtractionRecord] = [] |
|
|
| for model_id in model_ids: |
| model_slug = slugify_model_id(model_id) |
| model_output_dir = output_dir / model_slug |
| model_output_dir.mkdir(parents=True, exist_ok=True) |
|
|
| resolved_device = "dry-run" |
| model = None |
| tokenizer = None |
|
|
| if not dry_run: |
| import torch |
| from transformers import AutoTokenizer |
|
|
| is_llada = _is_llada_model_id(model_id) |
|
|
| if device == "auto": |
| resolved_device = "cuda" if torch.cuda.is_available() else "cpu" |
| else: |
| resolved_device = device |
|
|
| model_dtype = torch.float16 if resolved_device.startswith("cuda") else torch.float32 |
|
|
| tokenizer = AutoTokenizer.from_pretrained( |
| model_id, |
| use_fast=True, |
| local_files_only=local_files_only, |
| trust_remote_code=is_llada, |
| ) |
| if tokenizer.pad_token is None: |
| tokenizer.pad_token = tokenizer.eos_token |
|
|
| model = _load_causal_lm( |
| model_id, |
| model_dtype=model_dtype, |
| local_files_only=local_files_only, |
| ) |
|
|
| try: |
| model.to(resolved_device) |
| except torch.OutOfMemoryError as exc: |
| if device != "auto" or not resolved_device.startswith("cuda"): |
| raise RuntimeError( |
| "CUDA out of memory while moving model to device. " |
| "Retry with --feature-device cpu or reduce model size." |
| ) from exc |
|
|
| warnings.warn( |
| f"CUDA OOM while loading model {model_id}; falling back to CPU.", |
| RuntimeWarning, |
| ) |
|
|
| try: |
| del model |
| torch.cuda.empty_cache() |
| except Exception: |
| pass |
|
|
| resolved_device = "cpu" |
| model_dtype = torch.float32 |
| model = _load_causal_lm( |
| model_id, |
| model_dtype=model_dtype, |
| local_files_only=local_files_only, |
| ) |
| model.to(resolved_device) |
|
|
| model.eval() |
|
|
| for run in runs: |
| run_df = run_events_df[run_events_df["run"] == run].sort_values("word_index") |
| words = run_df["word"].astype(str).tolist() |
|
|
| npz_path = model_output_dir / f"run-{run:02d}_features.npz" |
| metadata_path = model_output_dir / f"run-{run:02d}_metadata.json" |
|
|
| if npz_path.exists() and metadata_path.exists() and not overwrite: |
| with metadata_path.open("r", encoding="utf-8") as handle: |
| metadata = json.load(handle) |
|
|
| summary_rows.append( |
| FeatureExtractionRecord( |
| model_id=model_id, |
| model_slug=model_slug, |
| run=int(run), |
| n_words=int(metadata["n_words"]), |
| n_layers=int(metadata["n_layers"]), |
| hidden_dim=int(metadata["hidden_dim"]), |
| unmatched_words=int(metadata.get("unmatched_words", 0)), |
| max_words_per_chunk=int(metadata.get("max_words_per_chunk", max_words_per_chunk)), |
| dry_run=bool(metadata.get("dry_run", dry_run)), |
| device=str(metadata.get("device", resolved_device)), |
| features_npz_path=str(npz_path), |
| metadata_json_path=str(metadata_path), |
| ) |
| ) |
| continue |
|
|
| if dry_run: |
| feature_map, diagnostics = _extract_dry_run_features_for_run( |
| words=words, |
| model_id=model_id, |
| run=int(run), |
| dry_run_n_layers=dry_run_n_layers, |
| dry_run_hidden_dim=dry_run_hidden_dim, |
| ) |
| else: |
| assert model is not None |
| assert tokenizer is not None |
| try: |
| feature_map, diagnostics = _extract_real_features_for_run( |
| words=words, |
| model=model, |
| tokenizer=tokenizer, |
| device=resolved_device, |
| layer_indices=layer_indices, |
| max_words_per_chunk=max_words_per_chunk, |
| ) |
| except torch.OutOfMemoryError as exc: |
| if device != "auto" or not resolved_device.startswith("cuda"): |
| raise RuntimeError( |
| "CUDA out of memory during feature extraction. " |
| "Retry with --feature-device cpu or reduce --max-words-per-chunk." |
| ) from exc |
|
|
| warnings.warn( |
| ( |
| f"CUDA OOM during feature extraction for model {model_id}, run={run}; " |
| "falling back to CPU and retrying." |
| ), |
| RuntimeWarning, |
| ) |
|
|
| torch.cuda.empty_cache() |
| resolved_device = "cpu" |
| model.to(resolved_device) |
|
|
| feature_map, diagnostics = _extract_real_features_for_run( |
| words=words, |
| model=model, |
| tokenizer=tokenizer, |
| device=resolved_device, |
| layer_indices=layer_indices, |
| max_words_per_chunk=max_words_per_chunk, |
| ) |
|
|
| np.savez( |
| npz_path, |
| **{f"layer_{layer_idx}": values for layer_idx, values in feature_map.items()}, |
| onset_s=run_df["onset_s"].to_numpy(dtype=np.float32), |
| offset_s=run_df["offset_s"].to_numpy(dtype=np.float32), |
| word_index=run_df["word_index"].to_numpy(dtype=np.int64), |
| ) |
|
|
| metadata = { |
| "model_id": model_id, |
| "model_slug": model_slug, |
| "run": int(run), |
| "n_words": int(diagnostics["n_words"]), |
| "n_layers": int(diagnostics["n_layers"]), |
| "hidden_dim": int(diagnostics["hidden_dim"]), |
| "unmatched_words": int(diagnostics["unmatched_words"]), |
| "n_tokens_total": int(diagnostics["n_tokens_total"]), |
| "selected_layers": [int(value) for value in diagnostics["selected_layers"]], |
| "max_words_per_chunk": int(max_words_per_chunk), |
| "dry_run": bool(dry_run), |
| "device": str(resolved_device), |
| "features_npz_path": str(npz_path), |
| } |
| with metadata_path.open("w", encoding="utf-8") as handle: |
| json.dump(metadata, handle, indent=2, sort_keys=True) |
|
|
| summary_rows.append( |
| FeatureExtractionRecord( |
| model_id=model_id, |
| model_slug=model_slug, |
| run=int(run), |
| n_words=int(diagnostics["n_words"]), |
| n_layers=int(diagnostics["n_layers"]), |
| hidden_dim=int(diagnostics["hidden_dim"]), |
| unmatched_words=int(diagnostics["unmatched_words"]), |
| max_words_per_chunk=int(max_words_per_chunk), |
| dry_run=bool(dry_run), |
| device=str(resolved_device), |
| features_npz_path=str(npz_path), |
| metadata_json_path=str(metadata_path), |
| ) |
| ) |
|
|
| if model is not None: |
| del model |
| del tokenizer |
|
|
| summary_df = pd.DataFrame([asdict(row) for row in summary_rows]) |
| if not summary_df.empty: |
| summary_df = summary_df.sort_values(["model_slug", "run"]).reset_index(drop=True) |
|
|
| feature_qc: dict[str, Any] = { |
| "n_models": int(len({row.model_slug for row in summary_rows})), |
| "n_model_run_rows": int(len(summary_rows)), |
| "dry_run": bool(dry_run), |
| "max_words_per_chunk": int(max_words_per_chunk), |
| } |
|
|
| if not summary_df.empty: |
| feature_qc["n_layers_min"] = int(summary_df["n_layers"].min()) |
| feature_qc["n_layers_max"] = int(summary_df["n_layers"].max()) |
| feature_qc["hidden_dim_min"] = int(summary_df["hidden_dim"].min()) |
| feature_qc["hidden_dim_max"] = int(summary_df["hidden_dim"].max()) |
| feature_qc["unmatched_words_total"] = int(summary_df["unmatched_words"].sum()) |
|
|
| return summary_df, feature_qc |
|
|
|
|
| def _multi_gpu_available(device: str, requested_workers: int) -> bool: |
| """Return True when CUDA exposes >=2 devices and the request is sane.""" |
| if requested_workers <= 1: |
| return False |
|
|
| device_lower = str(device).strip().lower() |
| if device_lower in {"cpu", "dry-run"}: |
| return False |
|
|
| try: |
| import torch |
| except Exception: |
| return False |
|
|
| if not torch.cuda.is_available(): |
| return False |
|
|
| return torch.cuda.device_count() >= 2 |
|
|
|
|
| def resolve_feature_num_workers(requested: int | str, device: str) -> int: |
| """Resolve --feature-num-workers ('auto' or int) to a concrete worker count. |
| |
| Returns 1 unless multiple CUDA devices are visible and ``device`` is auto/cuda. |
| """ |
| device_lower = str(device).strip().lower() |
| if device_lower in {"cpu", "dry-run"}: |
| return 1 |
|
|
| try: |
| import torch |
| n_gpus = int(torch.cuda.device_count()) if torch.cuda.is_available() else 0 |
| except Exception: |
| n_gpus = 0 |
|
|
| if isinstance(requested, str): |
| token = requested.strip().lower() |
| if token in {"", "auto"}: |
| return max(1, n_gpus) |
| try: |
| value = int(token) |
| except ValueError as exc: |
| raise ValueError(f"Invalid --feature-num-workers={requested!r}") from exc |
| else: |
| value = int(requested) |
|
|
| if value <= 1: |
| return 1 |
| if n_gpus <= 0: |
| return 1 |
| return min(value, n_gpus) |
|
|
|
|
| def _record_from_metadata( |
| *, |
| model_id: str, |
| model_slug: str, |
| run: int, |
| metadata: dict[str, Any], |
| npz_path: Path, |
| metadata_path: Path, |
| max_words_per_chunk: int, |
| ) -> FeatureExtractionRecord: |
| return FeatureExtractionRecord( |
| model_id=model_id, |
| model_slug=model_slug, |
| run=int(run), |
| n_words=int(metadata["n_words"]), |
| n_layers=int(metadata["n_layers"]), |
| hidden_dim=int(metadata["hidden_dim"]), |
| unmatched_words=int(metadata.get("unmatched_words", 0)), |
| max_words_per_chunk=int(metadata.get("max_words_per_chunk", max_words_per_chunk)), |
| dry_run=bool(metadata.get("dry_run", False)), |
| device=str(metadata.get("device", "cuda")), |
| features_npz_path=str(npz_path), |
| metadata_json_path=str(metadata_path), |
| ) |
|
|
|
|
| def _gpu_worker_entrypoint( |
| rank: int, |
| world_size: int, |
| payload_path: str, |
| ) -> None: |
| """Process entrypoint for one GPU worker. Runs in a spawned subprocess.""" |
| import pickle |
|
|
| |
| os.environ["CUDA_VISIBLE_DEVICES"] = str(rank) |
|
|
| |
| n_cpu = os.cpu_count() or 8 |
| threads = max(1, n_cpu // max(1, world_size)) |
| os.environ.setdefault("OMP_NUM_THREADS", str(threads)) |
| os.environ.setdefault("MKL_NUM_THREADS", str(threads)) |
| os.environ.setdefault("OPENBLAS_NUM_THREADS", str(threads)) |
| os.environ.setdefault("NUMEXPR_NUM_THREADS", str(threads)) |
|
|
| try: |
| import torch |
| torch.set_num_threads(threads) |
| except Exception: |
| pass |
|
|
| with open(payload_path, "rb") as handle: |
| payload: dict[str, Any] = pickle.load(handle) |
|
|
| all_runs: list[int] = payload["runs"] |
| my_runs = [r for idx, r in enumerate(all_runs) if idx % world_size == rank] |
| if not my_runs: |
| return |
|
|
| run_events_df: pd.DataFrame = payload["run_events_df"] |
| df_subset = run_events_df[run_events_df["run"].isin(my_runs)].reset_index(drop=True) |
| if df_subset.empty: |
| return |
|
|
| extract_and_cache_run_level_features( |
| run_events_df=df_subset, |
| model_ids=payload["model_ids"], |
| output_dir=Path(payload["output_dir"]), |
| layer_indices=payload["layer_indices"], |
| max_words_per_chunk=payload["max_words_per_chunk"], |
| dry_run=False, |
| dry_run_n_layers=0, |
| dry_run_hidden_dim=0, |
| device="cuda:0", |
| local_files_only=payload["local_files_only"], |
| overwrite=payload["overwrite"], |
| num_workers=1, |
| ) |
|
|
|
|
| def _dispatch_multi_gpu_feature_extraction( |
| *, |
| run_events_df: pd.DataFrame, |
| model_ids: list[str], |
| output_dir: Path, |
| layer_indices: list[int] | None, |
| max_words_per_chunk: int, |
| local_files_only: bool, |
| overwrite: bool, |
| num_workers: int, |
| runs: list[int], |
| ) -> tuple[pd.DataFrame, dict[str, Any]]: |
| """Spawn one worker per GPU; each worker handles a disjoint subset of runs.""" |
| import pickle |
| import tempfile |
| import torch.multiprocessing as mp |
|
|
| world_size = min(int(num_workers), len(runs)) |
|
|
| print( |
| f"[features] Multi-GPU feature extraction: world_size={world_size}, " |
| f"runs={runs}, models={len(model_ids)}", |
| flush=True, |
| ) |
|
|
| payload = { |
| "runs": runs, |
| "run_events_df": run_events_df, |
| "model_ids": list(model_ids), |
| "output_dir": str(output_dir), |
| "layer_indices": layer_indices, |
| "max_words_per_chunk": int(max_words_per_chunk), |
| "local_files_only": bool(local_files_only), |
| "overwrite": bool(overwrite), |
| } |
|
|
| with tempfile.NamedTemporaryFile( |
| mode="wb", suffix=".pkl", delete=False, dir=str(output_dir) |
| ) as handle: |
| pickle.dump(payload, handle) |
| payload_path = handle.name |
|
|
| try: |
| mp.spawn( |
| _gpu_worker_entrypoint, |
| args=(world_size, payload_path), |
| nprocs=world_size, |
| join=True, |
| ) |
| finally: |
| try: |
| os.unlink(payload_path) |
| except OSError: |
| pass |
|
|
| |
| summary_rows: list[FeatureExtractionRecord] = [] |
| for model_id in model_ids: |
| model_slug = slugify_model_id(model_id) |
| model_output_dir = output_dir / model_slug |
| for run in runs: |
| npz_path = model_output_dir / f"run-{run:02d}_features.npz" |
| metadata_path = model_output_dir / f"run-{run:02d}_metadata.json" |
| if not (npz_path.exists() and metadata_path.exists()): |
| raise RuntimeError( |
| f"Multi-GPU worker did not produce features for " |
| f"model={model_id} run={run}: missing {metadata_path} or {npz_path}" |
| ) |
| with metadata_path.open("r", encoding="utf-8") as handle: |
| metadata = json.load(handle) |
| summary_rows.append( |
| _record_from_metadata( |
| model_id=model_id, |
| model_slug=model_slug, |
| run=int(run), |
| metadata=metadata, |
| npz_path=npz_path, |
| metadata_path=metadata_path, |
| max_words_per_chunk=int(max_words_per_chunk), |
| ) |
| ) |
|
|
| summary_df = pd.DataFrame([asdict(row) for row in summary_rows]) |
| if not summary_df.empty: |
| summary_df = summary_df.sort_values(["model_slug", "run"]).reset_index(drop=True) |
|
|
| feature_qc: dict[str, Any] = { |
| "n_models": int(len({row.model_slug for row in summary_rows})), |
| "n_model_run_rows": int(len(summary_rows)), |
| "dry_run": False, |
| "max_words_per_chunk": int(max_words_per_chunk), |
| "multi_gpu_world_size": int(world_size), |
| } |
| if not summary_df.empty: |
| feature_qc["n_layers_min"] = int(summary_df["n_layers"].min()) |
| feature_qc["n_layers_max"] = int(summary_df["n_layers"].max()) |
| feature_qc["hidden_dim_min"] = int(summary_df["hidden_dim"].min()) |
| feature_qc["hidden_dim_max"] = int(summary_df["hidden_dim"].max()) |
| feature_qc["unmatched_words_total"] = int(summary_df["unmatched_words"].sum()) |
|
|
| return summary_df, feature_qc |
|
|
