准备阶段下载并启动AutoModel官方镜像dockerpull nvcr.io/nvidia/nemo-automodel:26.04sudodockerrun-it--rm--ipchost--gpusall--networkhost\-v/data/hmlp:/hmlp/data\--namenemo-automodel-test\nvcr.io/nvidia/nemo-automodel:26.04\/bin/bashsudodockerrun-it--ipchost--gpusall--networkhost\-v/data/hmlp:/hmlp/data\-eNCCL_P2P_DISABLE1\-eNCCL_SHM_DISABLE1\-eCUDA_VISIBLE_DEVICES0,1,2,3,4,5,6,7\-w/hmlp/data/finetune/196d63aa-84f8-42f6-9cff-722758276518/f1a28b50-b8e5-4318-8545-5a8b57b67be8\--nametrain-text\nvcr.io/nvidia/nemo-automodel:26.04\/bin/bash模型准备mkdir-p/hmlp/data/finetune/nemo-automodel-sft-testcd/hmlp/data/finetune/nemo-automodel-sft-test# data: 存储数据集; output: 存储微调后的checkpointmkdir-pdata output# 模型使用本地Qwen2.5-0.5B-Instruct/hmlp/data/storage/Qwen2.5-0.5B-Instruct/V1/Qwen2.5-0.5B-Instruct文本格式数据集继续预训练准备数据集可以通过hfd.sh脚本下载预处理文本格式数据集为megatron格式uv run /opt/Automodel/tools/preprocess_megatron_dataset.py\--input/hmlp/data/finetune/nemo-automodel-sft-test/data/c4_demo.jsonl\--json-keys text\--output-prefix domain_corpus\--output-path /hmlp/data/finetune/nemo-automodel-sft-test/data/megatron\--workers8\--pretrained-model-name-or-path /hmlp/data/storage/Qwen2.5-0.5B-Instruct/V1/Qwen2.5-0.5B-Instruct\--append-eod创建数据集缓存索引目录每次启动预训练Megatron都需要为庞大的数据集构建一套随机访问的索引包含文档索引、样本索引和打乱索引这一过程在TB级数据上可能耗时超过30分钟。index_mapping_dir 参数的作用正是 “缓存” 这些复杂的索引文件。避免重复计算一旦指定了目录系统会优先尝试从该路径加载已有的.npy格式索引文件。只有首次运行或数据集发生变化时才会重新构建。加速下次启动这种机制尤其适用于多次启动且训练数据不变的场景能大幅缩短后续的训练启动时间。如下创建index_mapping_dir参数需要使用的缓存文件mkdir-p/hmlp/data/finetune/nemo-automodel-sft-test/data/megatron/mapping_dir准备好训练配置文件-使用文本数据集vimtrain-Qwen-Qwen-2.5-0.5B-Instruct-V1-text.yaml# 文件内容如下recipe: TrainFinetuneRecipeForNextTokenPrediction step_scheduler: global_batch_size:32local_batch_size:1ckpt_every_steps:200val_every_steps:100num_epochs:1max_steps:1000dist_env: backend: nccl timeout_minutes:30rng: _target_: nemo_automodel.components.training.rng.StatefulRNG seed:1111ranked:truemodel: _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained pretrained_model_name_or_path: /hmlp/data/storage/Qwen2.5-0.5B-Instruct/V1/Qwen2.5-0.5B-Instruct torch_dtype: bf16 trust_remote_code:truecheckpoint: enabled:truecheckpoint_dir: /hmlp/data/finetune/nemo-automodel-sft-test/output model_save_format: safetensors save_consolidated:true# 恢复训练时打开下面这一行# restore_from: LATESTdistributed: strategy: fsdp2 dp_size: none tp_size:1cp_size:1pp_size:1sequence_parallel:falseactivation_checkpointing:trueloss_fn: _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy dataset: _target_: nemo_automodel.components.datasets.llm.megatron_dataset.MegatronPretraining paths: /hmlp/data/finetune/aa15502a-09a5-41bf-a37e-42a657dfd4f3/f1a28b50-b8e5-4318-8545-5a8b57b67be8/data/train-megatron/domain_corpus_*_text_document* index_mapping_dir: /hmlp/data/finetune/aa15502a-09a5-41bf-a37e-42a657dfd4f3/f1a28b50-b8e5-4318-8545-5a8b57b67be8/data/train-megatron/mapping_dir tokenizer: _target_: nemo_automodel._transformers.auto_tokenizer.NeMoAutoTokenizer.from_pretrained pretrained_model_name_or_path: /hmlp/data/storage/Qwen2.5-0.5B-Instruct/V1/Qwen2.5-0.5B-Instruct trust_remote_code:trueseq_length:2048split:1, 0, 0splits_to_build:traindataloader: _target_: torchdata.stateful_dataloader.StatefulDataLoader collate_fn: torch.utils.data.default_collate dataloader_type: single validation_dataset: _target_: nemo_automodel.components.datasets.llm.megatron_dataset.MegatronPretraining paths: /hmlp/data/finetune/aa15502a-09a5-41bf-a37e-42a657dfd4f3/f1a28b50-b8e5-4318-8545-5a8b57b67be8/data/eval-megatron/domain_corpus_*_text_document* index_mapping_dir: /hmlp/data/finetune/aa15502a-09a5-41bf-a37e-42a657dfd4f3/f1a28b50-b8e5-4318-8545-5a8b57b67be8/data/eval-megatron/mapping_dir tokenizer: _target_: nemo_automodel._transformers.auto_tokenizer.NeMoAutoTokenizer.from_pretrained pretrained_model_name_or_path: /hmlp/data/storage/Qwen2.5-0.5B-Instruct/V1/Qwen2.5-0.5B-Instruct trust_remote_code:trueseq_length:2048split:0, 1, 0splits_to_build:validationvalidation_dataloader: _target_: torchdata.stateful_dataloader.StatefulDataLoader collate_fn: torch.utils.data.default_collate dataloader_type: single optimizer: _target_: torch.optim.AdamW lr:5.0e-6 betas:[0.9,0.95]eps:1.0e-8 weight_decay:0.1lr_scheduler: lr_decay_style: cosine lr_warmup_steps:100min_lr:1.0e-6执行继续预训练# 推荐命令CUDA_VISIBLE_DEVICES0,1\automodel\--nnodes1\--nproc-per-node2\--node_rank0\--master_addr10.10.0.1\--master_port39500\train-Qwen-Qwen-2.5-0.5B-Instruct-V1-text.yamlCUDA_VISIBLE_DEVICES0,1automodel --nproc-per-node2train-Qwen-Qwen-2.5-0.5B-Instruct-V1-text.yaml# 不推荐只是简单测试CUDA_VISIBLE_DEVICES0,1\uv run torchrun\--nnodes1\--nproc-per-node2\--node_rank0\--master_addr10.10.0.1\--master_port29500\/opt/Automodel/examples/llm_pretrain/pretrain.py\--configtrain-Qwen-Qwen-2.5-0.5B-Instruct-V1-text.yaml# 也可以用这个CUDA_VISIBLE_DEVICES0,1torchrun\--nnodes1\--nproc-per-node2\--node_rank0\--master_addr10.10.0.1\--master_port29500\-mnemo_automodel._cli.app\pretrain llm-ctrain-Qwen-Qwen-2.5-0.5B-Instruct-V1-text.yamlnnodes总节点数node_rank当前节点是第几台节点从0开始master_addr主节点地址master_port主节点端口支持的完整微调配置文件# 配方# 选择哪个配方类来运行训练循环。# 使用短名称自动发现或完整的 Python 路径# recipe: nemo_automodel.recipes.llm.train_ft.TrainFinetuneRecipeForNextTokenPredictionrecipe:TrainFinetuneRecipeForNextTokenPrediction# 训练计划# 控制 epoch 数量、批量大小以及保存检查点/验证的频率。# 没有 _target_ — 这些是配方直接读取的普通值。step_scheduler:grad_acc_steps:4# 每次优化器步骤之前累积的微批次数。有效批量大小 grad_acc_steps × batch_size。ckpt_every_steps:10# 每 N 个梯度步骤保存一个检查点val_every_steps:10# 每 N 个梯度步骤运行一次验证循环num_epochs:1# 对训练数据集进行多少次完整遍历# 进程组# 初始化 PyTorch 分布式进程组。# 没有 _target_ — 由配方直接使用。# 通常不需要调整此项。dist_env:backend:nccl# 通信后端ncclGPU推荐或 glooCPUtimeout_minutes:1# 集合通信操作的超时时间对于初始化时间较长的大模型可增加此值# 分布式策略# 确定模型权重、数据和计算如何在 GPU 之间拆分。# 没有 _target_ — 由配方直接使用。# 详细信息请参阅“高级主题”中的“分布式训练TP、PP、CP 和 EP”。distributed:strategy:fsdp2# 并行策略fsdp2推荐、megatron_fsdp 或 ddp。# FSDP2 在数据并行组中对参数和优化器状态进行分片。dp_size:null# 数据并行组大小。null 从 world_size ÷ (tp_size × cp_size × pp_size) 自动检测。tp_size:1# 张量并行大小在 GPU 之间拆分权重矩阵。# 如果模型无法放在单个 GPU 上则设置为 2、4 或 8。# 应能整除注意力头的数量。cp_size:1# 上下文并行大小在 GPU 之间拆分输入序列。# 对于非常长的上下文例如 32k token请增加此值。sequence_parallel:false# 当为 true 时扩展 TP 以同时沿序列维度分片激活值从而进一步节省内存。# 随机数生成器# _target_ → StatefulRNG一个可保存检查点的 RNG确保训练重启时的序列相同。# seed 和 ranked 是传递给 StatefulRNG() 的关键字参数。rng:_target_:nemo_automodel.components.training.rng.StatefulRNGseed:1111# 用于可重现性的全局随机种子ranked:true# 当为 true 时每个 GPU 等级获得一个从种子派生的唯一 RNG 流# 因此每个 GPU 的数据打乱方式不同# 模型# _target_ → NeMoAutoModelForCausalLM.from_pretrained下载或从缓存加载预训练的 HuggingFace 模型# 并包装以支持 NeMo 分布式训练。接受任何 from_pretrained 关键字参数cache_dir、torch_dtype 等。model:_target_:nemo_automodel.NeMoAutoModelForCausalLM.from_pretrainedpretrained_model_name_or_path:meta-llama/Llama-3.2-1B# PEFT如需全参数 SFT请删除或注释整个部分# _target_ → PeftConfig一个描述哪些层获得 LoRA 适配器的数据类。# 配方将此配置传递给 build_model()后者将适配器附加到匹配的层上。peft:_target_:nemo_automodel.components._peft.lora.PeftConfigtarget_modules:*.proj# 与全限定层名称匹配的 glob 模式# *.proj 匹配每个以 proj 结尾的层dim:8# 低秩维度 r — 控制适配器的容量。# 值越大表达能力越强但使用更多内存。alpha:32# LoRA 缩放因子适配器输出乘以 alpha/dim。# 值越高适配器在训练期间的影响越大。use_triton:True# 使用优化的 Triton 内核进行 LoRA 前向/反向传播# 需要安装 triton 包# 检查点# 没有 _target_ — 由配方直接使用的普通键值组。checkpoint:enabled:true# 设置为 false 以完全跳过保存检查点checkpoint_dir:checkpoints/# 输出目录。Docker 用户请绑定挂载此路径# 例如 -v $(pwd)/checkpoints:/workspace/checkpoints# 以在容器重启后保留检查点。model_save_format:safetensors# safetensors推荐更快更安全或# torch_save传统的基于 pickle 的格式save_consolidated:True# 当为 true 时将一个与 HuggingFace 兼容的单个检查点写入 model/consolidated/# 可直接由 Transformers、vLLM 等加载。需要 safetensors 格式。# 训练数据集# _target_ → make_squad_dataset一个工厂函数用于下载 SQuAD 数据集、进行标记化并返回一个 torch Dataset。# 要使用不同的数据集请将 _target_ 更改为另一个工厂函数参见数据集指南。dataset:_target_:nemo_automodel.components.datasets.llm.squad.make_squad_datasetdataset_name:rajpurkar/squad# HuggingFace Hub 数据集 IDsplit:train# 使用哪个拆分train、validation、test# 验证数据集validation_dataset:_target_:nemo_automodel.components.datasets.llm.squad.make_squad_datasetdataset_name:rajpurkar/squadsplit:validationlimit_dataset_samples:64# 将验证集限制为 64 个样本以加快评估循环# 删除此行以使用完整验证集# 训练 DataLoader# _target_ → StatefulDataLoader来自 torchdata 的可保存检查点的 DataLoader# 在训练重启时保存和恢复迭代状态因此恢复的运行不会重新处理已见过的批次。dataloader:_target_:torchdata.stateful_dataloader.StatefulDataLoadercollate_fn:nemo_automodel.components.datasets.utils.default_collater# 将单个样本填充并批处理为张量的函数可替换为自定义整理函数batch_size:8# 每个 GPU 每个微批次的样本数shuffle:true# 每个 epoch 是否打乱数据集# 验证 DataLoadervalidation_dataloader:_target_:torchdata.stateful_dataloader.StatefulDataLoadercollate_fn:nemo_automodel.components.datasets.utils.default_collaterbatch_size:8# 损失函数# _target_ → MaskedCrossEntropy标准的交叉熵损失自动忽略填充 token使其不影响梯度。# 其他可用的损失函数替换 _target_ 以使用# - nemo_automodel.components.loss.chunked_ce.ChunkedCrossEntropy# 沿序列维度分块计算 CE以降低峰值内存。对于超长序列很有用。接受 chunk_len默认 32。# - nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy# 将最终的线性投影lm_head与 CE 计算融合避免生成完整的 logit 张量。# 对于大词汇表可显著节省**内存**。# - nemo_automodel.components.loss.te_parallel_ce.TEParallelCrossEntropy# 基于 TransformerEngine 的并行 CE使用 Triton 内核。专为 logits 在 TP 等级间分片的张量并行设置设计。loss_fn:_target_:nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy# 优化器# _target_ → torch.optim.Adam此处可使用任何 torch.optim 类例如 AdamW、SGD。# 其余所有键成为构造函数的参数。optimizer:_target_:torch.optim.Adamlr:1.0e-5# 学习率 — 最重要的可调超参数betas:[0.9,0.999]# Adam 动量系数β₁ 用于均值β₂ 用于方差eps:1e-8# 为保证数值稳定性加到分母上的小常数weight_decay:0# L2 正则化强度0 无正则化# 日志记录可选# 取消注释以启用 Weights Biases 实验跟踪。# wandb:# project: your_wandb_project # WB 项目名称# entity: your_wandb_entity # WB 团队或用户名# name: your_wandb_exp_name # 本次运行的显示名称# save_dir: your_wandb_save_dir # WB 工件的本地目录