Managing Thousands of DL Experiments and Not Getting Lost

Not many AI researchers or engineers talk about what happens behind the scenes: how to launch (and re-launch) experiments, manage training logs, and assemble final results. This topic is highly opinionated, and everyone does things their own way. Still, if some of those tips may help you manage your research better, or make you think "aha, this is neat,” my job is done.

Codebase Choice

I have seen many people justify their codebase choice by saying, “Well, it is the fastest on benchmarks.” This may be an important point to consider if your goal is to just train a model fast with an already pre-existing recipe but just substituting data. This is one of the few cases when "the fastest codebase" makes sense.

If you are tweaking model architecture, doing complex data processing pipelines, curriculum, optimizers research, training dynamics analysis, and a hundred more research-related things, the fastest codebase is not your objective. It is actively harmful to select a research codebase based only on efficiency. Think about how much time you spend writing, debugging, and changing code compared to launching a training job. Moreover, most of the research is done on relatively small model scales; code efficiency becomes secondary. Still, training speed is important, but there must be a balance.

Please, do not go with Megatron-LM for your project by default. While appealing from a speed perspective, it is terrible to tweak and play with. Consider more lightweight projects that are fast enough and research-friendly:

💡

It doesn’t matter which framework you pick. Just pay attention to its research-readiness. Consider Optimus-DL for this reason. 🙂

Meta Lingua — small (~20 files) project for fast iteration. It is definitely on the more research-y and less training-speed spectrum. Still, it can be a good choice for small experiments. It is definitely fairseq-inspired (as it is also no longer maintained)
TorchTitan — PyTorch-backed framework with cutting-edge PyTorch optimizations and good architecture. Definitely read its guides, sample experiments, and see if it feels right.
NanoChat — The simplest framework on the list, streamlined for GPT-2 training (which is a bit outdated for the newest research). Still, this codebase is one of the easiest to tweak but one of the worst from an architecture perspective.
Optimus-DL — I mean, I have to advertise my project, right? 😄 This framework was created to be a good balance between training speed and research convenience. Here I list just a few features, so make sure to check out the docs and the project itself!
1. Clean architecture, where you can replace any element for research purposes with a registry system.
2. Purely config-driven: model, trainer, criterion, optimizer groups, loggers— you name it—it will be config-defined. This is extremely important, as you want to be able to swap components easily and do easy ablations.
3. Extensive train / eval metrics system. We do not just train models; we analyze training data. A unified logging system allows you to calculate anything anytime with automatic distributed aggregations.
4. Rich data preprocessing system. I really hate when a framework requires you to provide tokenized uint16 files as data and then incurs a mental breakdown if you want to modify the data mixture mid-training (hi, Megatron). Optimus' data system allows you to define complex composable data processing pipelines with multi-dataset sampling.
5. Fast. It implements various optimizations (flat batching, torch compile, custom kernels, TP, SP, HSDP), showing a competitive performance despite being deeply configurable.

Launching Experiments

Nobody talks about this.

Yet this is one of the most important parts in the whole pipeline—no experiments, nothing to analyze, no output (perfect in a way). Alternatively, lost experiments are overlooked results or wasted compute.

It is hard to give platform-independent tips, but while working with several different systems, some principles are transferable.

Leverage tags

How to group related experiments? Luckily, humanity developed a way some hundreds of years ago. Use tags. Most platforms support tags. If not, prepend your experiment’s name with one.

One important rule I keep is having a so-called main tag that defines a minimal group of experiments, where all experiments within the group are comparable and vary only by some parameters. This saves a lot of time and reduces the possibility of human error by making sure you fetch all the experiments within some sweep suite.

In my logs, I have tags like width-depth-1, width-depth-2 and so on. Those define versioning with all versions up to the maximum one considered failed. This way, you can keep track of your experiments’ improvements without mixing up failed attempts (bugs) with final experiments. If only part of the experiments failed, nothing stops you from giving an experiment several main tags.

💡

Use tags and adopt a consistent strategy. This will help you retrieve and compare experiments easily. It is a good thing to also leave a short note for every new tag you make so that you do not forget what you were doing in a week’s time

Streamline exps naming strategy

A dead-simple approach to experiment names is uniqueness. One name = one experiment. Why wouldn’t it be? This makes matching between runs, configs, and names unambiguous and minimizes the risk of overwriting important experiments. It’s also good if all names are human-readable. You don’t need to dump the full config into the name, but the most important parameters are worth mentioning.

💡

Two things a good name conforms to: uniqueness and readability

Your Experiments Will Fail

If they will not fail, you may cancel something by accident, or the whole cluster will go down. Something will happen, and your launch scripts must be able to handle that. All my launch scripts include something like this:

status = get_run_status(run_name)
if status in {RUNNING, PENDING, FINISHED}:
  logger.info(f"Experiment {run_name} is {status}, skipping re-creation")
  return
elif status == FAILED:
  logger.info(f"Experiment {run_name} has failed! Check out the logs")
elif status is None:
  logger.info(f"Launching experiment {run_name}")
else:
  raise RuntimeError(f"Unknown {status = }")

It is pretty self-explanatory, simple, and surprisingly efficient. You place this code in the loop, and now you are able to launch parameter sweeps. It saves a lot of time not to hand-pick failed IDs, look up tested params, and hand-launch failed experiments. More importantly, it is not just saving your time but is also saving you from human error when re-launching something.

💡

Your launch scripts must be able to skip already running exps, re-launch failed ones, and spawn new experiments.

Ditch UI

I mean, it’s in the title. Imagine launching thousands of experiments through UI. It’s just inefficient if not impossible. Leverage CLI for whatever platform you use and ensure that you are able to launch experiments with just scripts. You can use bash, but I would encourage using python straight away, as bash starts to feel clunky as soon as you start doing complicated experiments, planning, calculations, etc.

💡

Using UI to launch experiments is just not scalable. I am not even talking about copy-pasting parameters and configs

Reproducibility

Launch your code twice. Do you get the same results?

Reproducibility of your runs is a good principle for a reason. Mainly, I view it as a sanity check. If your experiments depend on the moon phase, something is off. It can be model init, data sampling, reproducible torch flags—anything.

Additionally, most systems have some kind of preemption or eviction— aka a situation when you use too many resources and your jobs are killed to give room to other people. If your experiments are not reproducible, you will get different results depending on when preemption happened (the closest to the moon phase impact thing), which is bad.

Another good thing to check is whether training a model for 100 steps has the same result as training for 50 steps, stopping, restoring the checkpoint, and resuming until 100 steps (Optimus-DL has such tests!). If results vary, you are missing some information in your checkpoint.

💡

Reproducibility is not just a nice thing to have. It is also a sanity check.

Docker

An important thing about reproducibility is your environment. Luckily, most job-scheduling systems use some type of containers with fixed environments, so this is true for most setups.

If you do not want or cannot use Docker, at least use pyproject.toml with a lock file or requirements.txt dependencies. You do not want your results to differ on a random day just because PyTorch introduced a new bug.

Git All the Way

We all use git for code. I do not know why more people do not extend it to experiments. Store your experiments’ launch scripts in a git repository too. You can develop some structure that works for you: branch per experiment, tag per experiment, directory per experiment, etc. The exact way you do it may vary, but organize your experiments’ management with git too.

Neat trick: if you depend on some external repositories, do not just do git clone <repo> in your launch script. First of all, this clones the most recent version, making your results reproducible. Second, it may pull different code versions in between job preemptions, which just breaks all hell loose. So, the trick is to use git submodules for all your variable dependencies directly in your experiments management git repository.

You can add a submodule via:

git submodule add <dependency> ./<dependency path>

And in your job script, you clone the experiments management repo and check out to the exact commit that was used when launching the job:

git clone <repo> launch && cd launch
git checkout <commit when launched>
git submodule init
git submodule update

pip install <dependency path>

This way, all the repos you depend on are fixed, and each experiment always runs with the same dependency. If you bump up the version of the dependency, you will need to commit a submodule change so that it is used in the following launches.

💡

Git can be used not just to store your main codebase. Use it to store your experiments’ management scripts too! If you depend on external repositories, use gut sub modules to synchronize your launch scripts with external dependencies.

Tracking

If a tree falls in a forest and no one is around to hear it, does it make a sound? Well, the same applies to experiments. If they did not log the data, was it even worth it?

One thing that I’m constantly finding myself in is that it doesn’t matter what metrics tracking platform you use. All of them have very limited visualization capabilities to do anything serious. So, what actually matters is how resilient the logging is so that you do not lose any data.

And the truth is that everything fails; there’s not a single tracking service out there with 100.0% uptime. So, what I do is use several tracking services: mlflow, plain JSONL, wandb. All of those have their pros and cons. Wandb changes their API silently so that my old download scripts do not work anymore and produce wrong results. Mlflow does not easily support bulk export. JSONL may be lost if the service you store it in fails. So, the best option is to have them all in your framework!

Next, I just have scripts that download all data from all runs and use Python to do comparisons and plots.

💡

Let’s face it: all tracking services just cannot do complex comparison plots. So focus on making sure your data is preserved, and use Python to do advanced analysis.

Presenting and Analyzing Data

I will not say much about it, but it is just so painful when someone has interesting experiments implemented but then the data is so badly presented that you cannot understand anything. Even if it is an internal preliminary result that you decide to share, it takes exactly two more minutes to label axes, make a legend, and give it a title.

Seems minor, but the presentation of the data is what matters in the end, right?

All in All

Those tips were gradually developed from doing a large-scale ASR at Yandex, then throughout my master’s degree, and research papers. Perhaps there’s more to it, and not a single setup is perfect. Still, I wanted to share this to save time figuring out those basic truths for other researchers 😁.

Managing Thousands of DL Experiments and Staying Sane