如何从S3路径加载模型至SageMaker进行Hugging Face微调？

我需要对通过SageMaker训练作业生成并存储在S3的模型进行微调，但Hugging Face默认会从其官方仓库查找模型，导致报错。目前已经能在SageMaker Notebook中加载该模型进行推理和训练，但SageMaker训练作业会创建独立容器，因此需要通过train.py脚本在容器内处理模型加载逻辑来解决问题。当前用于启动训练作业的代码如下：

from sagemaker.huggingface import HuggingFace
batch_size = 16

metric_definitions = [
    {'Name': 'eval_loss', 'Regex': "'eval_loss': ([0-9]+(.|e\-)[0-9]+),?"},
    {'Name': 'loss', 'Regex': "'loss': ([0-9]+(.|e\-)[0-9]+),?"},

]

# 传入训练作业的超参数
hyperparameters={'epochs': 20,
                 'train_batch_size': batch_size,
                 'model_name': "model",
                 }

# 创建HuggingFace Estimator
huggingface_estimator = HuggingFace(
    entry_point='train.py',
    source_dir='./code',
    instance_type='ml.g5.xlarge',
    instance_count=1,
    role=role,
    transformers_version='4.26',
    pytorch_version='1.13',
    py_version='py39',
    metric_definitions=metric_definitions,
    hyperparameters = hyperparameters
)

# 启动训练作业，传入训练/测试数据集以及预训练模型的S3路径
huggingface_estimator.fit({'train': training_input_path, 'test': val_input_path,
                          "input_model": "s3://sagemaker-us-east-1-18**********/huggingface-pytorch-training-1/output/model.tar.gz"})

解决方案 #

核心逻辑

SageMaker训练容器会自动将fit()方法中指定的输入通道挂载到容器内的/opt/ml/input/data/目录，并且会自动解压model.tar.gz格式的模型包到对应通道路径下。无需手动编写S3下载逻辑，直接从本地挂载路径加载模型即可跳过Hugging Face Hub的查找流程。

补全train.py核心代码

在train.py中，通过SageMaker提供的环境变量获取模型路径，替换默认的Hub模型加载逻辑：

import argparse
import os
from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments
from datasets import load_from_disk

def main():
    parser = argparse.ArgumentParser()
    # 解析传入的超参数
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--train_batch_size", type=int, default=16)
    
    # 获取SageMaker容器内的输入路径环境变量
    parser.add_argument("--train_dir", type=str, default=os.environ.get("SM_CHANNEL_TRAIN"))
    parser.add_argument("--test_dir", type=str, default=os.environ.get("SM_CHANNEL_TEST"))
    parser.add_argument("--input_model_dir", type=str, default=os.environ.get("SM_CHANNEL_INPUT_MODEL"))

    args = parser.parse_args()

    # 从本地挂载路径加载模型和tokenizer，跳过Hub查找
    model = AutoModelForSequenceClassification.from_pretrained(args.input_model_dir)
    tokenizer = AutoTokenizer.from_pretrained(args.input_model_dir)

    # 加载预处理好的数据集（假设已存储为datasets格式）
    train_dataset = load_from_disk(args.train_dir)
    test_dataset = load_from_disk(args.test_dir)

    # 配置训练参数，指定模型输出路径为SageMaker要求的/opt/ml/model
    training_args = TrainingArguments(
        output_dir="/opt/ml/model",
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.train_batch_size,
        logging_dir="/opt/ml/output/logs",
        evaluation_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True
    )

    # 初始化Trainer并启动微调
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=test_dataset
    )
    trainer.train()

    # 确保模型和tokenizer保存到指定路径（Trainer会自动处理，此处为冗余保障）
    model.save_pretrained("/opt/ml/model")
    tokenizer.save_pretrained("/opt/ml/model")

if __name__ == "__main__":
    main()

关键细节

• 环境变量获取路径：SageMaker会自动为每个输入通道设置SM_CHANNEL_<通道名称>环境变量，比如input_model通道对应SM_CHANNEL_INPUT_MODEL，直接通过os.environ.get()读取即可拿到解压后的模型路径。
• 模型加载方式：使用from_pretrained()传入本地路径而非模型名称，避免触发Hugging Face Hub的模型查找逻辑，解决报错问题。
• 模型输出要求：必须将微调后的模型保存到/opt/ml/model目录，训练作业结束后SageMaker会自动将该目录打包为model.tar.gz并上传到S3的默认输出路径。

内容的提问来源于stack exchange，提问作者Nikhil Belure