bench_inference.py

import vllm
import os
from src.training.build_dataset import load_tl_model
PILE_PATH = "/root/pile_uncopyrighted"


from transformer_lens import HookedTransformer
from transformer_lens.hook_points import HookPoint
from typing import Optional
import pandas as pd
import numpy as np
import time
from tqdm.auto import tqdm
from functools import partial
import torch

def TL_inference(
    model, 
    tokens : torch.Tensor, 
    hook_datasets : dict[str, Optional[list[np.ndarray]]],
    stop_at_layer_idx : int
):    
    # Initialize datasets for each hook
    

    def hook_fn(
        tensor: torch.Tensor, 
        hook_name: str, 
        hook: HookPoint = None
    ):
        
        t0 = time.time()
        processed_tensors = tensor.detach().cpu().numpy()
        print("took t0", t0-time.time())

        
        if hook_datasets[hook_name] is None:
            hook_datasets[hook_name] = [processed_tensors]
        else:
            hook_datasets[hook_name].append(processed_tensors)
            
        return tensor


   
    with torch.no_grad():
        # Convert to numpy array first, then to tensor
        t0 = time.time()
        model.run_with_hooks(
            tokens,
            stop_at_layer = stop_at_layer_idx,
            fwd_hooks=[
                (hook_name, partial(hook_fn, hook_name=hook_name)) 
                for hook_name in hook_names
            ]
        )
        torch.cuda.synchronize()
        t1 = time.time()
        print(f"forward pass with cache took {t1-t0} seconds")
        tok_per_sec = tokens.numel() / (t1-t0)
        print(f"tokens per second: {tok_per_sec}")


def compute_gb(hook_acts: np.ndarray):
    assert isinstance(hook_acts, np.ndarray)
    return hook_acts.nbytes / 1e9  # Use nbytes and divide by 1e9 for GB

tl_bench = True
print("gpu memory", torch.cuda.memory_summary())

if tl_bench:
    device = "cuda"
    model : HookedTransformer = load_tl_model("google/gemma-2-2b")
    print("gpu memory", torch.cuda.memory_summary())
    
    hook_names = [f"blocks.{i}.hook_resid_post" for i in [3]]
    hook_datasets : dict[str, Optional[list[np.ndarray]]] = {
        hook_name: None for hook_name in hook_names
    }
    from tqdm import tqdm
    for _ in tqdm(range(100)):
        tokens = torch.randint(1, 100, (1, 512))
        TL_inference(model, tokens, hook_datasets, stop_at_layer_idx=4)