重要提示

您正在查看 NeMo 2.0 文档。此版本对 API 和新库 NeMo Run 进行了重大更改。我们目前正在将 NeMo 1.0 中的所有功能移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档，请参阅 NeMo 24.07 文档。

将预训练从 NeMo 1.0 迁移到 NeMo 2.0#

NeMo 1.0（先前版本）#

在 NeMo 1.0 中，预训练使用 megatron_gpt_config.yaml 进行配置，并使用 megatron_gpt_pretaining.py 启动。

NeMo 2.0（新版本）#

NeMo 2.0 引入了一种 Pythonic 和模块化的方法来配置实验。有关将 NeMo 1.0 YAML 配置迁移到 NeMo 2.0 的详细说明，请参阅本迁移指南中的其他文档

此外，NeMo 2.0 与 NeMo-Run 兼容，后者简化了 NeMo 实验的配置和执行。有关更多信息，请参阅 NeMo-Run 文档。

可以使用 llm.train API 在 NeMo 2.0 中运行预训练，如下所示

import torch
from nemo import lightning as nl
from nemo.collections import llm
from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer
from megatron.core.optimizer import OptimizerConfig

### set up your GPT model config
gpt_config = llm.GPTConfig(
   num_layers=12,
   hidden_size=384,
   ffn_hidden_size=1536,
   num_attention_heads=6,
   seq_length=2048,
   init_method_std=0.023,
   hidden_dropout=0.1,
   attention_dropout=0.1,
   layernorm_epsilon=1e-5,
   make_vocab_size_divisible_by=128,

)

### other docs in this section explain how to configure these
tokenizer = get_nmt_tokenizer("megatron", "GPT2BPETokenizer")
data = llm.PreTrainingDataModule(
   paths={
      "train": [0.75, '/my/traindata1', 0.25, '/my/traindata2'],
      "validation": '/my/validdata1',
      "test": '/my/testdata`',
   },
   global_batch_size=4,
   micro_batch_size=2,
   num_workers=8,
   pin_memory=True,
   seq_length=2048,
   tokenizer=tokenizer
)
model = llm.GPTModel(
   gpt_config,
   tokenizer=data.tokenizer
)
strategy = nl.MegatronStrategy(
   tensor_model_parallel_size=2,
   pipeline_model_parallel_size=2,
   virtual_pipeline_model_parallel_size=None,
   context_parallel_size=1,
   sequence_parallel=True,
   expert_model_parallel_size=1,
)
optim = nl.MegatronOptimizerModule(
   config=OptimizerConfig(
      optimizer="adam",
      lr=0.001,
      use_distributed_optimizer=True
   ),
   lr_scheduler=nl.lr_scheduler.CosineAnnealingScheduler(),
)
trainer = nl.Trainer(
   num_nodes=16,
   devices=8,
   accelerator="gpu",
   strategy=strategy,
   plugins=nl.MegatronMixedPrecision(precision="bf16-mixed"),
   max_epochs=None,
   max_steps=75000,
   max_time="05:00:00:00",
   log_every_n_steps=10,
   val_check_interval=2000,
   limit_val_batches=50,
   limit_test_batches=50,
   accumulate_grad_batches=1,
   gradient_clip_val=1.0,
)
nemo_logger = nl.NeMoLogger(
   log_dir="your/log/ckpt/dir/here"
)
resume = None # None for pretraining from scratch

llm.train(
   model=model,
   data=data,
   trainer=trainer,
   log=nemo_logger,
   tokenizer=tokenizer,
   resume=resume,
   optim=optim,
)

除了上面示例中使用的通用 GPTModel 之外，我们还支持 Gemma、Llama、Mistral 和 Mixtral 模型。对于其他模块，请参阅 NeMo 2.0 指南中的 分词器、数据模块、Megatron 策略、Megatron 优化器模块、训练器 和 NeMo 日志记录器。

命令行界面#

您还可以通过 NeMo CLI（由 NeMo-Run 提供）使用现有配方。

例如

nemo llm pretrain --factory llama3_8b

  • llama3_8b 可以替换为 NeMo 2.0 中包含的其他配方（例如 mixtral_8x7b）

  • 对于序列长度大于 16k 的长上下文配方，由于目前还没有实际的长上下文微调用例，我们只支持 pretrain 任务。

继续训练#

在 NeMo 2.0 中，您可以设置恢复路径以从 NeMo 检查点继续训练。例如

... # same as above
resume = llm.default_resume()
resume.resume_from_path = "path/to/NeMo/checkpoint/you/want/to/resume/from"

llm.train(
   model=model,
   data=data,
   trainer=trainer,
   log=nemo_logger,
   tokenizer=tokenizer,
   resume=resume,
   optim=optim,
)

如果未设置 resume.resume_from_path，则训练将尝试从日志记录器的 log_dir 恢复（如果检查点存在）。

命令行界面#

要使用 现有配方从检查点继续训练，您还可以将 resume.resume_from_path="to/some/path" 选项添加到命令中。用户还可以通过这种方式指定其他选项，例如序列长度。例如

nemo llm pretrain --factory llama3_8b resume.resume_from_path="to/some/path" model.config.seq_length=4096 data.seq_length=4096

迁移步骤#

  1. 使用迁移指南中的其他文档将您的 NeMo 1.0 YAML 迁移到 NeMo 2.0。

  2. 使用 llm.train API 运行预训练。

在本页中