在本项目中,微调脚本文件finetune.py 提供了一套全面的工具,用于对 DeepSeek-MoE 预训练语言模型进行微调。支持加载特定任务的数据、对数据进行预处理和编码,以及通过多种配置选项(如 LoRA 量化、分布式训练等)对模型进行高效训练。用户可以根据自己的需求,通过命令行参数或配置文件调整微调策略,以优化模型在特定任务或数据集上的性能。
在DeepSeek-MoE 项目中,文件finetune.py 通过以下技术对预训练模型进行微调。
文件finetune.py的具体实现流程如下:
(1)模型加载与配置
(2)数据加载与预处理
(3)训练准备
(4)训练过程
(5)模型保存与推理
通过上述流程,finetune.py能够在高效利用计算资源的同时,实现对 MoE 模型的快速微调,使其适应特定的下游任务。
定义函数 build_instruction_prompt(),用于生成包含指令和响应格式的提示文本。它接受一个字符串参数 instruction,并返回一个多行字符串,其中包含预定义的上下文说明、指令部分以及响应部分的模板。在生成的字符串中,instruction 的内容被插入到指定位置,以构建完整的提示文本。
# 定义常量
IGNORE_INDEX = -100 # 忽略的索引值
EOT_TOKEN = "<|EOT|>" # 响应结束标记
logger = logging.getLogger(__name__) # 创建日志记录器
# 构建指令提示的函数
def build_instruction_prompt(instruction: str):
return f'''
You are an AI assistant, developed by DeepSeek Company. For politically sensitive questions, security and privacy issues, you will refuse to answer.
### Instruction:
{instruction.strip()}
### Response:
'''.strip()
定义一个名为 ModelArguments 的数据类,用于配置模型微调的参数。其中包括可训练的模型参数列表、LoRA(低秩适应)相关设置、需要保存的模块、是否使用 LoRA、预训练模型的路径、注意力机制的实现方式,以及量化相关的配置,如是否使用双量化、量化数据类型和位数等。这些配置有助于在微调过程中灵活地调整模型的各项参数,以满足不同的训练需求。
# 定义模型参数类
@dataclass
class ModelArguments:
trainable: Optional[str] = field(
default="q_proj,v_proj,k_proj,o_proj,gate_proj,down_proj,up_proj", # 默认可训练的模型参数
metadata={"help": "Comma-separated list of model parameters to train."}
)
lora_rank: Optional[int] = field(
default=8, # LoRA 的秩
metadata={"help": "Rank of LoRA"}
)
lora_dropout: Optional[float] = field(
default=0.1, # LoRA 的 dropout 率
metadata={"help": "Dropout rate for LoRA."}
)
lora_alpha: Optional[float] = field(
default=32.0, # LoRA 的 alpha 值
metadata={"help": "Alpha value for LoRA."}
)
modules_to_save: Optional[str] = field(
default="embed_tokens,lm_head", # 需要保存的模块
metadata={"help": "Comma-separated list of modules to save."}
)
use_lora: Optional[bool] = field(
default=False, # 是否使用 LoRA
metadata={"help": "Whether to use LoRA."}
)
model_name_or_path: Optional[str] = field(
default="deepseek-ai/deepseek-moe-16b", # 模型路径
metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models."}
)
attn_implementation: Optional[str] = field(
default="flash_attention_2", # 注意力实现方式
metadata={"help": "Attention implementation to use."}
)
double_quant: bool = field(
default=True, # 是否使用双量化
metadata={"help": "Compress the quantization statistics through double quantization."}
)
quant_type: str = field(
default="nf4", # 量化数据类型
metadata={"help": "Quantization data type to use. Should be one of `fp4` or `nf4`."}
)
bits: int = field(
default=16, # 使用的位数
metadata={"help": "How many bits to use."}
)
定义数据类DataArguments,用于指定训练数据的路径。它包含一个名为 data_path 的字段,用于存储训练数据的位置。
# 定义数据参数类
@dataclass
class DataArguments:
data_path: str = field(
default=None, # 数据路径
metadata={"help": "Path to the training data."}
)
类TrainingArguments继承自 Hugging Face 的 transformers.TrainingArguments,用于配置训练过程中的超参数,如缓存目录、优化器类型和模型最大序列长度等。这些设置有助于控制训练过程的各个方面,确保模型以最佳方式进行微调。
# 定义训练参数类
@dataclass
class TrainingArguments(transformers.TrainingArguments):
cache_dir: Optional[str] = field(default=None) # 缓存目录
optim: str = field(
default="adamw_torch", # 优化器类型
metadata={"help": "Optimizer to use."}
)
model_max_length: int = field(
default=512, # 模型最大序列长度
metadata={"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."}
)
SavePeftModelCallback 是一个自定义的回调类,旨在与 Hugging Face 的 Trainer 一起使用,以在训练过程中正确保存 PEFT(参数高效微调)模型。该回调在模型保存时,仅保存 PEFT 模型的适配器权重,而不是整个基础模型,从而节省存储空间并提高加载效率。此外,在训练结束时,它会创建一个名为 completed 的文件,指示训练已成功完成。
# 定义保存 LoRA 模型回调类
class SavePeftModelCallback(transformers.TrainerCallback):
def save_model(self, args, state, kwargs):
logger.info('Saving PEFT checkpoint...') # 保存 LoRA 模型
if state.best_model_checkpoint is not None:
checkpoint_folder = os.path.join(state.best_model_checkpoint, "adapter_model")
else:
checkpoint_folder = os.path.join(args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}")
peft_model_path = os.path.join(checkpoint_folder, "adapter_model")
kwargs["model"].save_pretrained(peft_model_path)
kwargs["tokenizer"].save_pretrained(peft_model_path)
def on_save(self, args, state, control, **kwargs):
self.save_model(args, state, kwargs)
return control
def on_train_end(self, args, state, control, **kwargs):
def touch(fname, times=None):
with open(fname, 'a'):
os.utime(fname, times)
touch(os.path.join(args.output_dir, 'completed'))
self.save_model(args, state, kwargs)
在使用 PEFT 进行模型微调时,正确保存和加载适配器权重至关重要。根据 Hugging Face 的文档,建议在保存模型时,仅保存适配器权重,以避免不必要地保存整个基础模型。
另外,Hugging Face 的官方文档提供了有关如何使用回调自定义训练过程的详细信息。 通过实现 SavePeftModelCallback,您可以确保在训练过程中正确保存 PEFT 模型的适配器权重,并在训练结束时生成一个指示训练成功完成的文件。
函数get_last_checkpoint()用于获取指定目录中最新的检查点。首先检查目录是否存在一个名为 'completed' 的文件,如果存在,则表示训练已完成,返回 None。否则,会遍历目录中以 'checkpoint-' 开头的子目录,找到其中编号最大的一个,认为其是最新的检查点,并返回其路径。如果未找到任何符合条件的子目录,同样返回 None。
# 获取最新的检查点
def get_last_checkpoint(checkpoint_dir):
if os.path.isdir(checkpoint_dir):
is_completed = os.path.exists(os.path.join(checkpoint_dir, 'completed'))
if is_completed:
return None
max_step = 0
for filename in os.listdir(checkpoint_dir):
if os.path.isdir(os.path.join(checkpoint_dir, filename)) and filename.startswith(PREFIX_CHECKPOINT_DIR):
max_step = max(max_step, int(filename.replace(f'{PREFIX_CHECKPOINT_DIR}-', '')))
if max_step == 0:
return None
latest_ckpt_dir = os.path.join(checkpoint_dir, f'{PREFIX_CHECKPOINT_DIR}-{max_step}')
logger.info(f"Found a previous checkpoint at: {checkpoint_dir}")
return latest_ckpt_dir
return None
函数safe_save_model_for_hf_trainer()的主要功能是将 Hugging Face 的 Trainer 对象中的模型安全地保存到指定的目录中。在保存模型时,该函数首先将模型的状态字典(state_dict)从 GPU 内存转移到 CPU 内存,以减少 GPU 内存的占用。然后,它调用 Trainer 对象的 _save 方法,将模型的状态字典保存到指定的输出目录。这种方法确保了在使用分布式训练或混合精度训练时,模型能够被正确地保存,避免了可能的内存问题。
# 安全保存模型以供 HF Trainer 使用
def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str):
state_dict = trainer.model.state_dict()
if trainer.args.should_save:
cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()}
del state_dict
trainer._save(output_dir, state_dict=cpu_state_dict)
函数 _tokenize_fn()作用是对输入的一批字符串进行分词处理,并返回分词后的 ID 序列以及相关的长度信息。该函数通过 tokenizer 将文本转换为模型可理解的 ID 序列,并返回这些序列及其长度信息,为后续的模型训练和推理提供了数据支持。
# 分词函数
def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict:
tokenized_list = [
tokenizer(
text,
max_length=tokenizer.model_max_length,
truncation=True,
)
for text in strings
]
input_ids = labels = [np.array(tokenized.input_ids) for tokenized in tokenized_list]
input_ids_lens = labels_lens = [len(tokenized.input_ids) for tokenized in tokenized_list]
return {
"input_ids": input_ids,
"labels": labels,
"input_ids_lens": input_ids_lens,
"labels_lens": labels_lens,
}
函数preprocess()实现了对输入源文本和目标文本的预处理,主要用于为模型训练准备结构化的输入和标签数据。其核心功能是将源文本和目标文本拼接后进行分词,并通过标记源文本部分的标签为 -100(IGNORE_INDEX),使得模型在训练时仅对目标文本部分进行预测。
# 预处理数据
def preprocess(sources: Sequence[str], targets: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict:
examples = [s + t for s, t in zip(sources, targets)]
examples_tokenized, sources_tokenized = [_tokenize_fn(strings, tokenizer) for strings in (examples, sources)]
input_ids = examples_tokenized["input_ids"]
labels = copy.deepcopy(input_ids)
for label, source_len in zip(labels, sources_tokenized["input_ids_lens"]):
label[:source_len] = IGNORE_INDEX
return {
"input_ids": input_ids,
"labels": labels,
}
下面代码定义了一个名为 DataCollatorForSupervisedDataset 的数据收集器类,用于在监督微调过程中将多个数据实例整理成模型训练所需的批量数据格式。其主要功能是将输入的 input_ids 和 labels 转换为张量,并进行填充(padding)以确保批量数据的维度一致,同时生成 attention_mask 用于指示模型哪些位置是有效的输入数据。
# 定义监督微调的数据收集器类
@dataclass
class DataCollatorForSupervisedDataset(object):
tokenizer: transformers.PreTrainedTokenizer
def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))
input_ids = [torch.tensor(x) for x in input_ids]
input_ids = torch.nn.utils.rnn.pad_sequence(
input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
)
labels = [torch.tensor(x) for x in labels]
labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
return {
"input_ids": input_ids,
"labels": labels,
"attention_mask": input_ids.ne(self.tokenizer.pad_token_id),
}
下面代码定义了一个名为 train_tokenize_function()的函数,用于对训练数据进行分词和预处理。函数train_tokenize_function()接收 examples 和 tokenizer 作为输入,examples 包含用户提供的指令和对应的输出。函数首先调用 build_instruction_prompt 为每个指令构建提示文本,然后将输出拼接上结束标记 token(EOT_TOKEN)。接着,函数对这些提示和输出进行分词,并通过 preprocess 函数将分词后的 ID 序列转换为适合模型训练的格式。最后,函数返回一个包含输入 ID 和标签的数据字典,用于后续的训练过程。
# 定义训练分词函数
def train_tokenize_function(examples, tokenizer):
sources = [build_instruction_prompt(instruction) for instruction in examples['instruction']]
targets = [f"{output}\n{EOT_TOKEN}" for output in examples['output']]
data_dict = preprocess(sources, targets, tokenizer)
return data_dict
函数build_model()用于根据提供的参数构建和配置模型,支持量化和 LoRA(Low-Rank Adaptation)配置。该函数的主要功能和实现细节如下:
# 构建模型
def build_model(model_args, training_args, checkpoint_dir):
compute_dtype = torch.bfloat16 if training_args.bf16 else torch.float16 # 确定计算数据类型
model = transformers.AutoModelForCausalLM.from_pretrained(
model_args.model_name_or_path,
load_in_4bit=model_args.bits == 4,
load_in_8bit=model_args.bits == 8,
quantization_config=BitsAndBytesConfig(
load_in_4bit=model_args.bits == 4,
load_in_8bit=model_args.bits == 8,
llm_int8_threshold=6.0,
llm_int8_has_fp16_weight=False,
bnb_4bit_compute_dtype=compute_dtype,
bnb_4bit_use_double_quant=model_args.double_quant,
bnb_4bit_quant_type=model_args.quant_type,
) if model_args.use_lora else None,
torch_dtype=compute_dtype,
trust_remote_code=True,
)
if compute_dtype == torch.float16 and model_args.bits == 4:
if torch.cuda.is_bf16_supported():
logger.warning("=" * 80)
logger.warning("您的 GPU 支持 bfloat16,您可以通过添加参数 --bf16 来加速训练!")
logger.warning("=" * 80)
setattr(model, 'model_parallel', True)
setattr(model, 'is_parallelizable', True)
model.config.torch_dtype = torch.bfloat16 if training_args.bf16 else torch.float32
if model_args.use_lora and model_args.bits < 16:
model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=training_args.gradient_checkpointing)
if model_args.use_lora:
if checkpoint_dir is not None:
logger.info(f"从 {checkpoint_dir} 加载适配器。")
model = PeftModel.from_pretrained(model, checkpoint_dir, is_trainable=True)
else:
logger.info("初始化 LoRA 模块...")
target_modules = model_args.trainable.split(',')
modules_to_save = model_args.modules_to_save
if modules_to_save is not None:
modules_to_save = modules_to_save.split(',')
lora_rank = model_args.lora_rank
lora_dropout = model_args.lora_dropout
lora_alpha = model_args.lora_alpha
peft_config = LoraConfig(
task_type=TaskType.CAUSAL_LM,
target_modules=target_modules,
inference_mode=False,
r=lora_rank,
lora_alpha=lora_alpha,
lora_dropout=lora_dropout,
modules_to_save=modules_to_save
)
model = get_peft_model(model, peft_config)
for name, module in model.named_modules():
if isinstance(module, LoraLayer):
if training_args.bf16:
module = module.to(torch.bfloat16)
if 'norm' in name or 'gate' in name:
module = module.to(torch.float32)
if 'lm_head' in name or 'embed_tokens' in name:
if hasattr(module, 'weight'):
if training_args.bf16 and module.weight.dtype == torch.float32:
module = module.to(torch.bfloat16)
return model
函数train()是 DeepSeek-MoE 项目中的主训练函数,负责整个模型的训练流程。函数train()实现了一个完整的训练流程,用于对预训练的MoE(Mixture-of-Experts)语言模型进行微调。首先加载用户提供的数据,通过分词和预处理将其转化为模型可接受的输入格式。然后,构建模型时应用了LoRA(低秩适应)和量化技术,以优化显存占用和训练效率。在训练过程中,使用了梯度累积、动态损失缩放等技术,并通过检查点机制支持从断点恢复训练。训练结束后,模型的权重会被保存下来,以便后续的评估或推理使用。
# 定义训练函数
def train():
parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
log_level = training_args.get_process_log_level()
logger.setLevel(log_level)
datasets.utils.logging.set_verbosity(log_level)
transformers.utils.logging.set_verbosity(log_level)
transformers.utils.logging.enable_default_handler()
transformers.utils.logging.enable_explicit_format()
if training_args.local_rank == 0:
logger.info("=" * 100)
logger.info(training_args)
tokenizer = transformers.AutoTokenizer.from_pretrained(
model_args.model_name_or_path,
model_max_length=training_args.model_max_length,
padding_side="right",
use_fast=True,
trust_remote_code=True
)
logger.info(f"填充标记: {tokenizer.pad_token} ({tokenizer.pad_token_id})")
logger.info(f"开始标记: {tokenizer.bos_token} ({tokenizer.bos_token_id})")
logger.info(f"结束标记: {tokenizer.eos_token} ({tokenizer.eos_token_id})")
resume_from_checkpoint_dir = get_last_checkpoint(training_args.output_dir)
model = build_model(model_args, training_args, resume_from_checkpoint_dir)
raw_train_datasets = load_dataset(
'parquet',
data_files=data_args.data_path,
split="train",
cache_dir=training_args.cache_dir
)
if training_args.local_rank > 0:
torch.distributed.barrier()
train_dataset = raw_train_datasets.map(
train_tokenize_function,
batched=True,
batch_size=3000,
num_proc=32,
remove_columns=raw_train_datasets.column_names,
load_from_cache_file=True,
desc="运行编码",
fn_kwargs={"tokenizer": tokenizer}
)
if training_args.local_rank == 0:
torch.distributed.barrier()
if training_args.local_rank == 0:
logger.info(f"训练数据集样本数量: {len(train_dataset)}")
for index in random.sample(range(len(train_dataset)), 3):
logger.info(f"训练集样本 {index}: 输入ID - {train_dataset[index]['input_ids']}, 标签 - {train_dataset[index]['labels']}.")
logger.info(f"训练集样本 {index}: 问题: {train_dataset[index]['instruction']}\n回答: {train_dataset[index]['output']}")
data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
data_module = dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)
trainer = Trainer(model=model, tokenizer=tokenizer, args=training_args, **data_module)
if model_args.use_lora:
trainer.add_callback(SavePeftModelCallback)
trainer.train(resume_from_checkpoint=resume_from_checkpoint_dir)
trainer.save_state()
if not model_args.use_lora:
safe_save_model_for_hf_trainer(trainer=trainer, output_dir=training_args.output_dir)
# 主程序入口
if __name__ == "__main__":
train()
函数train()的具体实现流程如下:
(1)解析命令行参数:使用 transformers.HfArgumentParser 解析命令行参数,将模型参数、数据参数和训练参数分别解析为 ModelArguments、DataArguments 和 TrainingArguments 类的实例。
(2)设置日志记录:根据训练参数中的日志级别设置日志记录器的日志级别,并启用默认的日志处理器和显式格式。
(3)加载分词器:使用 transformers.AutoTokenizer.from_pretrained()方法加载预训练的分词器,并设置分词器的最大序列长度、填充方向、是否使用快速分词器等参数。
(4)加载模型:调用 build_model 函数构建和配置模型,包括加载预训练模型、设置量化参数、初始化 LoRA 模块等。
(5)加载训练数据:使用 datasets.load_dataset()方法加载训练数据集,并根据数据路径、缓存目录等参数进行配置。
(6)数据预处理:使用 train_tokenize_function()对训练数据进行预处理,包括构建指令提示、分词和标签处理等。
(7)数据收集器:创建 DataCollatorForSupervisedDataset 实例,用于在训练过程中将数据整理成模型所需的格式。
(8)初始化训练器:使用 Hugging Face 的 Trainer 类初始化训练器,传入模型、分词器、训练参数和数据模块等。
(9)添加回调函数:如果启用了 LoRA,添加 SavePeftModelCallback()回调函数,用于在训练过程中保存 LoRA 适配器的权重。
(10)开始训练:调用 trainer.train 方法开始训练模型,并根据需要从检查点恢复训练。
(11)保存训练状态:训练结束后,调用 trainer.save_state()方法保存训练状态,包括模型权重、优化器状态等。
(12)保存模型:如果未启用 LoRA,调用 safe_save_model_for_hf_trainer()函数将模型权重安全地保存到指定目录。
运行下面的命令,文件finetune.py将启动训练过程,加载指定的模型和数据,并按照配置进行微调。训练过程中,检查点和日志将保存到指定的输出目录,并且可以通过 TensorBoard 查看训练指标。
DATA_PATH=""
OUTPUT_PATH=""
MODEL_PATH=""
cd finetune
deepspeed finetune.py \
--model_name_or_path $MODEL_PATH \
--data_path $DATA_PATH \
--output_dir $OUTPUT_PATH \
--num_train_epochs 3 \
--model_max_length 1024 \
--per_device_train_batch_size 16 \
--per_device_eval_batch_size 1 \
--gradient_accumulation_steps 4 \
--evaluation_strategy "no" \
--save_strategy "steps" \
--save_steps 100 \
--save_total_limit 100 \
--learning_rate 2e-5 \
--warmup_steps 10 \
--logging_steps 1 \
--lr_scheduler_type "cosine" \
--gradient_checkpointing True \
--report_to "tensorboard" \
--deepspeed configs/ds_config_zero3.json \
--bf16 True \
--use_lora False
对上述各个命令参数的具体说明如下所示:
上述令使用 DeepSpeed 来加速和优化模型的微调过程,通过指定 DeepSpeed 配置文件(如 configs/ds_config_zero3.json),可以利用其 ZeRO 优化器阶段3(ZeRO Stage 3)来有效地管理内存和计算资源,从而支持大规模模型的训练。
此外,在上述命令中还设置了梯度累积、学习率调度、日志记录等参数,以控制训练过程的各个方面。这些参数的配置需要根据具体的硬件资源和任务需求进行调整,以获得最佳的训练效果。