Meta Lingua：一个专为研究而设计的极简快速的 LLM 训练和推理库-AI应用

Meta Lingua 是一个专为研究而设计的极简快速的 LLM 训练和推理库。Meta Lingua 使用易于修改的 PyTorch 组件来尝试新的架构、损失、数据等。我们的目标是让此代码能够实现端到端的训练、推理和评估，并提供工具来更好地了解速度和稳定性。虽然 Meta Lingua 目前正在开发中，但我们为您提供了多个代码库apps来展示如何使用此代码库。

1 快速启动

以下命令启动 slurm 作业，为 Meta Lingua 创建环境。创建环境大约需要 5 分钟（不包括下载次数）。

git clone https://github.com/facebookresearch/lingua
cd lingua

bash setup/create_env.sh
# or if you have access to a slurm cluster
sbatch setup/create_env.sh

完成后你就可以激活环境

conda activate lingua_<date>

并启动调试作业以检查一切是否正常。 提供的配置是模板，您需要对其进行调整才能使其正常工作（更改 dump_dir、数据根目录等…）

# stool stands for slurm tool !
python -m lingua.stool script=apps.main.train config=apps/main/configs/debug.yaml nodes=1 partition=<partition>
# if you want to launch locally you can use torchrun
torchrun --nproc-per-node 8 -m apps.main.train config=apps/main/configs/debug.yaml
# or you can also launch on 1 GPU
python -m apps.main.train config=apps/main/configs/debug.yaml

2 训练结果

我们在许多下游任务上取得了非常强劲的表现，并与DCLM 基线 1.0的性能相匹配。

2.1 1B models on 60B DCLM tokens

2.2 7B models

3 项目概况

Meta Lingua 的结构如下：

📦meta-lingua
 ┣ 📂lingua # Core library
 ┃ ┣ 📜args.py
 ┃ ┣ 📜checkpoint.py
 ┃ ┣ 📜data.py
 ┃ ┣ 📜distributed.py
 ┃ ┣ 📜float8.py
 ┃ ┣ 📜logger.py
 ┃ ┣ 📜metrics.py
 ┃ ┣ 📜optim.py
 ┃ ┣ 📜probe.py
 ┃ ┣ 📜profiling.py
 ┃ ┣ 📜stool.py
 ┃ ┣ 📜tokenizer.py
 ┃ ┗ 📜transformer.py
 ┣ 📂setup
 ┃ ┣ 📜create_env.sh
 ┃ ┗ 📜download_prepare_hf_data.py
 ┗ 📂apps # Apps that put components together
   ┣ 📂main # Main language modeling app with llama
   ┃ ┣ 📂configs
   ┃ ┣ 📜eval.py
   ┃ ┣ 📜generate.py
   ┃ ┣ 📜train.py
   ┃ ┗ 📜transformer.py
   ┣ 📂fastRNN 
   ┃ ┣ 📂component
   ┃ ┣ 📂hawk
   ┃ ┣ 📂minGRU
   ┃ ┣ 📂minLSTM
   ┣ 📂mamba
   ┣ 📂mtp # Multi token prediction
   ┗ 📂plots

文件lingua夹包含一些基本且可重复使用的组件，而文件apps夹包含将这些组件组合在一起的脚本。例如，主训练循环位于 中apps/main。我们强烈建议您将其用作模板，并根据您的需要对其进行修改以适合您的实验。

在 Meta Lingua 中，没有什么是神圣的。我们特意尝试使其尽可能易于修改！因此，请随意扩展并修改任何内容。

以下是最重要的文件和功能的简要说明：

• transformer.py：定义模型架构。这是纯 PyTorch nn.Module！这里没什么特别的。
• distribution.py：处理在多个 GPU 上分发模型。这是通过 parallelize_module 函数完成的，该函数包装您的 vanilla nn.Module 并应用几乎任何数据并行、完全分片数据并行、模型并行、torch.compile、激活检查点和 float8 的组合。
• data.py：用于LLM预训练的数据加载器。

• profiling.py：xformers分析器的小型包装器，提供自动 MFU 和 HFU 计算，并将分析跟踪转储到转储目录中的分析文件夹中。它还具有内存分析跟踪。
• checkpoint.py：管理模型检查点。它将模型以 .distcp 格式保存在转储目录中的 checkpoints 文件夹中，这是新的 PyTorch 分布式保存方法。此格式允许使用不同数量的 GPU 和不同的分片重新加载模型。您还可以将它们转换为正常的 PyTorch 检查点torch.distributed.checkpoint.format_utils.dcp_to_torch_save，反之亦然torch_save_to_dcp。
• args.py：用于配置的实用程序。

3 配置

大多数组件都需要配置，我们选择使用数据类来表示这些配置对象。args.py 有助于将 config.yaml 和 config 字典转换为相应的数据类。

因此例如 apps/main/train.py 中的 TrainArgs 具有 LMTransformerArgs、OptimArgs 等作为子项。

以下是将转换为 TrainArgs 的示例配置文件：

# This is where Meta Lingua will store anything related to the experiment. 
dump_dir: /path/to/dumpdir
name: "debug"
steps: 1000

seed: 12

optim:
    lr: 3e-4
    warmup: 2000
    lr_min_ratio: 0.000001
    clip: 10.0

distributed:
    fsdp_type: full_shard
    compile: true
    selective_activation_checkpointing: false

model:
    dim: 1024
    n_layers: 8
    n_heads: 8

data:
    root_dir: data/shuffled
    sources:
      wikipedia: 80.0
      arxiv: 20.0
    batch_size: 32
    seq_len: 1024
    load_async: true
    tokenizer:
        name: sp
        path: tokenizers/llama2.model

4 启动作业

4.1 命令行参数

所有脚本（train.py、eval.py、stool.py）中的命令行界面都使用OmegaConf ，它接受点列表形式的参数，因此，如果数据类看起来像

@dataclass
class DummyArgs:
    name: str = "blipbloup"
    mode: LMTransformerArgs = LMTransformerArgs()
    
@dataclass
class LMTransformerArgs:
    dim: int = 512
    n_layers: int = 12

然后您可以传递 model.dim=32 来更改 LMTransformerArgs 中的值，或者仅传递 name=tictac 来更改顶级属性。

train.py只是将配置文件的路径作为参数并加载该配置。此处的行为如下：

1. 我们用其默认值实例化 TrainArgs
2. 我们使用提供的配置文件中的值覆盖这些默认值
3. 我们使用命令行提供的附加参数覆盖结果

如果我们采用上面的 DummyArgs 示例，使用train.py config=debug.yaml model.dim=64 name=tictac debug.yaml 包含的内容调用 train.py

model:
    n_layers: 24

将使用配置启动训练

DummyArgs(name="tictac", LMTransformerArgs(dim=64, n_layers=24))

4.2 使用 slurm 启动

由于我们要进行分布式训练，因此我们需要train.py运行 N 次（N 为 GPU 的数量）

最简单的方法是通过 SLURM。为了简化操作，我们提供了lingua/stool.py一个简单的 Python 脚本，

1. 将提供的配置保存到 dump_dir
2. 将当前代码复制到 dump_dir 以便备份
3. 创建一个 sbatch 文件submit.slurm，然后使用提供的配置启动作业。

它可以通过命令行使用

python -m lingua.stool config=apps/main/configs/debug.yaml nodes=1 account=fair_amaia_cw_codegen qos=lowest

或者launch_job直接使用该函数。例如，这允许您在 jupyter 笔记本中创建许多任意配置（扫描参数、执行消融）并直接从那里启动作业。

由于配置文件被复制到 dump_dir，因此一个简单的迭代方法就是简单地更改配置文件并启动上面相同的命令。

5 调试

为了快速迭代，最好不要每次都等待 slurm 分配。您可以要求 slurm 为您分配资源，然后一旦分配完成，您就可以在同一分配上运行多个命令。

例如你可以这样做：

salloc --nodes 2 --cpus-per-gpu 16 --mem 1760GB --gres=gpu:8 --exclusive --time=72:00:00

这将使您能够访问当前终端中的 2 个节点。分配完成后，您将看到一些自动添加的 slurm 环境变量，例如 $SLURM_JOB_ID 等…这允许您在同一终端中执行例如

srun -n 16 python -m apps.main.train config=apps/main/configs/debug.yaml

这将在 16 个 GPU 中的每一个上运行python -m apps.main.train config=apps/main/configs/debug.yaml命令。如果崩溃或结束，您可以重新启动 srun，因为节点已经分配给您，您不必等待 slurm 再次为您提供资源。

这还将在同一个终端中显示所有这些命令的输出，这可能会变得麻烦。

相反，您可以直接使用 stool 配置日志，以便将日志分成每个 GPU 的不同文件。

python -m lingua.stool config=apps/main/configs/debug.yaml nodes=2 launcher=bash dirs_exists_ok=true

请注意，我们添加了launcher=bash，这基本上意味着生成的 submit.slurm 将直接执行，而不是通过 sbatch 提交。submit.slurm 也有一个 srun 命令，因此这与上面的 srun 命令非常相似。我们还添加了dirs_exists_ok=true来告诉 stool 可以覆盖现有文件夹中的内容（代码、配置等）

如果您想使用 pdb 逐步执行代码，则应使用 -n 1 仅在 1 个 GPU 上运行。

6 评估

评估可以在训练期间定期运行，或者您可以直接在给定的检查点启动评估，如下所示：

srun -n 8 python -u -m apps.main.eval config=apps/main/configs/eval.yaml

您需要在该配置中指定评估的检查点和转储目录

或者通过粪便

python -m lingua.stool script=apps.main.eval config=apps/main/configs/eval.yaml nodes=1 account=fair_amaia_cw_codegen qos=lowest

7 转储目录结构

📂example_dump_dir
 ┣ 📂checkpoints
 ┃ ┣ 📂0000001000
 ┃ ┣ 📂0000002000
 ┃ ┣ 📂0000003000
 ┃ ┣ 📂0000004000
 ┃ ┣ 📂0000005000
 ┃ ┣ 📂0000006000
 ┃ ┣ 📂0000007000 # Checkpoint and train state saved every 1000 steps here
 ┃ ┃ ┣ 📜.metadata
 ┃ ┃ ┣ 📜__0_0.distcp
 ┃ ┃ ┣ 📜__1_0.distcp
 ┃ ┃ ┣ 📜params.json
 ┃ ┃ ┣ 📜train_state_00000.json
 ┃ ┃ ┗ 📜train_state_00001.json
 ┣ 📂code # Backup of the code at the moment the job was launched
 ┣ 📂logs
 ┃ ┗ 📂166172 # Logs for each GPU in this slurm job.
 ┃ ┃ ┣ 📜166172.stderr
 ┃ ┃ ┣ 📜166172.stdout
 ┃ ┃ ┣ 📜166172_0.err
 ┃ ┃ ┣ 📜166172_0.out
 ┃ ┃ ┣ 📜166172_1.err
 ┃ ┃ ┗ 📜166172_1.out
 ┣ 📂profiling
 ┃ ┣ 📂memory_trace_plot # Trace of memory usage through time for all GPUs
 ┃ ┃ ┣ 📜000102_h100-192-145_451082.html
 ┃ ┃ ┣ 📜000102_h100-192-145_451083.html
 ┃ ┗ 📂profile_CPU_CUDA_000104 # Profiling traces for all GPUs
 ┃ ┃ ┣ 📜h100-192-145_451082.1720183858874741723.pt.trace.json.gz
 ┃ ┃ ┗ 📜h100-192-145_451083.1720183858865656716.pt.trace.json.gz
 ┣ 📜base_config.yaml
 ┣ 📜config.yaml
 ┣ 📜metrics.jsonl
 ┗ 📜submit.slurm

8 相关存储库

这里我们重点介绍一些与此相关的补充工作。最重要的是torchtitan和torchtune。

Lingua 专为想要尝试 LLM 预训练新想法并获得训练/推理速度和下游基准的快速反馈的研究人员而设计。我们的目标是通过提供轻量级且专注的代码库来降低 LLM 研究的进入门槛。

我们认为 torchtitan、torchtune 和 lingua 是互补工具。Torchtitan 非常适合大规模工作，因为它具有 3D 并行性，并且由于与 PyTorch 团队的密切联系，很可能更快地集成最新的 PyTorch 分布式训练功能。另一方面，Torchtune 擅长微调，尤其是在 GPU 资源有限的情况下，它提供各种微调策略，如 LoRA、QLoRA、DPO 和 PPO。

典型的工作流程可能如下所示：您可能首先在 Lingua 中测试新想法，然后使用 Torchtitan 进一步扩展它，最后使用 Torchtune 进行指导或偏好微调。

虽然这些代码库之间肯定存在一些重叠，但我们认为针对 LLM 工作的不同方面有专门的工具很有价值。例如，Torchtitan 旨在以简洁、极简的代码库展示 PyTorch 的最新分布式训练功能，但对于大多数研究来说，您实际上并不需要 PyTorch 提供的所有功能或在 4096 个 GPU 上扩展到 100B 参数的能力。例如，我们认为 FSDP + torch compile 将满足研究人员 90% 的需求。对于 lingua，我们试图问“需要哪些最小功能集才能得出关于想法 X 可扩展性的可靠结论？”

我们相信这种有针对性的方法可以帮助研究人员更快地取得进展，而无需花费脑力去使用许多可能不需要的技术。