Tips and Code for Empirical Research Workflows

Ethan Perez

Some additional opinionated tips:

Cache LLM responses! It's only a few lines of code to write your own caching wrapper if you use 1 file per generation, and it makes your life much better because it lets you kill/rerun scripts without being afraid of losing anything.
Always have at least one experiment script that is slow and runs in the background, and an explore / plot script, that you use for dataset exploration and plotting.
- If you have n slow steps (e.g. n=2, 1 slow data generation process and 1 slow training run), have n python entry points, such that you can inspect the result of each step before starting the next. Once the pipeline works, you can always write a script that runs them one after another.
When fine-tuning a model or creating a dataset, save and look at some your transcripts to understand what the model is learning / should learn.
- If wandb makes it mildly annoying, don't use wandb, just store json/jsonl locally
In general, using plotting / printing results as the main debugging tool: the evolution of train/eval during training loss should look fine, the transcript should look reasonable, the distribution of probabilities should look reasonable, etc. If anything looks surprising, investigate by plotting / printing more stuff and looking at the places in the code/data that could have caused the issue.
- Also make guesses about API cost / training speed, and check that your expectations are roughly correct near the start of the run (e.g. use tqdm on any slow for loop / gather). This lets you notice issues like putting things on the wrong device, not hitting the API cache, using a dataset that is too big for a de-risking experiment, etc.
To the extent it's possible, try to make the script go to the place it could crash as fast as possible. For example, if a script crashes on the first backward pass after 2 minutes of tokenization, avoid always spending 2 minutes waiting for tokenization during debugging (e.g. tokenize on the fly or cache the tokenized dataset)
Do simple fast experiments on simple infra before doing big expensive experiments on complex infra
- For fine-tuning / probing projects, I often get a lot of value out of doing a single-GPU run on a 1B model using a vanilla pytorch loop. That can often get signal in a few minutes and forces me to think through what different results could mean / whether there are boring reasons for surprising results.
- If things go wrong, simplify even further: e.g. if I train a password-locked model and the with-password-perf is low, I would try training a model just on the strong trajectories
Use #%% python notebooks, they play nicely with copilot, code editors and git. And you need notebooks to do printing / plotting.
Use prompting experiments before doing SL. Do SL before doing RL. And you often don't need SL or RL (e.g. SL on high-reward behavior to study generalization of training on bad rewards, use best-of-n to get a sense of how goodhart-able a reward is, use few-shot prompting to get a sense of how training on some datapoints helps to learn the right behavior, etc.)
- Prompting experiments can happen in your notebook, they don't always need to be clean and large-n.
Use as few libraries you don't understand as possible. I like just using torch+huggingface+tqdm+matplotlib. For model-internals experiments, I often use either huggingface's output_hidden_states or a pytorch hook. For saving data, I often just use json.

Mistakes to avoid:

Not sweeping learning rates when fine-tuning a model (I would usually sweep around the expected lr. For example Adam often expects something like lr=1e-4, so I would at least try something like 1e-5, 3e-5, 1e-4 and 3e-4, and then adjust during the next experiment based on train loss and val accuracy, using smaller gaps if I want better performance and not just a reasonable lr)
Not shuffling your dataset (it's free, don't rely on someone else shuffling your dataset)
- This is important not just for training datasets, but also when taking n elements from an evaluation dataset - datasets like Alpaca are not shuffled, so taking the first n elements will no be the same as taking n random points from the dataset.
When distilling one model into another, not sampling at temperature 1 (sampling at temperature 1 and training on it is what should copy the original model probs)
When evaluating a model, only sampling at temperature 1 (sampling at a lower temperature often works better)
Using an instruction-tuned model without using the official prompt template (usually chat_apply_template)
Not including the end-of-sequence token when fine-tuning on an instruction-tuning dataset
Not sweeping in logspace when sweeping hyperparams like temperature, lr, etc.
Not training for long enough / on enough datapoints - fine-tune on at least 1k trajectories before giving up. Small <1k trajectories runs are not that informative.
Using a huggingface model forward pass with padding side left without prepare_inputs_for_generation
Using await in a for loop (in particular when awaiting LLMs calls, always use gather)
If your LLMs request benefits from prefix caching, not using a semaphore to make sure the queries with shared prefix happen at <5 minutes of interval
Using python's built-in hash instead of md5 when you want a hash that is consistent over runs
Using random.Random(some object that is not a str or an int) (it sometimes uses the object's id as seed, which varies between runs)
Not using vllm and other specialized inference libraries instead of huggingface model.generate (which can be 30x slower)
- Using model.generate is fine for small mid-run evals. If you do this, at least batch your generations.

This looks extremely comprehensive and useful, thanks a lot for writing it! Some of my favourite tips (like clipboard managers and rectangle) were included, which is always a good sign. And I strongly agree with "Cursor/LLM-assisted coding is basically mandatory".

I passed this on to my mentees - not all of this transfers to mech interp, in particular the time between experiments is often much shorter (eg a few minutes, or even seconds) and often almost an entire project is in de-risking mode, but much of it transfers. And the ability to get shit done fast is super important

[-]John Hughes11moΩ220

Thanks Neel! I'm glad you found it helpful. If you or your scholars recommend any other tools not mentioned in the post, I'd be interested to hear more.

[-]Neel Nanda11mo40

I've been really enjoying voice to text + LLMs recently, via a great Mac App called Super Whisper (which can work with local speech to text models, so could also possibly be used for confidential stuff) - combining Super Whisper and Claude and Cursor means I can just vaguely ramble at my laptop about what experiments should happen and they happen, it's magical!

[-]Sheikh Abdur Raheem Ali10mo10

What keybinding do you set for it?

IIRC, ⌥+space conflicts with the default for ChatGPT and Alfred.app (which I use for clipboard history).

[-]Neel Nanda10mo30

Control space

[-]Sheikh Abdur Raheem Ali10mo10

For anyone else who stumbles across this thread: when modifying the superwhisper toggle settings, hit spacebar then control, instead of control then spacebar. Also, it turns out that Control + Space is the default shortcut for switching keyboard input sources (at least on macOS Sequoia 15.3.1), make sure to disable that by going to System Settings → Keyboard → Keyboard Shortcuts → Input Sources.

[-]Quinn11mo30

Tmux allows you to set up multiple panes in your terminal that keep running in the background. Therefore, if you disconnect from a remote machine, scripts running in tmux will not be killed. We tend to run experiments across many tmux panes (especially overnight).

Does no one use suffix & disown which sends a command to a background process that doesn't depend on the ssh process, or prefix nohup which does the same thing? You have to make sure any logging that goes to stdout goes to a log file instead (and in this way tmux or screen are better)

Your remark about uv: forgot to mention that it's effectively a poetry replacement, too

Like htop: btm and btop are a little newer, nicer to look at

Also for json: jq. cat file.json | jq for pretty printing json to terminal

[-]shawnghu9mo10

I didn't learn about disown or nohup until recently, because there was no impetus to, because I'd been using tmux. (My workflow also otherwise depended on tmux; when developing locally I liked its method of managing terminal tabs/splits.)

[-]artkpv10mo20

I find it useful to employ text based tools more than UI/UX based tools as they integrate faster with LLMs. For example, AiChat (https://github.com/sigoden/aichat/) does many things: chat with most (all?) models from command line / your text editor, upload files (pdfs, jpg, etc.), execute bash commands and more. It can take stdin and outputs into stdout so you can chain your tools.

[-]Sheikh Abdur Raheem Ali10mo10

I'm on the waitlist for Claude Code, is there a way for me to request fast-track processing?

[-]legobridge10mo10

This is extremely useful for people like me who are just starting out with alignment research (especially parts 3 and 4). Thanks a lot for sharing!

[-]Jan Betley11mo*10

I have one question:

asyncio is very important to learn for empirical LLM research since it usually involves many concurrent API calls

I've lots of asyncio experience, but I've never seen a reason to use it for concurrent API calls, because concurrent.futures, especially ThreadPoolExecutor, work as well for concurrent API calls and are more convenient than asyncio (you don't need await, you don't need the loop etc).

Am I missing something? Or is this just a matter of taste?

[-]Isaac Dunn11mo40

I recently switched from using threads to using asyncio, even though I had never used asyncio before.

It was a combination of:

Me using cheaper "batch" LLM API calls, which can take hours to return a result
Therefore wanting to run many thousands of tasks in parallel from within one program (to make up for the slow sequential speed of each task)
But at some point, the thread pool raised a generic "can't start a new thread" exception, without giving too much more information. It must have hit a limit somewhere (memory? hardcoded thread limit?), although I couldn't work out where.

Maybe the general point is that threads have more overhead, and if you're doing many thousands of things in parallel, asyncio can handle it more reliably.

[-]John Hughes11mo30

Threads are managed by the OS and each thread has an overhead in starting up/switching. The asyncio coroutines are more lightweight since they are managed within the Python runtime (rather than OS) and share the memory within the main thread. This allows you to use tens of thousands of async coroutines, which isn't possible with threads AFAIK. So I recommend asyncio for LLM API calls since often, in my experience, I need to scale up to thousands of concurrents. In my opinion, learning about asyncio is a very high ROI for empirical research.

[+][comment deleted]1y1-2

[+][comment deleted]1y10

LESSWRONG
LW

LESSWRONG
LW

96

Tips and Code for Empirical Research Workflows

96

Ω 40

96

Ω 40

Quick Summary

Part 1: Workflow Tips

Terminal

Integrated Development Environment (IDE)

Git, GitHub and Pre-Commit Hooks

Part 2: Useful Tools

Software/Subscriptions

LLM Tools

LLM Providers

Command Line and Python Packages

Part 3: Experiment Tips

De-risk and extended project mode

Tips for both modes

Tips for extended project mode

Part 4: Shared AI Safety Tooling Repositories

Repo 1: safety-tooling

Repo 2: safety-examples

Acknowledgements