引言：简化机器学习实验的新工具------JetTrain
训练模型是机器学习的核心任务。训练模型涉及调整模型内部参数以最小化预测或决策中的错误。一个实验需要多次训练运行，用于验证关于模型改进的某些假设。
然而，随着大型语言模型（LLM）时代的到来，训练或微调现代机器学习模型需要大量硬件资源。因此，机器学习工程师编写代码和启动实验时需要考虑在多种不同的硬件设施上切换，导致机器学习实验工具变得过于复杂。
为了解决这一问题，JetBrains提出了一种新的集成开发环境（IDE）集成方法来做机器学习实验，称为JetTrain。
论文信息概览
1. 论文标题： JetTrain: IDE-Native Machine Learning Experiments
2. 作者与机构： Artem Trofimov 等，JetBrains
3. 论文链接： arxiv.org/pdf/2402.10...
现有机器学习实验接口的局限性
在机器学习的实验流程中，训练模型是核心环节之一。然而，随着大型语言模型（LLM）时代的到来，训练或微调现代机器学习模型需要复杂多样的硬件资源，这导致了机器学习实验工具的复杂化和频繁的上下文切换。本章节将探讨现有机器学习实验工具的局限性，包括SSH连接和虚拟机、Jupyter Notebooks、Pipeline工具以及任务调度工具。
1. SSH连接和虚拟机的缺点
使用SSH连接到租用的虚拟机（VMs）或本地服务器是启动实验的最直接方法。使用者需要安装所有必要的库，下载数据，并运行本地准备的代码。尽管可以在集成开发环境（IDE）中使用远程开发功能，但这种方法的可扩展性有限。在处理大规模实验时，这种方法的局限性为无法高效的利用硬件。
2. Jupyter Notebooks的不足之处
Jupyter Notebooks可以完全控制代码的执行，其高度可定制并且可以与IDE集成。Notebooks还具有丰富的可视化功能。然而，Notebooks中的代码通常难以测试和部署。此外，Notebooks运行代码的速度通常慢于将相同代码放入py文件中运行。
3. Pipeline工具的复杂性
Pipeline工具，如KF Pipelines、Ray、Metaflow等，主要便于流程的生产化，但缺乏调试功能。这些工具要求使用者调整代码以适应特定框架，提高了入门门槛。此类工具的问题在于，其导致实验的编写过于复杂，虽然其能方便地进行部署，但并不是每次实验所训练出的模型都可用。
4. 任务调度工具的限制
任务调度工具，如SkyPilot、DStack、MosaicML等。使用者不需要更改现有代码；他们可以编写一个YAML配置，并通过CLI界面运行实验。YAML与CLI界面的结合缺乏远程调试、终端和其他IDE提供的功能。
| 工具类别 | 描述 | 主要不足之处 |
|---|---|---|
| SSH连接和虚拟机 | 使用SSH连接到租用的虚拟机或本地服务器，需要安装必要的库，下载数据并运行代码。 | 无法高效的利用硬件，可扩展性有限。 |
| Jupyter Notebooks | 允许完全控制代码执行，高度可定制，可与IDE集成，具有丰富的可视化功能。 | 代码测试和部署困难，运行速度慢于直接在.py文件中运行相同代码。 |
| Pipeline工具 | 如KF Pipelines、Ray、Metaflow等，主要用于流程生产化，但缺乏调试功能。 | 编写实验过于复杂，提高了入门门槛，不是每次实验的模型都可用。 |
| 任务调度工具 | 如SkyPilot、DStack、MosaicML等，使用者可以编写YAML配置，并通过CLI界面运行实验。 | 缺乏远程调试功能，与终端和其他IDE的集成有限。 |
JetTrain的创新之处
1. JetTrain的核心理念
我们先总结一下JetTrain的核心：JetTrain是一种创新的集成开发环境（IDE）集成方法，用于启动机器学习（ML）实验。它的核心是降低熟悉IDE用户的入门门槛，并减少上下文切换的负面影响。JetTrain旨在填补简单接口（如SSH和Jupyter）与更复杂的替代方案（如Pipeline和任务调度工具）之间的空白。JetTrain的目标是在保持简单用户体验的同时，提供成熟的调度工具。
2. JetTrain如何简化实验流程
JetTrain通过将实验的启动和管理直接集成到IDE中，简化了实验流程。用户可以在IDE中完成以下操作：编写代码、本地运行和调试代码、编写启动实验的命令、选择能够运行实验的硬件设置、指示需要远程挂载的数据，并在远程硬件上启动实验。
为了确保高效的硬件利用率，JetTrain需要与云调度器集成。调度器需要支持自定义协议连接，以维护调试和终端协议，并且能够从用户定义的存储（例如S3）挂载数据，以便用户可以为训练带来自己的数据。JetTrain选择使用TeamCity作为云调度器，因为它具有各种云连接器，并且针对长时间运行的任务进行了优化。
3. JetTrain的IDE集成特点
JetTrain的IDE集成特点体现在它如何将本地开发与远程实验无缝转换。所有开发工作都保持在本地，而计算资源仅在需求时为实验分配。即使在本地IDE运行配置中，我们几乎拥有迁移到远程机器所需的所有信息：要同步的工作目录、要安装的环境以及要运行的命令。我们只需添加一些额外的属性，例如要挂载的外部数据和硬件配置参数（GPU类型、GPU数量等）。
JetTrain的IDE集成特点还包括对异步调试的支持，这是在远程机器可能在用户启动调试会话后的某个时间点才被分配时所必需的。为了支持这种情况，JetTrain必须缓冲调试协议消息，并重新连接到现有的远程调试会话。这种集成使得用户可以在熟悉的IDE环境中享受到本地开发的体验，同时又能够利用远程硬件资源进行大规模的机器学习实验。
技术挑战与解决方案
在JetTrain项目的实施过程中，我们遇到了几个主要的技术挑战，这些挑战对于确保机器学习实验的顺利进行至关重要。以下是我们面临的挑战以及我们提出的解决方案。
1. 代码和数据同步的挑战
在进行机器学习实验时，确保远程机器上运行的代码与本地机器上的代码完全一致是至关重要的。传统的同步方法，如通过SSH连接使用rsync，可能在地理位置较远的情况下受到互联网连接速度的限制。为了解决这个问题，我们提出了使用地理分布式存储，如Amazon S3，来进行数据同步。首先，数据被同步到离本地机器最近的存储桶，然后通过CDN服务传输到目标地理区域。这种方法可以有效地解决远程同步的问题，确保数据在不同地点之间高效传输。
2. 实验可复现性的重要性
机器学习实验的可复现性是确保实验结果可靠性的关键。工程师通常使用特殊的跟踪工具（如Weights and Biases或MLFlow）来持久化实验指标，但同时也需要能够恢复实验的代码。一种方法是将git提交修订版附加到实验中，但这对于参数选择任务来说可能并不方便，并可能导致大量无意义的提交。我们采用了从本地机器同步用户代码和数据的方法，并利用这一行为来增强可复现性：每次运行都会形成工作目录的快照。为了减少从本地机器到远程机器的网络流量，我们可以先从之前的快照中复制数据，然后再上传本地的更改。
总结
JetTrain预期将改进机器学习实验流程，具体包括：
- 低入门门槛和上下文保持：用户无需离开IDE，降低了学习和使用的难度，同时保持了工作上下文的连续性。
- 集成开发基础设施功能：利用IDE已集成的调试器、终端等功能，提升开发体验。
- 高效硬件利用：通过与云调度器（如TeamCity）集成，确保高效的硬件资源利用。
- 代码和数据同步：通过使用分布式存储（如Amazon S3）和CDN服务，解决了代码和数据同步的问题，提高了实验的可重复性。
- 异步调试支持：解决了由于GPU短缺或有限的本地资源导致的远程机器分配延迟问题，支持断点和调试会话的缓冲和重连。