Add Python 3.11 to CI tests (#91)

* load dset

* use Datasets.from_generator for multiprocessing

* revert some debugging stuff

* remove accidentally included files

* switch to correct multiprocessing lib + switch to spawn

* split dataset before multiprocessing

* nicer prints

* Replace custom method with np.array_split

* Move multiprocessing inside generator

* change yield to put

* change to use model_str

* fix connectionreseterror

* switch extract to return dataset instead of saving to disk

* Switch to torch.multiprocessing.spawn

* Start testing on Python 3.11

---------

Co-authored-by: Alexander Wan <[email protected]>

.github/workflows/cpu_ci.yml

@@ -46,3 +46,25 @@ jobs:
 
       - name: Run CPU Tests
         run: pytest -m cpu
+
+  run-tests-python3_11:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+
+      - name: Install Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.11"
+
+      - name: Upgrade Pip
+        run: python -m pip install --upgrade pip
+
+      - name: Install Dependencies
+        run: pip install -e .[dev]
+
+      - name: Type Checking
+        uses: jakebailey/pyright-action@v1
+
+      - name: Run CPU Tests
+        run: pytest -m cpu

elk/__main__.py

@@ -4,8 +4,6 @@
 from .files import args_to_uuid
 from .list import list_runs
 from argparse import ArgumentParser
-from contextlib import nullcontext, redirect_stdout
-import logging
 import warnings
 
 
@@ -40,6 +38,7 @@ def run():
         list_runs(args)
         return
 
+    # Import here and not at the top to speed up `elk list`
     from transformers import AutoConfig, PretrainedConfig
 
     config = AutoConfig.from_pretrained(args.model)
@@ -67,47 +66,30 @@ def run():
     # Import here and not at the top to speed up `elk list`
     from .extraction.extraction_main import run as run_extraction
     from .training.train import train
-    import os
-    import torch.distributed as dist
-
-    # Check if we were called with torchrun or not
-    local_rank = os.environ.get("LOCAL_RANK")
-    if local_rank is not None:
-        dist.init_process_group("nccl")
-        local_rank = int(local_rank)
-
-    with redirect_stdout(None) if local_rank else nullcontext():
-        # Print CLI arguments to stdout
-        for key, value in vars(args).items():
-            print(f"{key}: {value}")
 
-        if local_rank:
-            logging.getLogger("transformers").setLevel(logging.CRITICAL)
-
-        if args.command == "extract":
+    # Print CLI arguments to stdout
+    for key, value in vars(args).items():
+        print(f"{key}: {value}")
+
+    if args.command == "extract":
+        run_extraction(args)
+    elif args.command == "elicit":
+        # The user can specify a name for the run, but by default we use the
+        # MD5 hash of the arguments to ensure the name is unique
+        if not args.name:
+            args.name = args_to_uuid(args)
+
+        try:
+            train(args)
+        except (EOFError, FileNotFoundError):
             run_extraction(args)
-        elif args.command == "elicit":
-            # The user can specify a name for the run, but by default we use the
-            # MD5 hash of the arguments to ensure the name is unique
-            if not args.name:
-                args.name = args_to_uuid(args)
-
-            try:
-                train(args)
-            except (EOFError, FileNotFoundError):
-                run_extraction(args)
-
-                # Ensure the extraction is finished before starting training
-                if dist.is_initialized():
-                    dist.barrier()
-
-                train(args)
-
-        elif args.command == "eval":
-            # TODO: Implement evaluation script
-            raise NotImplementedError
-        else:
-            raise ValueError(f"Unknown command {args.command}")
+            train(args)
+
+    elif args.command == "eval":
+        # TODO: Implement evaluation script
+        raise NotImplementedError
+    else:
+        raise ValueError(f"Unknown command {args.command}")
 
 
 if __name__ == "__main__":

elk/extraction/extraction.py

@@ -2,19 +2,36 @@
 
 from ..utils import pytree_map
 from .prompt_collator import Prompt, PromptCollator
+from dataclasses import dataclass
 from einops import rearrange
 from torch.utils.data import DataLoader
-from tqdm.auto import tqdm
-from transformers import BatchEncoding, PreTrainedModel, PreTrainedTokenizerBase
-from typing import cast, Iterable, Literal, Sequence
+from transformers import (
+    BatchEncoding,
+    PreTrainedModel,
+    PreTrainedTokenizerBase,
+    AutoModel,
+)
+from typing import cast, Literal, Sequence
+import logging
+import numpy as np
 import torch
-import torch.distributed as dist
+import torch.multiprocessing as mp
+
+
+@dataclass
+class ExtractionParameters:
+    model_str: str
+    tokenizer: PreTrainedTokenizerBase
+    collator: PromptCollator
+    batch_size: int = 1
+    layers: Sequence[int] = ()
+    prompt_suffix: str = ""
+    token_loc: Literal["first", "last", "mean"] = "last"
+    use_encoder_states: bool = False
 
 
-@torch.autocast("cuda", enabled=torch.cuda.is_available())  # type: ignore
-@torch.no_grad()
 def extract_hiddens(
-    model: PreTrainedModel,
+    model_str: str,
     tokenizer: PreTrainedTokenizerBase,
     collator: PromptCollator,
     *,
@@ -24,40 +41,85 @@ def extract_hiddens(
     prompt_suffix: str = "",
     token_loc: Literal["first", "last", "mean"] = "last",
     use_encoder_states: bool = False,
-) -> Iterable[tuple[torch.Tensor, list[int]]]:
-    """Run inference on a model with a set of prompts, yielding the hidden states.
-
-    Args:
-        model: The model to run inference on.
-        tokenizer: The tokenizer to use for tokenization.
-        collator: The PromptCollator to use for generating prompts.
-        batch_size: The batch size to use for inference.
-        layers (Sequence[int]): The layers to extract hidden states from.
-        prompt_suffix (str): A string to append to the end of each prompt.
-        token_loc: The location of the token to extract hidden states from.
-            can be either "first", "last", or "mean". Defaults to "last".
-        use_encoder_states: Whether to use the encoder states instead of the
-            decoder states. This allows simplification from an encoder-decoder
-            model to an encoder-only model. Defaults to False.
+):
+    """Run inference on a model with a set of prompts, yielding the hidden states."""
+
+    ctx = mp.get_context("spawn")
+    queue = ctx.Queue()
+
+    num_gpus = torch.cuda.device_count()
+    params = ExtractionParameters(
+        model_str=model_str,
+        tokenizer=tokenizer,
+        collator=collator,
+        batch_size=batch_size,
+        layers=layers,
+        prompt_suffix=prompt_suffix,
+        token_loc=token_loc,
+        use_encoder_states=use_encoder_states,
+    )
+
+    # Spawn a process for each GPU
+    ctx = torch.multiprocessing.spawn(
+        _extract_hiddens_process,
+        args=(num_gpus, queue, params),
+        nprocs=num_gpus,
+        join=False,
+    )
+    assert ctx is not None
+
+    # Yield results from the queue
+    for _ in range(len(collator)):
+        yield queue.get()
+
+    # Clean up
+    ctx.join()
+
+
+@torch.no_grad()
+def _extract_hiddens_process(
+    rank: int,
+    world_size: int,
+    queue: mp.Queue,
+    params: ExtractionParameters,
+):
     """
+    Do inference on a model with a set of prompts on a single process.
+    To be passed to Dataset.from_generator.
+    """
+    print(f"Process with rank={rank}")
+    if rank != 0:
+        logging.getLogger("transformers").setLevel(logging.CRITICAL)
+
+    num_choices = params.collator.num_classes
+    shards = np.array_split(np.arange(len(params.collator)), world_size)
+    params.collator.select_(shards[rank])
+
+    # AutoModel should do the right thing here in nearly all cases. We don't actually
+    # care what head the model has, since we are just extracting hidden states.
+    model = AutoModel.from_pretrained(params.model_str, torch_dtype="auto").to(
+        f"cuda:{rank}"
+    )
 
-    device = model.device
-    num_choices = collator.num_classes
+    if params.use_encoder_states and not model.config.is_encoder_decoder:
+        raise ValueError(
+            "use_encoder_states is only compatible with encoder-decoder models."
+        )
 
     # TODO: Make this configurable or something
     # Token used to separate the question from the answer
-    sep_token = tokenizer.sep_token or "\n"
+    sep_token = params.tokenizer.sep_token or "\n"
 
     # TODO: Maybe also make this configurable?
     # We want to make sure the answer is never truncated
-    tokenizer.truncation_side = "left"
-    if not tokenizer.pad_token:
-        tokenizer.pad_token = tokenizer.eos_token
+    params.tokenizer.truncation_side = "left"
+    if not params.tokenizer.pad_token:
+        params.tokenizer.pad_token = params.tokenizer.eos_token
 
     def tokenize(strings: list[str]):
         return pytree_map(
-            lambda x: x.to(device),
-            tokenizer(
+            lambda x: x.to(f"cuda:{rank}"),
+            params.tokenizer(
                 strings,
                 padding=True,
                 return_tensors="pt",
@@ -69,7 +131,7 @@ def tokenize(strings: list[str]):
     # each question-answer pair. After inference we need to reshape the results.
     def collate(prompts: list[Prompt]) -> tuple[BatchEncoding, list[int]]:
         choices = [
-            prompt.to_string(i, sep=sep_token) + prompt_suffix
+            prompt.to_string(i, sep=sep_token) + params.prompt_suffix
             for prompt in prompts
             for i in range(num_choices)
         ]
@@ -83,7 +145,7 @@ def collate_enc_dec(
         )
         tokenized_answers = tokenize(
             [
-                prompt.answers[i] + prompt_suffix
+                prompt.answers[i] + params.prompt_suffix
                 for prompt in prompts
                 for i in range(num_choices)
             ]
@@ -99,22 +161,22 @@ def reduce_seqs(
         # Unflatten the hiddens
         hiddens = [rearrange(h, "(b c) l d -> b c l d", c=num_choices) for h in hiddens]
 
-        if token_loc == "first":
+        if params.token_loc == "first":
             hiddens = [h[..., 0, :] for h in hiddens]
-        elif token_loc == "last":
+        elif params.token_loc == "last":
             # Because of padding, the last token is going to be at a different index
             # for each example, so we use gather.
             B, C, _, D = hiddens[0].shape
             lengths = attention_mask.sum(dim=-1).view(B, C, 1, 1)
             indices = lengths.sub(1).expand(B, C, 1, D)
             hiddens = [h.gather(index=indices, dim=-2).squeeze(-2) for h in hiddens]
-        elif token_loc == "mean":
+        elif params.token_loc == "mean":
             hiddens = [h.mean(dim=-2) for h in hiddens]
         else:
-            raise ValueError(f"Invalid token_loc: {token_loc}")
+            raise ValueError(f"Invalid token_loc: {params.token_loc}")
 
-        if layers:
-            hiddens = [hiddens[i] for i in layers]
+        if params.layers:
+            hiddens = [hiddens[i] for i in params.layers]
 
         # [batch size, layers, num choices, hidden size]
         return torch.stack(hiddens, dim=1)
@@ -123,7 +185,7 @@ def reduce_seqs(
     # we don't need to run the decoder at all. Just strip it off, making the problem
     # equivalent to a regular encoder-only model.
     is_enc_dec = model.config.is_encoder_decoder
-    if is_enc_dec and use_encoder_states:
+    if is_enc_dec and params.use_encoder_states:
         # This isn't actually *guaranteed* by HF, but it's true for all existing models
         if not hasattr(model, "get_encoder") or not callable(model.get_encoder):
             raise ValueError(
@@ -134,17 +196,16 @@ def reduce_seqs(
 
     # Whether to concatenate the question and answer before passing to the model.
     # If False pass them to the encoder and decoder separately.
-    should_concat = not is_enc_dec or use_encoder_states
+    should_concat = not is_enc_dec or params.use_encoder_states
 
     dl = DataLoader(
-        collator,
-        batch_size=batch_size,
+        params.collator,
+        batch_size=params.batch_size,
         collate_fn=collate if should_concat else collate_enc_dec,
     )
 
     # Iterating over questions
-    rank = dist.get_rank() if dist.is_initialized() else 0
-    for batch in tqdm(dl, position=rank):
+    for batch in dl:
         # Condition 1: Encoder-decoder transformer, with answer in the decoder
         if not should_concat:
             questions, answers, labels = batch
@@ -154,12 +215,29 @@ def reduce_seqs(
                 output_hidden_states=True,
             )
             # [batch_size, num_layers, num_choices, hidden_size]
-            yield torch.stack(outputs.decoder_hidden_states, dim=2), labels
+            # need to convert hidden states to numpy array first or
+            # you get a ConnectionResetErrror
+            queue.put(
+                {
+                    "hiddens": torch.stack(outputs.decoder_hidden_states, dim=2)
+                    .cpu()
+                    .numpy(),
+                    "labels": labels,
+                }
+            )
 
         # Condition 2: Either a decoder-only transformer or a transformer encoder
         else:
             choices, labels = batch
 
             # Skip the input embeddings which are unlikely to be interesting
             h = model(**choices, output_hidden_states=True).hidden_states[1:]
-            yield reduce_seqs(h, choices["attention_mask"]), labels
+
+            # need to convert hidden states to numpy array first or
+            # you get a ConnectionResetErrror
+            queue.put(
+                {
+                    "hiddens": reduce_seqs(h, choices["attention_mask"]).cpu().numpy(),
+                    "labels": labels,
+                }
+            )