$100K or 100 Days: Trade-offs when Pre-Training with Academic Resources

Dear [QD5y], thank you very much for your review. We appreciate your shared belief about the importance of our research problem. We comment further on the novelty of our work in our general comment.

Dear [DTxJ], thank you for your thoughtful and detailed review! We respond to your concerns below:

• We intentionally exclude multi-node settings from our search space (L84-87; 316-317), because our survey clearly indicated that most academics do not have access to multiple nodes. (We similarly don’t have the multi-node resources to dedicate to this experiment.) But this analysis would make for great future work; for those with resources, our methodology/codebase can be easily extended for this setting.
• We describe our method for measuring the maximum batch size in L154-157. We will do our best to make this even more clear. You are correct that the maximum batch size is a function of the other settings (the model, GPUs, and training methods) and so we measure this independently for each combination of settings. To do so, we compute a training step with batch size = 1, 2, 4, … (consuming increasing amounts of GPU memory) until we run out of memory and consequently identify this maximum.
• We address your comment about novelty in our general comment. Regarding the emphasis on engineering work, we’d also note that “engineering” is strictly invited in COLM’s call for papers (likewise “technological impact” in reviewing guidelines).
• Thanks for bringing Echo (and the other works in training simulation it refers to) to our attention! We would be happy to refer to these in our paper as well. In general, it seems that simulating training step time (by modeling primitive operations) is far more complex than taking a measurement (as we do); hence, that seems to be an open area of research/engineering. Our approach is both accurate and relatively inexpensive for our purposes and scale. It also permits our framework to be extensible (to arbitrary new training methods; and so that our measurement is not tied to any low-level implementation). We continue to believe that “measuring training step time” is a good design decision for our purposes.
• We will try to make it as simple as possible to adopt the codebase (and provide clearer documentation, etc). In general, our framework is an abstraction over HuggingFace Transformers, so we can easily inherit models and all training methods with minimal implementation (or glue) on the side of our framework. We have recently added a standalone training script (using the discovered configuration from our benchmarking) into the codebase. Thus, an end-user could (1) test their desired model in our codebase and then (2) pre-train their model in their own project/codebase (by migrating the training script). Their pre-training should thus be no more difficult than doing pre-training with HuggingFace Transformers. We do have ongoing projects that indeed pre-train models using our framework, but will release these at a later time.

Dear [Ja5N], thanks for your review and comments! We’d just note that:

• Our benchmark and framework are designed such that they can be expanded to new GPUs, models, and training methods. Accordingly, this work can be continually updated to reflect the current state of models and resources. We will do our best to do so and this can also be a community effort.

Dear [ufEP], thank you for your thoughtful review! We are glad you enjoyed reading our paper. Please see our responses below:

• We respond to your note about technical contributions and methodological novelty in our general comment.

• We would like to clarify that our intent is not necessarily to enable academics to train competitive LLMs from scratch. Instead, we want to provide clarity and transparency about what is feasible in academia. Such information is currently not readily available to most academics. That said, we do see some promising research avenues that are available and likely overlooked. E.g. academics may begin conducting controlled studies involving pre-training (smaller) LLMs. I.e. to measure the causal effects of individual variables during pre-training (a scientific practice that is strongly neglected in this area). Such studies could include: testing new training objectives (as in the CLIP example; L24-30), data diets, architectures, optimizers, and so forth. Of course, all experiments require some hyper-parameter tuning, but it should be possible to estimate the resources for these additional runs as well (using our framework) and design studies with this budget in mind. Again, the idea is for academics to estimate what is feasible: a consequence of that will be that there are indeed some infeasible experiments. But some are feasible, and we’d like academics to be in a position to tell the difference.

• You’re correct. We claim that our framework can be easily expanded, but avoided describing this in detail in the paper body. We will definitely elaborate further in revisions. In general, our framework is an abstraction over HuggingFace Transformers. Thus, we can inherit any training method they support. To add new methods (e.g. from new versions of Transformers) into our framework, one must edit src/train.py and experiments/config.py, which jointly map our framework’s search space parameters to Trainer and TrainingArguments from transformers. New models can also be added by implementing the abstract model classes in src/models/__init__.py. We believe these are relatively easy interventions and we will work to make these even easier. These changes automatically propagate into our experiment runners and CLIs. In newer iterations, we also provide a standalone training script (receiving configurations identified by our benchmarking), so that end-users can use our codebase for testing and easily migrate the pre-training experiment to their own projects.

• I appreciate the overview of your codebase. However, it is still challenging to evaluate the effort required to add new optimizations. An acceptable, albeit imperfect, metric to quantify this effort could be lines of code. For instance, if you were to add a new optimization not included in your codebase, such as pipeline parallelism (supported by DeepSpeed), how many lines of code would need to be added/modified?

  Many optimizations also require tuning several knobs to achieve optimal performance. Continuing with the pipeline parallelism example, performance can depend on parameters such as number of pipeline stages, micro-batch size and scheduling strategy. How does the proposed automation framework support tuning such parameters, and including them in the search space?

• The addition of new optimizations causes the search space to grow factorially. Evaluating combinations of 5 optimizations (activation checkpointing, ZeRO stages 0, 1, 2 and offloading) require “a few hours per model-GPU setting”. Adding even a couple more optimizations can likely require multiple days for an exhaustive optimal search. Could you comment on the feasibility of the automation framework with growing optimizations?

• I appreciate the quantification of net improvements observed using optimal combination of memory-saving techniques. However, it would be helpful to know which specific combinations proved optimal. Could you share, for instance, the combination that performed best for the Pythia 160M model on 8 A100 GPUs? Which combination performed the "worst"?

  Could you also comment on whether a reasonably experienced ML researcher could have identified these optimal (or near-optimal) combinations without the aid of your automation framework? I understand that “experience” is subjective, but any insights would be valuable.

  Separately, publishing the optimal combinations you found for different models would be immediately useful to the ML community. Open sourcing this information, building an organized and searchable database, and encouraging academicians to contribute to this database would be highly impactful and serve as a long term solution.