[draft, low prio] infra: Fix flash-attn

[tomtseng]

Changes

Our implementation of GCG uses flash attention. PR #53 added the flash attention package, but actually it's the wrong one — we want flash-attn instead of flash_attention. However flash-attn is tricky to install, it needs to be compiled. So the Docker image needs to be updated.

It's not clear we care about GCG and we can prob turn off flash attention in GCG. so this PR is low prio, I think I'll put it aside for now.

TODO: trying to build dockerfile but currently it's OOMing when compiling locally, even when I max out the Docker Desktop memory cap setting and I set MAX_JOBS=1. Maybe compile on flamingo?

Testing

TODO build image, launch a new devbox and get this script to work:

"""Test script to verify flash-attn works with transformers."""

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"CUDA version: {torch.version.cuda}")

print(f"\nLoading {MODEL_ID} with Flash Attention 2...")

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2",
    device_map="auto",
)

print(f"Model loaded successfully!")
print(f"Attention implementation: {model.config._attn_implementation}")

# Quick inference test
input_text = "Hello, world!"
inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=10)
result = tokenizer.decode(outputs[0], skip_special_tokens=True)

print(f"\nInference test:")
print(f"  Input: {input_text}")
print(f"  Output: {result}")
print("\nFlash Attention 2 is working!")

[tomtseng]

needed to bump CUDA version or else flash-attn would have a compile error. decided to bump PyTorch version to the latest compatible version too

[tomtseng]

GCG refactor in #82 removed the reference to flash-attn, this is no longer relevant. Though I may take some of the Dockerfile refactors and open a separate PR for that