原文地址:https://poolside.ai/blog/introducing-the-model-factory

作者:Joe Rowell(Founding Engineer, R&D)、Eiso Kant(co-CEO & co-Founder)

本文为「Model Factory」系列的第一篇。以下为中英双语对照版本。

TL;DR

传统的基础模型训练方式需要大量人工干预、缺乏有效的改进信号,且迭代周期缓慢。为了解决这些问题,我们构建了 Model Factory——poolside 内部用于快速训练、扩展和实验新型基础模型的系统框架。在本文中,我们分享构建 Model Factory 的方法论及其背后的思考。

Introducing the Model Factory|Model Factory 简介

TL;DR: Traditional foundation model training approaches require substantial manual interaction, little signal for improvement, and slow iteration times. To resolve these issues, we built the Model Factory, Poolside’s internal systems framework for quickly training, scaling, and experimenting with novel foundation models. In this post, we share our methodology and reasoning for building the Model Factory.

摘要:传统的基础模型训练方式需要大量人工介入、几乎没有可用于改进的信号,且迭代周期缓慢。为了解决这些问题,我们构建了 Model Factory——poolside 内部用于快速训练、扩展并实验新型基础模型的系统框架。在本文中,我们将分享构建 Model Factory 的方法论及其背后的推理逻辑。

Introduction|引言

The world of AI is in the midst of a Cambrian explosion: new models, ideas, and techniques surface daily, and keeping up has become a full-time job. New ideas are developed based on intuition, deployed on supercomputer-scale clusters and pushed into production long before the associated theory has caught up. This pace of advancement presents a serious challenge to companies pursuing AGI—namely, how do we scale our ability to evaluate, adopt, and deploy these innovations fast enough to stay ahead?

AI 领域正处于一场”寒武纪大爆发”之中:新模型、新想法和新技术每天层出不穷,光是跟上进度就已经成了一份全职工作。新想法往往基于直觉产生,早在相关理论跟上之前,就被部署到超算规模的集群上并推向生产环境。这样的推进速度给追求 AGI 的公司带来了严峻的挑战——即,我们如何才能足够快地扩展自身评估、采纳和部署这些创新的能力,从而保持领先?

Organizations typically scale along two axes: scaling with people, and scaling with engineering. Hiring top-tier engineers and researchers can push the boundaries of what is possible, but additional engineers can only scale productivity linearly, at most. The pace of AI, by comparison, is growing at least exponentially; linear growth just will not cut it. In other words, hiring ever more engineers and researchers would not allow us to stay ahead without substantial productivity improvements. This requirement clashes with the traditional, linear approach to training models whereby substantial amounts of engineering time was typically dedicated to manually handling training runs.

组织通常沿两个维度进行扩展:靠人力扩展,以及靠工程扩展。招募顶尖的工程师和研究员能够拓展可能性的边界,但增加工程师最多只能让生产力呈线性增长。相比之下,AI 的发展速度至少是指数级的;线性增长根本不够看。换句话说,如果没有实质性的生产力提升,单靠不断招人是无法让我们保持领先的。这一需求与传统的、线性的模型训练方式相冲突——在那种方式下,大量的工程时间通常都被投入到手动处理训练任务上。

Since day one of Poolside, we’ve known that we needed something more: a way to keep up with the pace of our ideas. We needed the ability to test new ideas automatically, benchmark them against reliable baselines, and spin up new models with confidence and speed. We needed the ability to track changes in outcomes back to earlier decisions, a way to conduct hyperparameter searches easily, and a way to inject rigor and reproducibility back into our model training. We even needed a way to fairly schedule all of these competing tasks across a cluster of GPUs.

从 poolside 成立的第一天起,我们就知道自己需要一些更强大的东西:一种能跟上我们想法节奏的方式。我们需要能够自动测试新想法、以可靠的基线对其进行基准测评,并且充满信心且快速地启动新模型。我们需要能够把结果的变化回溯到更早期的决策,需要一种便捷进行超参数搜索的方式,以及一种把严谨性和可复现性重新注入模型训练的方式。我们甚至需要一种在 GPU 集群上公平调度所有这些相互竞争的任务的方式。

What we needed was a model factory. And since one didn’t exist, we built our own.

我们所需要的,正是一座”模型工厂”(model factory)。而既然它并不存在,我们就自己动手造了一座。

In this blog series, we’ll give you an in-depth tour of Poolside’s internal systems framework for training foundation models, which we call the Model Factory. We will show you the Factory’s layout, the high-level goals it solves for, and its full assembly line. We dig into how we use it, the machines it’s made up of, and how we incorporate new machinery. We’ll discuss some useful conveyor belts we’ve built for tying machines together, like GPU-to-GPU weight transfer. We’ll even talk about some ideas that we tried but ultimately rejected, like fully auto-tuned matrix multiplications.

在本系列博客中,我们将带你深入了解 poolside 用于训练基础模型的内部系统框架,我们称之为 Model Factory。我们会向你展示这座工厂的布局、它所要解决的高层目标,以及它完整的”流水线”。我们会深入探讨我们如何使用它、它由哪些”机器”组成,以及我们如何引入新的”机械设备”。我们会讨论一些为串联各台机器而搭建的实用”传送带”,比如 GPU 到 GPU 的权重传输。我们甚至会聊聊一些我们尝试过但最终放弃的想法,例如完全自动调优的矩阵乘法。

But first, let’s get started with a discussion of how the Model Factory came to be, what it enables us to do, and how we power it under the hood.

不过首先,让我们从 Model Factory 的由来谈起:它是如何诞生的、它让我们得以做到哪些事,以及我们在底层是如何驱动它运转的。

A Tour of the Model Factory|Model Factory 巡览

In order to explain why the Model Factory enables us to handle combinatorial complexity, we first need to understand how it evolved, and how that evolution has led to its current design. In order to do this, let’s consider the following scenario: imagine we want to train a foundation model as optimally as possible.

为了解释 Model Factory 为何能让我们应对”组合式复杂度”(combinatorial complexity),我们首先需要理解它是如何演化而来的,以及这一演化过程如何塑造了它当前的设计。为此,让我们设想以下场景:假设我们想要尽可能以最优的方式训练一个基础模型。

Iteration One: Traditional Foundation Model Training|迭代一:传统的基础模型训练

As a first approximation, we’ll start by training some pre-determined architecture on a pre-determined dataset. For the sake of simplicity, we’ll follow the common approach of training models: we load data from a data store, run the model’s forward pass, compute the loss, run the backwards pass, and repeat this process until we’ve finished processing all of our data. Once we’ve finished training the model, we’ll then apply fine tuning to add basic instruction-following capabilities to the model. This leads to a pipeline that looks something like this:

作为第一个近似方案,我们先从在一个预先确定的数据集上训练某个预先确定的架构开始。为简单起见,我们遵循训练模型的常见做法:从数据存储中加载数据,执行模型的前向传播,计算损失,执行反向传播,并不断重复这一过程,直到处理完所有数据。模型训练完成后,我们再进行微调,为模型添加基本的指令遵循能力。这就形成了一条大致如下的流水线:

This is a good first approach, but it has many shortcomings. First, notice that because everything is pre-determined ahead of time, we have trained just one possible model. This model may turn out to be exceptionally good at a variety of tasks—we may, of course, get lucky—but we have no way of knowing how this model could be better. After all, loss curves are great, but they only really make sense in the context of a particular model and dataset. In other words, there is no general evaluation signal with this pipeline. Thus, the solution is clear: in order to improve, we should add an evaluation signal to our pipeline.

这是一个不错的起步方案,但它存在诸多不足。首先请注意,由于所有一切都是事先确定好的,我们训练出来的只是一个可能的模型。这个模型或许在各类任务上都表现出色——我们当然有可能走运——但我们无从得知它究竟还能怎样变得更好。毕竟,损失曲线固然很好,但它只有在特定模型和特定数据集的语境下才真正有意义。换句话说,这条流水线中并没有通用的评估信号。因此,解决方案很清晰:为了改进,我们应当在流水线中加入评估信号。

Iteration Two: Automated Model Evaluations|迭代二:自动化模型评估

Let’s take another look at the pipeline as it is now. The first thing we’re missing is a rigorous way to automatically evaluate the model’s ability on certain downstream tasks. User feedback is certainly king, but it’s not suitable for pre-training scale insights. While there are many benchmarks available for evaluating the capabilities of foundation models, we don’t need to decide on the evaluations themselves at this stage; we just need to introduce a way to evaluate our model at regular (adjustable) intervals. Ideally, we would also be able to adjust which evaluations we use for different training jobs. We can even go one step further and evaluate how the model’s quality changes during pre-training, rather than just during fine tuning. With this in mind, our adjusted pipeline now looks something like this:

让我们再来看看当前这条流水线。我们所缺失的第一样东西,是一种严谨的、能够自动评估模型在特定下游任务上能力的方法。用户反馈固然至关重要,但它并不适合用于获取预训练规模的洞察。虽然目前有许多可用于评估基础模型能力的基准,但在这个阶段我们并不需要确定具体使用哪些评估;我们只需引入一种以固定(且可调节)间隔评估模型的机制即可。理想情况下,我们还应能针对不同的训练任务调整所使用的评估。我们甚至可以更进一步,去评估模型质量在预训练过程中(而不仅仅是在微调过程中)是如何变化的。有了这一考量,我们调整后的流水线现在大致如下:

Running evaluations during pre-training also requires us to have an inference service that we can invoke during pre-training. This presents several technical constraints: for instance, we need to be able to transfer weights between nodes that are dedicated for training and those that we will use for inference. In the context of regular pre-training, we can get away with using a traditional checkpointing system but we will soon see that paying attention to the speed of this transfer is essential for more advanced training techniques.

在预训练期间运行评估还要求我们拥有一个可以在预训练过程中调用的推理服务。这带来了若干技术约束:例如,我们需要能够在专用于训练的节点和用于推理的节点之间传输权重。在常规预训练的情境下,我们尚可将就着使用传统的检查点(checkpointing)系统,但我们很快就会看到,对于更高级的训练技术而言,关注这种传输的速度是至关重要的。

Iteration Three: Reinforcement Learning from Code Execution|迭代三:基于代码执行的强化学习

Now that we have an evaluation system, we can think further about how we might improve our models. One of the ways that humans learn and grow is through external feedback. Similarly, we can improve our model’s abilities by going beyond next-token prediction and using extensive reinforcement learning techniques that lie outside of typical approaches to fine tuning. This approach has required us to treat reinforcement learning as a primary concern, and we have now implemented a separate, modular reinforcement learning system that contains several reinforcement learning techniques (and allows for both synchronous and asynchronous reinforcement learning).

既然我们已经有了评估系统,就可以进一步思考如何改进模型。人类学习和成长的方式之一便是借助外部反馈。类似地,我们可以超越单纯的下一个 token 预测,借助大量超出典型微调范畴的强化学习技术来提升模型的能力。这一路径要求我们把强化学习作为首要考量,为此我们如今实现了一套独立、模块化的强化学习系统,其中包含多种强化学习技术(并且同时支持同步和异步强化学习)。

Our modular reinforcement learning system also includes several actor components that we can dynamically instantiate into our workflows. This has a direct benefit of allowing us to spin up new agents and rewards without modifying the surrounding machines in the Model Factory.

我们这套模块化的强化学习系统还包含若干 actor 组件,我们可以将它们动态地实例化到工作流中。这带来的一个直接好处是:我们能够启用新的智能体(agent)和奖励机制,而无需改动 Model Factory 中周边的其他”机器”。

Additionally, we have also put a substantial amount of effort into building a secure and isolated code environment for our foundation models to self play and explore. This environment includes roughly a million public repositories (and most of their revisions) containerized with their test suite executable for us to use to improve our models. Tasks themselves are served to our reinforcement learning algorithms through a highly-scalable and dedicated task engine, enabling easy orchestration of millions of complicated agentic workflows in parallel. The task engine itself is highly versatile and suitable for use in online reinforcement learning workloads. By combining the task engine with our code execution environment, we are able to teach agentic workflows to handle real-world software development.

此外,我们还投入了大量精力,为基础模型构建了一个安全且隔离的代码环境,供其进行自我博弈(self play)与探索。该环境包含约一百万个公共代码仓库(及其大部分历史版本),它们连同可执行的测试套件一起被容器化,供我们用来改进模型。任务本身通过一个高度可扩展的专用任务引擎被分发给我们的强化学习算法,从而能够轻松地并行编排数百万个复杂的智能体工作流。这个任务引擎本身用途极其广泛,非常适合用于在线强化学习的工作负载。通过将任务引擎与我们的代码执行环境相结合,我们得以教会智能体工作流去处理真实世界的软件开发任务。

Iteration Four: Automated Architecture Ablations|迭代四:自动化架构消融实验

We can now turn our attention to how we experiment with model architectures. It’s clear that experimenting with architecture will lead to changes in model performance. After all, a model’s architecture will strongly influence how it performs at a given task. However, varying the model’s architecture comes with some tough choices. On the one hand, not all aspects of a model are dictated solely by technical performance: for example, small models can provide low latency inference across massive context lengths, but at the cost of performing poorly compared to larger models. Yet these requirements are actually a gift, as they allow us to constrain an otherwise unbounded search space. What’s more, we still have a huge search space to explore across model layouts, parameter sizes, layers and so on. We can even consider novel architectures, like hybrids between transformers and linear attention variants. If you can imagine a tweak, we can efficiently try it with the Factory. But there’s still the elephant in the room: experimentation with large language models is an incredibly time consuming and costly endeavour. In order to optimize the usage of our GPU cluster, we will need to consider every possible optimization and observability technique during pre-training. Thankfully, we’ve collected a huge bag of tricks at Poolside for optimizing our pre-training codebase. We also make use of several techniques that allow us to experiment on small models in order to predict their scale-up performance, such as mu-P.

现在我们可以把注意力转向如何对模型架构进行实验。显然,对架构进行实验会带来模型性能的变化。毕竟,一个模型的架构会强烈影响它在特定任务上的表现。然而,改变模型架构也伴随着一些艰难的取舍。一方面,模型的并非所有方面都仅由技术性能决定:例如,小模型能够在超长上下文长度下提供低延迟推理,但代价是其表现不如更大的模型。然而这些需求其实是一份”礼物”,因为它们让我们得以约束一个原本无界的搜索空间。更重要的是,在模型布局、参数规模、层数等方面,我们仍有巨大的搜索空间有待探索。我们甚至可以考虑一些新颖的架构,比如 transformer 与线性注意力变体之间的混合体。只要你能想到某种改动,我们就能借助这座工厂高效地去尝试它。但仍有一个显而易见却被回避的问题:“对大语言模型进行实验”是一件极其耗时且昂贵的事。为了优化 GPU 集群的使用,我们需要在预训练过程中考量一切可能的优化和可观测性技术。所幸,在 poolside 我们已经积累了一大堆用于优化预训练代码库的”锦囊妙计”。我们还运用了若干技术,使我们能够在小模型上做实验,从而预测其扩大规模后的性能,例如 mu-P。

At this stage it’s worth discussing how we might handle the huge complexity of scheduling all of this additional work on a fixed-size cluster of GPUs. At first glance it may appear that we can simply schedule jobs arbitrarily across our GPUs using a preexisting scheduler, like SLURM. However, this is actually insufficient for our setting because we now have dependencies between tasks: our evaluations framework requires access to inference, which needs to wait for a checkpoint to be produced by the model currently undergoing pre-training. In order to manage this complexity, we designed and implemented a single Kubernetes orchestrator that handles all of these moving pieces in a consistent and rigorous way. As a side benefit, we can use such an orchestrator to schedule additional, useful tasks when GPUs are idle, such as synthetic data generation.

在这个阶段,值得讨论一下我们该如何应对在固定规模的 GPU 集群上调度所有这些额外工作所带来的巨大复杂度。乍看之下,我们似乎可以直接使用现成的调度器(如 SLURM)在 GPU 之间任意调度任务。然而,这对我们的场景实际上是不够用的,因为如今任务之间存在依赖关系:我们的评估框架需要访问推理服务,而推理又需要等待当前正在预训练的模型产出一个检查点。为了管理这种复杂度,我们设计并实现了一个单一的 Kubernetes 编排器(orchestrator),以一致且严谨的方式处理所有这些相互牵动的环节。作为一个附带好处,我们可以利用这样的编排器,在 GPU 空闲时调度额外的、有用的任务,比如合成数据生成。

Iteration Five: Automated Data Refinement|迭代五:自动化数据精炼

Looking back, we can see that there is one untouched piece of the original pipeline: the initial dataset. Much like raw materials in a real factory, the output of our Model Factory depends heavily on the quality and quantity data that we feed into our machines. On the one hand, it’s well known that AI models are orders of magnitude less efficient at learning than humans, and we typically need to train on trillions of tokens to see the emergence of latent abilities. However, wantonly training on trillions of tokens is not a recipe for success; after all, no one has ever learned to code by reading Great Expectations. Instead, we also need to make sure that our model sees enough well-mixed, high quality data during pre-training. This is further complicated by the fact that while data cleaning and refining techniques are well understood, it’s often difficult to understand the interactions between different data refinement steps and downstream performance. Put differently, what we really need is a form of lineage across our dataset refinements, the ability to link lineage back to model performance and large scale semi-automated experimentation.

回过头看,我们会发现原始流水线中还有一个未被触及的部分:初始数据集。就像真实工厂中的原材料一样,我们 Model Factory 的产出在很大程度上取决于我们喂给”机器”的数据的质与量。一方面,众所周知,AI 模型的学习效率比人类低好几个数量级,我们通常需要在数以万亿计的 token 上训练,才能看到潜在能力的涌现。然而,毫无节制地在万亿级 token 上训练并非成功之道;毕竟,从没有人是靠读《远大前程》学会编程的。相反,我们还需要确保模型在预训练过程中见到足够多、混合良好且高质量的数据。更棘手的是,尽管数据清洗与精炼的技术本身已经相当成熟,但要理解不同数据精炼步骤与下游性能之间的相互作用往往十分困难。换句话说,我们真正需要的是一种贯穿数据集精炼过程的”血缘”(lineage)机制,一种把血缘回溯到模型性能的能力,以及大规模的半自动化实验能力。

There are also other practical concerns when it comes to the roles of our data pipelines. On the one hand, the sheer quantity of data needed for pre-training necessitates that we pay careful attention to the speed of ingestion and processing. Frankly speaking, using a slow or unscalable approach would be so disastrous to training performance that we might never actually get around to evaluating a new idea. Equally, poor quality data ingestion can massively harm the ability of a foundation model to generalise and perform; an ingestion pipeline must not only be fast, it must also be as accurate as possible. On the flip side, we regularly need to generate hundreds of billions of new synthetic data tokens, as well as ingest, re-ingest and rewrite large parts of pre-existing web data, requiring us to process massive amounts of data. Yet, as we’ve already built an inference engine in our Factory, we can simply add a synthetic data generation machine and connect it to the inference engine.

在数据流水线的作用方面,还存在其他一些现实考量。一方面,预训练所需数据的庞大数量,迫使我们必须格外关注数据摄取与处理的速度。坦率地说,采用缓慢或不可扩展的方式,对训练性能的破坏会大到我们可能永远也轮不到去评估一个新想法。同样地,低质量的数据摄取会严重损害基础模型的泛化和表现能力;一条摄取流水线不仅要快,还必须尽可能地准确。另一方面,我们经常需要生成数千亿个新的合成数据 token,同时还要摄取、重新摄取并改写大量已有的网络数据,这要求我们处理海量的数据。不过,由于我们已经在工厂中构建了推理引擎,我们只需添加一台合成数据生成”机器”,并将其与推理引擎相连即可。

Iteration Six: Automated Data Mixing|迭代六:自动化数据混合

Now that we’ve produced the relevant datasets, we also need a way of evaluating how these datasets should be mixed together and supplied to the model. We can, of course, simply shuffle the data for training, but we also need a clever way of determining what the data blend should be, as well as the order in which the data is delivered. This all adds up to multiple new components in our Factory. A further concern is that data is typically supplied in training as a single, massive dataset, making it difficult to adjust for mistakes or changes in scope. Additionally, supplying such a dataset introduces many practical concerns that make training difficult. For instance, materialising a single dataset requires us to fix our cluster size ahead of time; and training is delayed until the entire dataset has been produced. Moreover, the cost of simply distributing a singular massive dataset across a GPU cluster—and the cost of re-materializing—is a substantial barrier to efficient training. From an experimentation perspective, serializing said dataset also couples the data blend and data order to the dataset, necessitating that we redistribute the dataset each time we run an experiment with a new blend configuration. We address these issues by instead streaming data into training from our upstream data source, allowing us to easily change data mixes or quantities spontaneously across different experiments and even live during model training.

既然我们已经产出了相关的数据集,我们还需要一种方法来评估这些数据集应该如何混合并供给模型。我们当然可以简单地把数据打乱后用于训练,但我们还需要一种巧妙的方式来确定数据的配比(data blend)以及数据供给的顺序。这一切累加起来,就为我们的工厂增添了多个新组件。另一个需要关注的问题是:数据在训练中通常是以单一、庞大的数据集形式供给的,这使得针对错误或范围变更进行调整变得很困难。此外,供给这样一个数据集还带来了许多让训练变得棘手的现实问题。例如,物化(materialise)单个数据集要求我们提前固定集群规模;而且训练必须等到整个数据集生成完毕后才能开始。更重要的是,仅仅是把一个庞大的单一数据集分发到 GPU 集群上的开销——以及重新物化的开销——就构成了高效训练的一大障碍。从实验的角度看,将该数据集序列化还会把数据配比和数据顺序与数据集本身耦合在一起,使得我们每次用新的配比配置做实验时都不得不重新分发数据集。为了解决这些问题,我们转而采用从上游数据源以流式(streaming)方式将数据送入训练的做法,从而能够在不同实验之间——甚至在模型训练进行中——轻松、即时地更改数据混合方式或数据量。

Now that we have all of the building blocks in place, we’re at the point where we can actually build models. However, training models is still a long and tricky process, and we still need to provide some form of observability to our Factory. We achieve this observability by adding dashboards and dials to the machines in our Factory, streaming a large number of metrics to various dashboard services. This approach allows us to keep an eye on everything to do with model training, from loss curves to iteration time and everything in between. We also capture a large number of production metrics, enabling us to ensure that changes made to the Factory itself are both stable and beneficial.

既然所有的构建模块都已就位,我们现在便来到了可以真正构建模型的阶段。然而,训练模型仍然是一个漫长而棘手的过程,我们还需要为工厂提供某种形式的可观测性。我们通过为工厂中的机器添加仪表盘和”旋钮”来实现这种可观测性,将大量指标以流式方式送入各种仪表盘服务。这种做法让我们能够密切关注与模型训练相关的一切,从损失曲线到迭代时间,以及介于两者之间的所有内容。我们还采集了大量的生产指标,从而能够确保对工厂本身所做的改动既稳定又有益。

We also need a way of doing some form of quality assurance (i.e. vibe checking) on the models that come out of our Factory. This sort of quality assurance goes beyond automated evaluation: we want to manually try and verify each and every model that comes off the production line. We do this by adding an additional verification platform that deploys freshly produced models to our inference service, allowing us to stress test produced models in our daily working lives.

我们还需要一种方式,对从工厂产出的模型进行某种形式的质量保证(即”体感检查”,vibe checking)。这类质量保证超出了自动化评估的范畴:我们希望亲手试用并验证每一个下线的模型。为此,我们增设了一个额外的验证平台,将新产出的模型部署到我们的推理服务上,使我们能够在日常工作中对产出的模型进行压力测试。

There is one, final thing to discuss about the Model Factory: namely, automation. We’ve added a considerable amount of efficiency when it comes to building new AI models, but if this effort requires substantially more engineering time, it will all have been for nothing. In order to handle this complexity, we represent our data and architectural experiments as programmatic assets. Concretely, we manifest our changes to data refinement and model architecture as structured experiment files that are automatically queued and executed by the orchestrator. Following this approach gives us two key benefits. On the one hand, by representing configurations as code, we gain the ability to fully reproduce experimental results in a rigorous fashion. Secondly, this approach leads to an extremely quick feedback loop. Scheduling a new sweep, for instance, such as over batch size or learning rates, takes fewer than ten minutes, as opposed to days. This reduces the friction of trying out new ideas, making it easier for us to move new innovations into production.

关于 Model Factory,还有最后一件事值得讨论:自动化。在构建新 AI 模型方面,我们已经大幅提升了效率,但如果这份成果需要大量额外的工程时间来支撑,那么这一切就都白费了。为了应对这种复杂度,我们把数据实验和架构实验表示为可编程的资产(programmatic assets)。具体来说,我们把对数据精炼和模型架构的改动体现为结构化的实验文件,由编排器自动排队并执行。采用这种方式带来了两大关键好处。一方面,通过将配置表示为代码,我们获得了以严谨方式完全复现实验结果的能力。其次,这种方式带来了极快的反馈循环。举例来说,调度一次新的扫描实验(例如对批大小或学习率的扫描)所需时间不到十分钟,而不再是数天。这降低了尝试新想法的阻力,使我们更容易将新的创新推向生产环境。

What does the Model Factory actually give us?|Model Factory 究竟为我们带来了什么?

We can now begin to think about what this pipeline gives us beyond what we have already discussed. First, notice that the subsystems in this pipeline are rather uncoupled, and thus each subsystem can be arbitrarily complicated. This decoupling is extremely useful in practice: we can, for instance, adapt or modify the data pipelines without needing to adjust the rest of the Factory in any way. In a similar vein, we can adjust data mixes, model architectures, or fine tuning techniques at will, without needing to adjust anything else. This enables us to handle the combinatorial complexity of building foundation models. By simply adapting a single piece of the Factory, we can suddenly explore new component mixes and architectures. Lastly, we note that the pipeline described above is fully modular, allowing for meta strategies when it comes to training models. For instance, we may wish to vary the amount of fine tuning that we use on a particular model. In other words our ensemble of machines can be altered at will, enabling us to be purely pragmatic and experiment-driven. We can even insert additional machines or systems as we need them, allowing the factory to grow almost at will, further handling and supporting our needs.

现在我们可以开始思考,除了前面已经讨论过的内容之外,这条流水线还为我们带来了什么。首先请注意,这条流水线中的各个子系统是相当解耦的,因此每个子系统都可以任意复杂。这种解耦在实践中极为有用:例如,我们可以调整或修改数据流水线,而完全无需改动工厂的其余部分。同理,我们可以随心所欲地调整数据配比、模型架构或微调技术,而无需改动其他任何东西。这使我们能够应对构建基础模型时的组合式复杂度。仅仅通过调整工厂中的某一个部件,我们就能够立刻探索新的组件组合和架构。最后,我们要指出,上述流水线是完全模块化的,从而在训练模型时允许采用”元策略”(meta strategy)。例如,我们可能希望改变对某个特定模型所使用的微调量。换句话说,我们这一整套”机器组合”可以被随意改动,使我们能够纯粹以务实和实验驱动的方式行事。我们甚至可以按需插入额外的机器或系统,让工厂几乎能随意扩张,从而进一步满足并支撑我们的需求。

It’s hard to overstate the impact the Model Factory has had on our ability to operate at the frontier of AI. Our previous deduplication efforts, for instance, used to take over a week’s worth of engineering time to schedule and run; now, we can deduplicate datasets effortlessly with the click of a button. Experiments that used to take weeks to organise and schedule are now handled by the orchestrator in under an hour—often as few as ten minutes, unlocking a new level of productivity for our team. Qualitatively speaking, we’ve also noticed that the Factory allows for a far nicer developer experience compared to our previous pipelines, making it much easier to onboard new engineers.

Model Factory 对我们在 AI 前沿开展工作的能力所产生的影响,怎么强调都不为过。举例来说,我们以前的去重(deduplication)工作,光是安排和运行就要花费一周以上的工程时间;如今,我们只需点一下按钮就能轻松完成数据集去重。过去需要数周来组织和调度的实验,现在由编排器在一小时之内——往往只需十分钟——就能处理完毕,为我们团队解锁了全新水平的生产力。从定性的角度来说,我们还注意到,相比以往的流水线,这座工厂带来了好得多的开发者体验,使得新工程师的上手变得容易得多。

Above all, the Factory gives us a huge amount of leverage to rapidly try out new ideas and build new foundation models. Instead of waiting all week for results, we conduct supercomputer-scale experiments over a weekend. Instead of manually triggering pipeline steps, we can simply describe an experimental configuration and the Factory handles everything else for us, even triggering new steps automatically.

最重要的是,这座工厂赋予了我们巨大的杠杆效应,让我们能够快速尝试新想法并构建新的基础模型。我们不再需要苦等一整周才能拿到结果,而是能在一个周末内完成超算规模的实验。我们也不再需要手动触发流水线的各个步骤,只需描述一份实验配置,工厂就会替我们处理好其余的一切,甚至会自动触发新的步骤。

What does the Model Factory look like in practice?|Model Factory 在实践中是什么样子?

Now that we have a good overview of how the Model Factory works in theory, we can begin to look at how we actually built it at Poolside. We’ll do a deeper dive in future posts, so for now let’s stick to the high-level details.

既然我们已经从理论上对 Model Factory 的运作方式有了不错的整体认识,现在我们可以开始了解在 poolside 我们究竟是如何将它构建出来的。我们会在后续的文章中进行更深入的剖析,因此现在我们先聚焦于高层次的细节。

High-level layout of the Model Factory|Model Factory 的高层布局

In practice, the Model Factory can be viewed as several systems stacked under a single Kubernetes orchestrator. The orchestrator is responsible for deploying and prioritising jobs across our 10K H200 GPU cluster. This process, it turns out, requires the building of several additional constraints and handling. For instance, some parts of our data pipelines (such as deduplication) do not actually require GPUs to run successfully; indeed, we can run them entirely on CPU nodes, saving our precious GPUs for other tasks. In order to do this, we extended our orchestrator to handle multiple different types of node pools, allowing us to easily schedule certain tasks on different types of nodes. We note, of course, that the Factory’s design is broadly agnostic to the underlying accelerator that is used, and the exact number of GPUs or nodes can be adjusted at will without changing the underlying ideas. The deployment of the Factory to the cluster is fully automated using Helm charts, with Terraform being used to manage the cluster itself.

在实践中,Model Factory 可以被看作是堆叠在单个 Kubernetes 编排器之下的若干系统。该编排器负责在我们由 1 万块 H200 GPU 组成的集群上部署任务并确定其优先级。事实证明,这一过程需要构建若干额外的约束与处理逻辑。例如,我们数据流水线中的某些部分(如去重)实际上并不需要 GPU 就能成功运行;确实,我们可以完全在 CPU 节点上运行它们,从而把宝贵的 GPU 留给其他任务。为此,我们扩展了编排器,使其能够处理多种不同类型的节点池(node pool),从而让我们能够轻松地把某些任务调度到不同类型的节点上。当然,我们要指出,工厂的设计对所使用的底层加速器基本上是无感知(agnostic)的,GPU 或节点的具体数量可以随意调整,而无需改变其底层思路。工厂到集群的部署完全通过 Helm charts 自动化完成,而集群本身则使用 Terraform 进行管理。

Now that we have an orchestrator, we need to think about how our experiments are structured. As previously mentioned, all of our experiments (for both data processing and architecture) are codified (i.e. every change to either model architecture or data processing steps is represented as a piece of versioned code). Concretely, we write all data experiments as Spark data frames, providing effortless scalability across all of our data experiments. In contrast, our model architecture experiments are typically specified as Python dictionaries that are then programmatically validated and converted into structured Markdown files. This conversion enables us to easily extend previous experiments in new directions—we simply need to extend an existing dictionary to cover new parameters. In both cases, once an experiment is defined it is immediately deployed as part of our continuous integration workflows.

既然我们有了编排器,接下来就需要思考我们的实验是如何组织的。如前所述,我们所有的实验(无论是数据处理还是架构方面的)都被代码化了(即对模型架构或数据处理步骤的每一处改动,都被表示为一段受版本控制的代码)。具体来说,我们把所有数据实验都写成 Spark data frame,从而为所有数据实验提供了轻松的可扩展性。相比之下,我们的模型架构实验通常被指定为 Python 字典,随后经过程序化验证并转换为结构化的 Markdown 文件。这种转换使我们能够轻松地把以往的实验朝新的方向扩展——我们只需扩展一个现有的字典,使其涵盖新的参数即可。在这两种情况下,一旦实验被定义完成,它就会作为我们持续集成(CI)工作流的一部分被立即部署。

Other than providing a seamless way to work, executing experiments as part of continuous integration also provides a huge level of reproducibility: after all, we can recreate the experiment by simply re-triggering the job. Lastly, for data experiments in particular, we can manifest the resulting data as assets in an Iceberg data lake, providing easy versioning and reliability, which is critical for machine-learning research and implementation.

除了提供一种无缝的工作方式之外,把实验作为持续集成的一部分来执行还带来了极高的可复现性:毕竟,我们只需重新触发相应任务,就能重建该实验。最后,特别是对于数据实验,我们可以把生成的数据作为资产(asset)物化到一个 Iceberg 数据湖中,从而提供便捷的版本管理和可靠性,而这对机器学习的研究与实现而言至关重要。

Dataset management|数据集管理

Crucially, Iceberg tables also allow us to query any retained version of any data asset, complete with full metadata and lineage. In fact, any Engine in the Factory that is Apache Iceberg-compatible can also query data assets, enabling powerful automatic workflows. In practice, we query our Iceberg tables by using Starburst’s SQL capabilities, allowing us to easily visualize and inspect the results of our queries.

至关重要的是,Iceberg 表还允许我们查询任何数据资产的任何被保留下来的版本,并附带完整的元数据和血缘信息。事实上,工厂中任何兼容 Apache Iceberg 的引擎(Engine)也都能查询这些数据资产,从而支撑起强大的自动化工作流。在实践中,我们借助 Starburst 的 SQL 能力来查询 Iceberg 表,使我们能够轻松地可视化并检视查询结果。

At this stage, it’s worth mentioning that training runs typically consume hundreds of data assets, with each asset representing a full dataset. We achieve this by using the Dagster platform, which acts as the beating heart of the Model Factory. Dagster provides an easy way to orchestrate different data pipeline steps and model training. Additionally, Dagster also allows us to conditionally run code on produced assets, which in turn allows us to quickly check and exclude a certain dataset if the quality is too low, or if data packing was otherwise inefficient. More generally, Dagster also enables us to launch training runs directly from the web-UI and to specify automatic sub-steps during training, such as offloading checkpoints into inference or triggering evaluations during our workflow. These changes are incredibly easy to make: we simply need to provide a corresponding configuration file. We can even use Dagster to quickly branch on existing assets and workflows, automatically launch new workflows according to either a pre-existing schedule or during other workflows, and visualize the results of existing experiments.

在这个阶段,值得一提的是,一次训练通常会消耗数百个数据资产,其中每个资产都代表一个完整的数据集。我们借助 Dagster 平台来实现这一点,它扮演着 Model Factory 跳动的”心脏”的角色。Dagster 提供了一种便捷的方式来编排不同的数据流水线步骤和模型训练。此外,Dagster 还允许我们有条件地在已产出的资产上运行代码,进而使我们能够快速检查并排除某个质量过低、或数据打包(data packing)效率不佳的数据集。更一般地说,Dagster 还使我们能够直接从 Web UI 启动训练,并在训练过程中指定自动化的子步骤,比如把检查点卸载到推理服务中,或在工作流进行中触发评估。做出这些改动极其简单:我们只需提供一份对应的配置文件即可。我们甚至可以用 Dagster 在已有的资产和工作流上快速分支,根据预设的时间表或在其他工作流进行时自动启动新的工作流,并可视化已有实验的结果。

One indispensable area of Dagster’s automatic job launching is our evaluations workflow. Instead of manually triggering evaluations, we configure evaluations to be launched on an experiment-by-experiment basis. These evaluations are typically executed automatically every 100-1000 training steps, with the exact value set in the experiment’s config file.

Dagster 的自动任务启动有一个不可或缺的应用领域,那就是我们的评估工作流。我们不再手动触发评估,而是将评估配置为按实验逐个启动。这些评估通常每隔 100 至 1000 个训练步就自动执行一次,具体数值在该实验的配置文件中设定。

We typically run a variety of evaluations against benchmarks that are indicative of general foundation model capabilities, ranging from general knowledge and reasoning capabilities to ability at coding-specific tasks, with all results automatically stored and reported. Crucially, because evaluations are a somewhat lightweight inference workload, we allow evaluation workflows to either be pre-empted in favor of long running jobs, or prioritized in order to provide even faster feedback loops.

我们通常会针对一系列能够反映基础模型通用能力的基准运行各种评估,涵盖从通用知识和推理能力到编程专项任务能力等各个方面,所有结果都会被自动存储和上报。至关重要的是,由于评估属于相对轻量的推理工作负载,我们允许评估工作流要么被抢占(以让位于长时间运行的任务),要么被提升优先级(以提供更快的反馈循环)。

Measure all the things|度量一切

As previously mentioned, we capture a huge quantity of metrics across our Model Factory. On the one hand, we typically stream training metrics—loss curves, gradients, evaluations—into the Neptune data platform. Crucially, Neptune allows us to compare metrics across different runs of the Factory, which gives us much deeper insights into how our ideas play out in practice during training. However, training metrics only go so far, and the proof is very much in the pudding when it comes to model quality. To enable us to more easily introspect produced models, our team built a suite of internal tools, collectively called Podium. Podium allows us to manually query and interact with models, compare models against each other with evaluations, and interact with diverse fine-tuning datasets for further model improvement. We stress that model evaluation—or “vibe checking”—is a company-wide effort at Poolside; everyone, from engineers to members of the go-to-market team, regularly provide feedback on models that come off the Model Factory’s production line. This feedback really helps us to train models that are better than we could otherwise achieve by simply chasing better metrics alone.

如前所述,我们在整座 Model Factory 中采集了海量的指标。一方面,我们通常会把训练指标——损失曲线、梯度、评估结果——以流式方式送入 Neptune 数据平台。至关重要的是,Neptune 让我们能够跨工厂的不同运行来对比指标,这使我们对自己的想法在训练过程中的实际表现有了深得多的洞察。然而,训练指标的作用是有限的,谈到模型质量,终究还是”是骡子是马拉出来遛遛”。为了让我们能够更轻松地内省(introspect)所产出的模型,我们团队构建了一整套内部工具,统称为 Podium。Podium 让我们能够手动查询模型并与之交互、通过评估将不同模型相互比较,以及与多样化的微调数据集交互以进一步改进模型。我们要强调的是,模型评估——即”体感检查”——在 poolside 是一项全公司参与的工作;从工程师到市场推广(go-to-market)团队的成员,每个人都会定期为下线的模型提供反馈。这些反馈切实地帮助我们训练出比单纯追逐更漂亮的指标所能得到的更好的模型。

On the other hand, we also capture production metrics for all of the sub-systems in the Factory and stream the results into a series of Grafana dashboards. In contrast to our Neptune metrics, our Grafana metrics are typically focused on the reliability and performance of the systems in the Factory, and not the models that are being trained. These metrics serve as useful bellweathers for changes; we can directly check how new ideas or tweaks impact the running of the Factory in practice. In the past, these Grafana metrics have been crucial in helping us to diagnose issues with new techniques, as well as helping us to identify failing nodes.

另一方面,我们还采集工厂中所有子系统的生产指标,并将结果以流式方式送入一系列 Grafana 仪表盘。与我们的 Neptune 指标不同,Grafana 指标通常关注的是工厂中各系统的可靠性和性能,而非正在训练的模型。这些指标充当了衡量变化的有用”风向标”;我们可以直接查看新想法或调整在实践中对工厂运行的影响。过去,这些 Grafana 指标在帮助我们诊断新技术的问题、以及帮助我们识别出故障节点方面,都发挥了至关重要的作用。

Pre-training at scale|大规模预训练

Moving on from evaluations, we land on our pre-training codebase, Titan. Titan is a state-of-the-art pre-training library based on PyTorch. Similar to other leading training frameworks, Titan is equipped with many low-level tweaks to ensure peak performance, along with additional observability tools that are not commonly found in public pre-training codebases. On the performance front, we make substantial use of PyTorch’s “torch.compile” feature to ensure that our auto-generated kernels are sufficiently efficient, and we also make heavy use of Triton and CUDA kernels where necessary. In addition, Titan provides the exceptionally useful feature of being embedded deeply inside our inference codebase, Atlas. This design decision enables us to quickly move new model architectures into production, as we can always ensure that we have a close match between our inference implementation and our training implementation. In fact, we typically task the Model Factory with automatically checking this equivalence for us during experiments, leaving us to focus on the low-level optimisations necessary for peak performance. Atlas itself is a large wrapper library around other popular inference libraries. We note that Atlas serves multiple different workloads—evaluations, customer inference in production, pre-training experiments—across a wide range of hardware platforms. Indeed, Atlas runs natively on NVIDIA, AMD and Trainium accelerators, giving us a large deal of choice when it comes to deployment.

从评估继续往下,我们来到了我们的预训练代码库 Titan。Titan 是一个基于 PyTorch 的最先进(state-of-the-art)的预训练库。与其他领先的训练框架类似,Titan 配备了许多底层优化以确保峰值性能,还配有一些在公开预训练代码库中并不常见的额外可观测性工具。在性能方面,我们大量使用 PyTorch 的 “torch.compile” 特性,以确保自动生成的算子(kernel)足够高效,并在必要时大量使用 Triton 和 CUDA 算子。此外,Titan 还提供了一个极其有用的特性:它被深度嵌入到我们的推理代码库 Atlas 之中。这一设计决策使我们能够快速将新的模型架构推向生产环境,因为我们始终可以确保推理实现与训练实现之间高度吻合。事实上,我们通常会交由 Model Factory 在实验过程中自动为我们检查这种等价性,从而让我们能够专注于实现峰值性能所必需的底层优化。Atlas 本身是一个包裹在其他流行推理库之外的大型封装库。我们要指出,Atlas 服务于多种不同的工作负载——评估、生产环境中的客户推理、预训练实验——并跨越广泛的硬件平台。确实,Atlas 能够原生运行在 NVIDIA、AMD 和 Trainium 加速器上,为我们在部署时提供了极大的选择空间。

Practically speaking, we either schedule an inference task as a dedicated job—requiring a fixed number of nodes in our cluster—or as an elastic task that scales up and down on demand. This automatic scaling both ensures that our GPUs are never idle, and that important experiments are not blocked by lower priority tasks.

从实际操作来看,我们要么把推理任务作为一个专用任务来调度——这需要在集群中占用固定数量的节点——要么把它作为一个可按需弹性伸缩的任务来调度。这种自动伸缩既确保了我们的 GPU 永不空闲,也确保了重要的实验不会被优先级较低的任务所阻塞。

Code execution at scale|大规模代码执行

We now turn our attention to our code execution platform. Our code execution platform provides a secure and isolated environment for models to solve tasks by writing and executing code. In order to provide such an environment for our models, we’ve built a platform that automatically builds GitHub repositories inside OCI containers, with the exact build mechanism depending on a series of heuristics. In practice, we build containers depending on a series of heuristic rules, with build tasks scheduled on an arbitrarily-sized node pool. We also use a robust observability system for reporting build failures, enabling us to further refine, debug, and improve our build heuristics. In terms of scale, our code execution environment contains around a million repositories (and many of their revisions) across many different programming languages, and we constantly add and build new repositories into the environment. We stress that the process of building containers is constantly evolving, and we have recently seen a huge amount of success from using agent-based building. This is a particularly exciting step for us, as an agent built using the Factory is now working to improve the Factory that built it. More broadly, our code execution platform also serves as a core component in our model evaluations, providing a strong feedback signal for model improvement on software development tasks.

现在我们把注意力转向代码执行平台。我们的代码执行平台为模型提供了一个安全且隔离的环境,让模型通过编写和执行代码来解决任务。为了给模型提供这样的环境,我们构建了一个能够在 OCI 容器内自动构建 GitHub 代码仓库的平台,其具体的构建机制取决于一系列启发式规则。在实践中,我们依据一系列启发式规则来构建容器,构建任务被调度到规模可任意设定的节点池上。我们还使用了一套稳健的可观测性系统来上报构建失败,从而使我们能够进一步优化、调试和改进我们的构建启发式规则。就规模而言,我们的代码执行环境包含约一百万个代码仓库(及其许多历史版本),涵盖多种不同的编程语言,而且我们还在不断地把新的仓库添加并构建进这个环境。我们要强调,构建容器的过程仍在持续演进,而最近我们从基于智能体(agent-based)的构建方式中收获了巨大的成功。这对我们来说是尤为激动人心的一步,因为一个用工厂构建出来的智能体,如今正在努力改进那座构建了它自己的工厂。更宽泛地说,我们的代码执行平台同样是模型评估中的一个核心组件,为模型在软件开发任务上的改进提供了强有力的反馈信号。

Finally, all of these things fit together. As part of our Kubernetes orchestrator, we build each individual piece of the Model Factory as a Docker container, which we then deploy on various node pools. The exact details of the deployment depend entirely on the cluster’s current workload and the workflow itself, but our orchestrator allows us to preempt low-priority jobs in favor of more pressing jobs. Moreover, the orchestrator itself also supports backfill, enabling quick jobs to be finished quickly.

最后,所有这些部分被拼合在了一起。作为 Kubernetes 编排器的一部分,我们把 Model Factory 的每一个独立组件都构建为一个 Docker 容器,随后将其部署到各种节点池上。部署的具体细节完全取决于集群当前的工作负载以及工作流本身,但我们的编排器允许我们抢占低优先级的任务,以让位于更紧迫的任务。此外,编排器本身还支持回填(backfill),使得耗时短的任务能够被快速完成。

Thank you|致谢

We hope you enjoyed this tour of our Model Factory. Check out our next installment for a deep dive into our data pipelines, our workflows, and a discussion of how we process and refine massive datasets for use in LLM training.

我们希望你喜欢这次对 Model Factory 的巡览。请继续关注我们的下一篇文章,届时我们将深入剖析我们的数据流水线、我们的工作流,并探讨我们如何处理和精炼海量数据集以用于 LLM 训练。