在企業組織努力更大限度地提高其 生成式 AI 投資的價值時,獲取最新的模型開發成果對于持續取得成功至關重要。通過在首發日使用先進的模型,團隊可以高效利用這些創新、保持相關性并提高競爭力。
在過去的一年里,開源社區發布了一系列激動人心的模型系列,包括 Meta Llama、Google Gemma、Mistral Codestral、Codestral Mamba、Large 2、Mixtral、Qwen 3、2 和 2.5、Deepseek R1、NVIDIA Nemotron 和 NVIDIA Llama Nemotron。這些模型通常在 Hugging Face Hub 上提供,以便更廣泛的社區輕松訪問。
發布后不久,許多用戶專注于評估模型功能和探索潛在應用。為了了解模型的潛力并發現創新機會,針對特定用例進行微調通常成為一個關鍵的優先事項。
NVIDIA NeMo 框架使用 NVIDIA Megatron-Core 和 Transformer-Engine (TE) 后端,在數千個 NVIDIA GPU 上實現高吞吐量和 Model Flops Utilization (MFU),從而實現卓越性能。但是,將新模型架構集成到 NeMo 框架中需要使用 Megatron-Core 基元進行多階段模型轉換,然后驗證不同階段,包括監督式和參數高效型微調、模型評估以及 Hugging Face 到 NeMo 轉換。這會在模型發布和優化訓練/訓練后配方開發之間造成時間延遲。
為確保最新模型獲得首發日支持,NeMo 框架引入了 Automatic Model (AutoModel) 功能。
在 NVIDIA NeMo 框架中引入 AutoModel
AutoModel 是一個高級接口,旨在簡化對預訓練模型的支持。作為 NeMo 框架的一部分,它使用戶能夠無縫微調任何 Hugging Face 模型,以進行快速實驗。AutoModel 目前涵蓋文本生成和視覺語言模型類別,并計劃擴展到更多類別,如視頻生成。

AutoModel 功能提供開箱即用的支持:
- 支持擴展的模型并行 – 目前通過 Fully-Sharded Data Parallelism 2 (FSDP2) 和 Distributed Data Parallel (DDP) 實現,Tensor Parallelism (TP) 和 Context Parallelism (CP) 即將推出。
- 通過 JIT 編譯增強 PyTorch 性能。
- 無縫過渡到由 Megatron-Core 提供支持的最新優化訓練和后訓練方法,隨著它們變得可用。
- 導出到 vLLM 以優化推理,NVIDIA TensorRT-LLM 導出即將推出。
通過使用 Hugging Face 生態系統,AutoModel 可以輕松集成大量 LLM,而無需顯式檢查點重寫。所有模型均受原生支持,其中最熱門的一個子集也獲得優化的 Megatron-Core 支持。

Megatron-Core 后端 | AutoModel 后端 | |
---|---|---|
覆蓋范圍 | 最熱門的 LLM,包含經過專家調優的方法 | 首發日 Hugging Face Text 支持的所有模型 |
訓練 Throughput 性能 | 使用 Megatron-Core 內核優化吞吐量 | 借助 Liger 內核、cut cross entropy 和 PyTorch JIT 實現出色性能 |
可擴展性 | 多達 1,000 個 GPU,具有完整的 4-D 并行性 (TP, PP, CP, EP) | 使用 PyTorch 原生 TP、CP 和 FSDP2 實現可比可擴展性,訓練吞吐量略有降低 |
推理路徑 | 導出到 TensorRT-LLM、vLLM 或直接導出到 NVIDIA NIM | 導出到 vLLM |
表 1。NeMo 框架中兩個后端 Megatron-Core 和 AutoModel 的比較
如何使用 AutoModel
要在 NeMo 框架中使用 AutoModel 加載并運行 LoRA 和 Supervised Finetuning (SFT) ,請執行以下高級步驟:
- 實例化 Hugging Face 模型 :使用
llm.HFAutoModelForCausalLM
加載任何 Hugging Face 模型,并指定model_id
參數以選擇所需的模型。 - 添加適配器: 利用
llm.peft.LoRA
為模型添加適配器。- 指定 LoRA 目標模塊: 使用 target_modules 識別要適應的模塊,并通過 Fully Qualified Names (FQN) 上的正則表達式進行靈活匹配。
- 配置
None
,使用 SFT 調整所有參數
- 準備數據:通過 llm.HFDatasetDataModule 利用 Hugging Face 的數據集。
- 配置并行:使用 DDP 和 FSDP2 指定模型并行和分片策略,以跨多個節點進行擴展。
請參閱以下偽代碼。您可以在 NeMo 框架 GitHub 中找到完整的參考示例。
from datasets import load_dataset dataset = load_dataset( "rajpurkar/squad" , split = "train" ) dataset = dataset. map (formatting_prompts_func) llm.api.finetune( # Model & PEFT scheme model = llm.HFAutoModelForCausalLM(model_id), # Setting peft=None will run full parameter SFT peft = llm.peft.LoRA( target_modules = [ '*_proj' , ‘linear_qkv’], # Regex-based selector dim = 32 , ), # Data data = llm.HFDatasetDataModule(dataset), # Optimizer optim = fdl.build(llm.adam.pytorch_adam_with_flat_lr(lr = 1e - 5 )), # Trainer trainer = nl.Trainer( devices = args.devices, max_steps = args.max_steps, strategy = args.strategy, # choices= [None, ‘ddp’, FSDP2Strategy] ... ), ) |
切換到 Megatron-Core 支持的路徑以實現最大吞吐量非常簡單,只需更改最少的代碼,并通過一致的 API 實現。
- 模型類
- 而非 AutoModel:
model=llm.HFAutoModelForCausalLM(model_id)
- 使用 Megatron-Core:
model=llm.LlamaModel(Llama32Config1B())
- 優化器模塊
- 而非 AutoModel:
optim=fdl.build(llm.adam.pytorch_adam_with_flat_lr(lr=1e-5))
- 使用 Megatron-Core :
optim=MegatronOptimizerModule(config=opt_config, ...)
- 訓練器策略
- 而非 AutoModel:
strategy=args.strategy,#choices= [None, ‘ddp’, ‘fsdp2’]
- 使用 Megatron-Core :
strategy=nl.MegatronStrategy(ddp="pytorch", …)
這可在訓練和后訓練中實現出色性能,并將 overhead 降至最低。
在 NeMo 中添加新的 AutoModel 類
目前,NeMo AutoModel 支持 AutoModelForCausalLM
類進行文本生成。
如果您想添加對其他任務 (例如 Sequence2SequenceLM
) 的支持,請創建類似于 HFAutoModelForCausalLM 的子類,為您的特定用例調整初始化器、模型配置、訓練/ 驗證步驟以及保存/ 加載方法。此外,實施適當的檢查點處理,并為數據集創建具有自定義批量預處理功能的新數據模塊。
有關更全面的步驟,請參閱 NeMo 框架文檔 。通過遵循它們并使用現有類作為參考,您可以快速擴展 NeMo AutoModel 以支持新任務和模型!
總結
NeMo 框架中的 AutoModel 功能支持快速試驗高性能實現,無需模型轉換即可原生支持 Hugging Face 模型。此外,它還提供了無縫的高性能 Megatron-core 路徑“可選”,讓用戶只需更改最少的代碼即可輕松切換到優化訓練。
AutoModel 隨 NeMo 框架 25.02 版本推出 。要開始使用,請參閱 AutoModel 教程 Notebook,了解 PEFT LoRA 、 SFT 和 多節點擴展 。我們還邀請開發者社區 分享反饋 , 貢獻代碼 ,幫助塑造 AutoModel 的未來發展。
?