code/a1_pipeline/features.py

"""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:
        # Best effort; continue with standard load flow.
        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:
        # Best effort; continue with standard load flow.
        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

    # Pin this process to a single GPU before importing torch in the child.
    os.environ["CUDA_VISIBLE_DEVICES"] = str(rank)

    # Avoid BLAS thread oversubscription across workers.
    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

    # Aggregate summary by reading metadata files written by workers.
    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