在当今人工智能蓬勃发展的时代，语言模型的性能优化和定制化成为研究与应用的关键方向。本文聚焦于 AWS SageMaker 平台上对 DeepSeek-R1-Distilled-Llama-8B 模型的精调实践，详细探讨这一过程中的技术细节、操作步骤以及实践价值。
一、实验背景与目标
（一）语言模型发展趋势
随着自然语言处理技术的不断演进，语言模型从基础架构的创新逐渐转向在特定场景下的性能优化和功能拓展。开发者们期望通过对预训练模型的精细调整，使其在特定领域、任务中表现更为出色，以满足多样化的应用需求。
（二）DeepSeek-R1-Distilled-Llama-8B 模型优势
DeepSeek-R1-Distilled-Llama-8B 模型融合了 DeepSeek-R1 的先进知识蒸馏技术与 Llama-8B 模型的架构优势。知识蒸馏使得小模型能够学习到大型模型的优秀特征，在保持较小参数规模的同时，具备强大的语言理解和生成能力。8B 的参数规模在性能与计算资源消耗之间达到了较好的平衡，适合在多种场景下进行部署和应用。
（三）实验目标
本次实验旨在利用 AWS SageMaker 平台的强大计算资源和丰富工具，对 DeepSeek-R1-Distilled-Llama-8B 模型进行精细调整，提升其在特定任务（如问答、文本生成等）上的性能表现。通过实验，探索如何优化模型训练过程，提高模型的泛化能力和准确性，为实际应用提供有力支持。
二、AWS SageMaker 平台概述
（一）SageMaker 核心功能
AWS SageMaker 是一个全托管的机器学习平台，为开发者提供了从数据预处理、模型训练、模型部署到模型监控的一站式解决方案。它集成了多种主流的机器学习框架，如 TensorFlow、PyTorch 等，支持在不同规模的计算实例上进行训练和推理。在本次实验中，我们将借助 SageMaker 的分布式训练功能，加速 DeepSeek-R1-Distilled-Llama-8B 模型的精调过程。
（二）SageMaker 在模型开发中的优势
- 高效资源管理：SageMaker 能够自动配置和管理计算资源，根据模型训练的需求动态调整实例类型和数量。这不仅提高了资源利用率，还降低了开发者的运维成本，使他们能够专注于模型开发和优化。
- 分布式训练支持：对于大规模的语言模型训练，分布式训练是关键。SageMaker 支持数据并行和模型并行等多种分布式训练策略，能够充分利用多个计算节点的计算能力，显著缩短训练时间。在精调 DeepSeek-R1-Distilled-Llama-8B 模型时，分布式训练可以加快模型收敛速度，提高训练效率。
- 内置算法与工具：SageMaker 提供了一系列内置的算法和工具，如超参数调优、模型评估等。这些工具简化了模型开发流程，帮助开发者快速找到最优的模型配置，提升模型性能。
三、实验步骤详解
在这个演示中，我们展示了如何使用 PyTorch FSDP、QLoRA、Hugging Face PEFT 和 bitsandbytes 对 DeepSeek-R1-Distill-Llama-8B 模型进行微调。
（一）我们正在使用 SageMaker 远程装饰器在 Amazon SageMaker 训练作业上运行微调任务
JupyterLab 实例类型：ml.t3.medium
Python 版本：3.11
微调：
- 实例类型：ml.g5.12xlarge
- 安装所需的库，包括 Hugging Face 库，并重启内核。
%pip install -r requirements.txt --upgrade %pip install -q -U python-dotenv
（二）设置配置文件路径
我们正在设置 config.yaml 文件所在的目录，以便远程装饰器可以通过 SageMaker Defaults 使用这些设置。
本次实验使用 Hugging Face 容器用于 us-east-1 区域。确保您使用的是适合您 AWS 区域的正确镜像，否则请编辑 config.yaml。容器镜像可在此处获取。
from dotenv import load_dotenv
import os
# Use .env in case of hidden variables
load_dotenv()
# Set path to config file
os.environ["SAGEMAKER_USER_CONFIG_OVERRIDE"] = os.getcwd()（三）可视化并上传数据集
我们将加载 rajpurkar/squad 数据集
import sagemaker
sagemaker_session = sagemaker.Session()
bucket_name = sagemaker_session.default_bucket()
default_prefix = sagemaker_session.default_bucket_prefix
from datasets import load_dataset
import pandas as pd
dataset = load_dataset("rajpurkar/squad")
df = pd.DataFrame(dataset['train'])
df = df.iloc[0:1000]
df['answer'] = [answer['text'][0] for answer in df['answers']]
df = df[['context', 'question', 'answer']]
df.head()
from sklearn.model_selection import train_test_split
train, test = train_test_split(df, test_size=0.1, random_state=42)
print("Number of train elements: ", len(train))
print("Number of test elements: ", len(test))创建一个提示模板，并加载一个随机样本的数据集以尝试摘要。
from random import randint
# custom instruct prompt start
prompt_template = f"""<|begin_of_text|><|start_header_id|>user<|end_header_id|>\nContext:\n{
{context}}\n\n{
{question}}<|eot_id|><|start_header_id|>assistant<|end_header_id|>{
{answer}}<|end_of_text|><|eot_id|>"""
# template dataset to add prompt to each sample
def template_dataset(sample):
sample["text"] = prompt_template.format(context=sample["context"],
question=sample["question"],
answer=sample["answer"])
return sample使用 Hugging Face Trainer 类对模型进行微调。定义我们想要使用的超参数。我们还创建一个 DataCollator，它将负责填充我们的输入和标签。
from datasets import Dataset, DatasetDict
train_dataset = Dataset.from_pandas(train)
test_dataset = Dataset.from_pandas(test)
dataset = DatasetDict({"train": train_dataset, "test": test_dataset})
train_dataset = dataset["train"].map(template_dataset, remove_columns=list(dataset["train"].features))
print(train_dataset[randint(0, len(dataset))]["text"])
test_dataset = dataset["test"].map(template_dataset, remove_columns=list(dataset["test"].features))为了训练我们的模型，我们需要将输入（文本）转换为标记 ID。这是通过 Hugging Face Transformers Tokenizer 完成的。除了 QLoRA，我们还将使用 bitsandbytes 4 位精度将冻结的 LLM 量化为 4 位，并在其上附加 LoRA 适配器。
定义训练函数
model_id = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
from accelerate import Accelerator
import datetime
from huggingface_hub import snapshot_download
import os
from peft import AutoPeftModelForCausalLM, LoraConfig, get_peft_model, prepare_model_for_kbit_training
from sagemaker.remote_function import remote
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, set_seed
import transformers
import traceback
# Start training
@remote(
keep_alive_period_in_seconds=0,
volume_size=100,
job_name_prefix=f"train-{model_id.split('/')[-1].replace('.', '-')}",
use_torchrun=True,
nproc_per_node=4
)
def train_fn(
model_name,
train_ds,
test_ds=None,
lora_r=8,
lora_alpha=16,
lora_dropout=0.1,
per_device_train_batch_size=8,
per_device_eval_batch_size=8,
gradient_accumulation_steps=1,
learning_rate=2e-4,
num_train_epochs=1,
fsdp="",
fsdp_config=None,
gradient_checkpointing=False,
merge_weights=False,
seed=42,
token=None
):
def init_distributed():
# Initialize the process group
torch.distributed.init_process_group(
backend="nccl", # Use "gloo" backend for CPU
timeout=datetime.timedelta(seconds=5400)
)
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
return local_rank
if torch.cuda.is_available() and (torch.cuda.device_count() > 1 or int(os.environ.get("SM_HOST_COUNT", 1)) > 1):
# Call this function at the beginning of your script
local_rank = init_distributed()
# Now you can use distributed functionalities
torch.distributed.barrier(device_ids=[local_rank])
os.environ.update({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
set_seed(seed)
accelerator = Accelerator()
if token is not None:
os.environ.update({"HF_TOKEN": token})
accelerator.wait_for_everyone()
print("Downloading model ", model_name)
os.makedirs("/tmp/tmp_folder", exist_ok=True)
snapshot_download(repo_id=model_name, local_dir="/tmp/tmp_folder")
print(f"Downloading model {model_name} under /tmp/tmp_folder")
model_name = "/tmp/tmp_folder"
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Set Tokenizer pad Token
tokenizer.pad_token = tokenizer.eos_token
with accelerator.main_process_first():
# tokenize and chunk dataset
lm_train_dataset = train_ds.map(
lambda sample: tokenizer(sample["text"]), remove_columns=list(train_ds.features)
)
print(f"Total number of train samples: {len(lm_train_dataset)}")
if test_ds is not None:
lm_test_dataset = test_ds.map(
lambda sample: tokenizer(sample["text"]), remove_columns=list(test_ds.features)
)
print(f"Total number of test samples: {len(lm_test_dataset)}")
else:
lm_test_dataset = None
torch_dtype = torch.bfloat16
# Defining additional configs for FSDP
if fsdp != "" and fsdp_config is not None:
bnb_config_params = {
"bnb_4bit_quant_storage": torch_dtype
}
model_configs = {
"torch_dtype": torch_dtype
}
fsdp_configurations = {
"fsdp": fsdp,
"fsdp_config": fsdp_config,
"gradient_checkpointing_kwargs": {
"use_reentrant": False
},
"tf32": True
}
else:
bnb_config_params = dict()
model_configs = dict()
fsdp_configurations = dict()
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch_dtype,
**bnb_config_params
)
model = AutoModelForCausalLM.from_pretrained(
model_name,
trust_remote_code=True,
quantization_config=bnb_config,
attn_implementation="flash_attention_2",
use_cache=not gradient_checkpointing,
cache_dir="/tmp/.cache",
**model_configs
)
if fsdp == "" and fsdp_config is None:
model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=gradient_checkpointing)
if gradient_checkpointing:
model.gradient_checkpointing_enable()
config = LoraConfig(
r=lora_r,
lora_alpha=lora_alpha,
target_modules="all-linear",
lora_dropout=lora_dropout,
bias="none",
task_type="CAUSAL_LM"
)
model = get_peft_model(model, config)
trainer = transformers.Trainer(
model=model,
train_dataset=lm_train_dataset,
eval_dataset=lm_test_dataset if lm_test_dataset is not None else None,
args=transformers.TrainingArguments(
per_device_train_batch_size=per_device_train_batch_size,
per_device_eval_batch_size=per_device_eval_batch_size,
gradient_accumulation_steps=gradient_accumulation_steps,
gradient_checkpointing=gradient_checkpointing,
logging_strategy="steps",
logging_steps=1,
log_on_each_node=False,
num_train_epochs=num_train_epochs,
learning_rate=learning_rate,
bf16=True,
ddp_find_unused_parameters=False,
save_strategy="no",
output_dir="outputs",
**fsdp_configurations
),
data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
trainer.train()
if trainer.accelerator.is_main_process:
trainer.model.print_trainable_parameters()
if trainer.is_fsdp_enabled:
trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")
if merge_weights:
output_dir = "/tmp/model"
# merge adapter weights with base model and save
# save int 4 model
trainer.model.save_pretrained(output_dir, safe_serialization=False)
if accelerator.is_main_process:
# clear memory
del model
del trainer
torch.cuda.empty_cache()
# load PEFT model
model = AutoPeftModelForCausalLM.from_pretrained(
output_dir,
torch_dtype=torch.float16,
low_cpu_mem_usage=True,
trust_remote_code=True,
)
# Merge LoRA and base model and save
model = model.merge_and_unload()
model.save_pretrained(
"/opt/ml/model", safe_serialization=True, max_shard_size="2GB"
)
else:
trainer.model.save_pretrained("/opt/ml/model", safe_serialization=True)
if accelerator.is_main_process:
tokenizer.save_pretrained("/opt/ml/model")
accelerator.wait_for_everyone()
train_fn(
model_id,
train_ds=train_dataset,
test_ds=test_dataset,
per_device_train_batch_size=2,
per_device_eval_batch_size=1,
gradient_accumulation_steps=2,
gradient_checkpointing=True,
num_train_epochs=1,
fsdp="full_shard auto_wrap offload",
fsdp_config={
'backward_prefetch': 'backward_pre',
'cpu_ram_efficient_loading': True,
'offload_params': True,
'forward_prefetch': False,
'use_orig_params': False
},
merge_weights=True
)（四）加载微调模型
注意：使用 merge_weights=True 运行 train_fn 以合并训练好的适配器
下载模型
import boto3
import sagemaker
sagemaker_session = sagemaker.Session()
model_id = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
bucket_name = sagemaker_session.default_bucket()
default_prefix = sagemaker_session.default_bucket_prefix
job_prefix = f"train-{model_id.split('/')[-1].replace('.', '-')}"
def get_last_job_name(job_name_prefix):
sagemaker_client = boto3.client('sagemaker')
search_response = sagemaker_client.search(
Resource='TrainingJob',
SearchExpression={
'Filters': [
{
'Name': 'TrainingJobName',
'Operator': 'Contains',
'Value': job_name_prefix
},
{
'Name': 'TrainingJobStatus',
'Operator': 'Equals',
'Value': "Completed"
}
]
},
SortBy='CreationTime',
SortOrder='Descending',
MaxResults=1)
return search_response['Results'][0]['TrainingJob']['TrainingJobName']
job_name = get_last_job_name(job_prefix)
job_name推理配置
import sagemaker
from sagemaker import get_execution_role
from sagemaker import Model
instance_count = 1
instance_type = "ml.g5.4xlarge"
health_check_timeout = 700
image_uri = sagemaker.image_uris.retrieve(
framework="djl-lmi",
region=sagemaker_session.boto_session.region_name,
version="latest"
)
image_uri
if default_prefix:
model_data = f"s3://{bucket_name}/{default_prefix}/{job_name}/{job_name}/output/model.tar.gz"
else:
model_data = f"s3://{bucket_name}/{job_name}/{job_name}/output/model.tar.gz"
model = Model(
image_uri=image_uri,
model_data=model_data,
role=get_execution_role(),
env={
'HF_MODEL_ID': "/opt/ml/model", # path to where sagemaker stores the model
'OPTION_TRUST_REMOTE_CODE': 'true',
'OPTION_ROLLING_BATCH': "vllm",
'OPTION_DTYPE': 'bf16',
'OPTION_TENSOR_PARALLEL_DEGREE': 'max',
'OPTION_MAX_ROLLING_BATCH_SIZE': '1',
'OPTION_MODEL_LOADING_TIMEOUT': '3600',
'OPTION_MAX_MODEL_LEN': '4096'
}
)
model_id = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
endpoint_name = f"{model_id.split('/')[-1].replace('.', '-')}-djl"
predictor = model.deploy(
endpoint_name=endpoint_name,
initial_instance_count=instance_count,
instance_type=instance_type,
container_startup_health_check_timeout=health_check_timeout,
model_data_download_timeout=3600
)预测
import sagemaker
sagemaker_session = sagemaker.Session()
model_id = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
endpoint_name = f"{model_id.split('/')[-1].replace('.', '-')}-djl"
predictor = sagemaker.Predictor(
endpoint_name=endpoint_name,
sagemaker_session=sagemaker_session,
serializer=sagemaker.serializers.JSONSerializer(),
deserializer=sagemaker.deserializers.JSONDeserializer(),
)
base_prompt = f"""<|begin_of_text|><|start_header_id|>user<|end_header_id|>{
{question}}<|eot_id|><|start_header_id|>assistant<|end_header_id|>"""
prompt = base_prompt.format(question="What statue is in front of the Notre Dame building?")
response = predictor.predict({
"inputs": prompt,
"parameters": {
"temperature": 0.2,
"top_p": 0.9,
"return_full_text": False,
"stop": ['<|eot_id|>', '<|end_of_text|>']
}
})
response = response["generated_text"].split("<|end_of_text|>")[0]
response删除端点
predictor.delete_model()
predictor.delete_endpoint(delete_endpoint_config=True)四、实验结果与分析
（一）性能指标对比
将精调后的 DeepSeek-R1-Distilled-Llama-8B 模型与预训练模型在相同任务上进行性能对比。结果显示，精调后的模型在准确率、F1 值等指标上有显著提升，表明精调过程有效地提高了模型在特定任务上的性能表现。
（二）结果分析
- 优势分析：精调后的模型能够更好地理解和处理特定任务的数据，生成更准确、更符合需求的回答。这得益于在精调过程中，模型学习了目标任务的特定模式和知识，增强了对任务的适应性。
- 局限性分析：尽管精调后的模型性能有所提升，但仍然存在一些局限性。例如，在处理一些复杂的、语义模糊的问题时，模型的回答可能不够准确或完整。这可能是由于数据集的覆盖范围有限，或者模型的架构在处理复杂语义时存在一定的困难。
五、总结与展望
（一）实验总结
本次实验在 AWS SageMaker 平台上成功对 DeepSeek-R1-Distilled-Llama-8B 模型进行了精调，通过合理配置实验环境、选择合适的精调算法和优化策略，提升了模型在特定任务上的性能。实验结果表明，SageMaker 平台为语言模型的精调提供了强大的支持，QLoRA 算法在减少计算成本的同时有效地提高了模型性能。
（二）未来展望
- 模型改进方向：进一步探索更先进的精调算法和技术，如基于注意力机制的优化方法，以提高模型对复杂语义的理解和处理能力。同时，尝试增加数据集的规模和多样性，提高模型的泛化能力。
- 应用拓展：将精调后的模型应用到更多的实际场景中，如智能客服、智能写作等领域，验证模型的实用性和有效性。通过实际应用的反馈，不断优化模型，提升用户体验。
