LoRA（低秩自适应）是一种针对大型预训练模型的参数高效微调技术。与需要更新全部模型参数的标准全参数微调不同，LoRA 会冻结整个模型，仅引入少量可训练参数。这些参数被添加到模型的特定层或模块中，从而以最小的内存开销实现高效适配。
由于 LoRA 仅需存储可训练参数的优化器状态和梯度，其内存消耗远低于全参数微调。但由于其他模型参数保持冻结状态，该方法无法适配新词元——任何新增词元都将对应未经训练的嵌入向量。
在先前文章中，我们探讨过如何在完全重训词元嵌入层和语言建模头的情况下使用 LoRA 微调。
这种方法使模型能有效处理特殊词元（例如聊天模板和特定领域中的专用词元）。虽然该方案效果显著，但由于需要将嵌入层和语言建模头参数设为可训练状态，会显著增加内存消耗。
本文将探讨 Hugging Face PEFT（参数高效微调）库提供的一种全新替代方案。与传统方法需要重新训练完整嵌入层和语言建模头部不同，该技术仅专注于更新模型需要学习的特殊标记对应的嵌入向量。我们将首先解析该技术的工作原理、局限性及其内存效率优势，最后与完整重训练方法进行性能对比。
LoRA 可训练标记技术原理剖析
在 PEFT 库此次更新之前，修改词元嵌入需要重新训练整个嵌入矩阵，而该矩阵的规模可能大得惊人。这个矩阵通常包含数十万个词元，每个词元都由高维向量表示。以 Llama 3.1 8B 模型为例：它拥有约 128k 个词元（ vocab_size ），嵌入维度（ embed_dim ）为 4,096，这意味着总参数量达到 128k * 4,096 = 5.24 亿。
训练这些参数会消耗大量内存：
- 16 位权重：5.24 亿 * 2 字节 = 1GB
- 这些参数的 32 位优化器状态：5.24 亿 * (2 * 4 字节) = 4GB
总内存使用量：1 GB（权重） + 4 GB（优化器状态） = 5 GB
这不仅增加了计算和存储需求，还可能对与任务无关的嵌入引入不必要的更新。
“可训练词元”通过引入稀疏更新机制提供了一种更高效的替代方案。它不再更新完整的 (vocab_size, embed_dim) 嵌入矩阵，而是创建一个更小的 (n, embed_dim) 矩阵，其中 n 代表您希望修改的词元数量。在前向传播过程中，这些稀疏嵌入会被叠加在原始矩阵之上，使得只有指定的词元索引可训练，其余部分保持冻结状态。这种设计既能针对性地调整词元表征，又能保持原始模型的整体完整性，同时将内存和计算成本控制在最低水平。
有两种主要方法可以将此功能集成到模型训练工作流程中，两者都依赖于 PEFT 框架。
1、该方法专为稀疏嵌入更新设计。首先需要定义一个 CustomTokensConfig ，其中明确指定 target_modules （通常设置为 ['embed_tokens'] ）和 token_indices （即 tokenizer.json 中定义的需要微调的 token ID 列表）。配置完成后，使用 get_peft_model(model, peft_config) 封装基础模型。这会创建一个内部管理稀疏嵌入更新的封装器，确保仅训练指定的 token。
peft_config = CustomTokensConfig(target_modules=['embed_tokens'], token_indices=[0, 1, 2])
peft_model = get_peft_model(model, peft_config)2、 结合 LoRA：
若您正使用 LoRA 调整模型中的线性层，还可将其功能扩展至支持词元级嵌入更新。这通过 LoraConfig 中的 trainable_token_indices 参数实现：传入一个字典，将目标模块名称（如 'embed_tokens' ）映射至应可训练的词元索引列表。该方法仅微调指定索引关联的嵌入，同时继续对其他模型组件应用 LoRA。
peft_config = LoraConfig(
lora_alpha=16,
lora_dropout=0.05,
r=16,
bias="none",
task_type="CAUSAL_LM",
target_modules= ['k_proj', 'q_proj', 'v_proj', 'o_proj', "gate_proj", "down_proj", "up_proj"],
modules_to_save = ["embed_tokens", "lm_head"] if full_retrain_vocab else None,
trainable_token_indices = {'embed_tokens':[128006, 128007, 128008, 128009], 'lm_head':[128006, 128007, 128008, 128009]},
)两种方法均兼容 PEFT 的保存/加载机制，可独立于基础模型导出和重新加载微调后的权重。
在接下来的章节中，我们将重点探讨如何将其应用于 LoRA 微调。
可训练标记 vs. 嵌入层的完全重训练
在本节中，我们将探讨一个常见场景：需要微调基础模型以使用聊天模板。这种情况下，我们必须在词汇表中添加特殊标记（例如对话分隔符），或者使用词汇表中已定义的特殊标记。
即便像 Llama 3.x 和 Qwen2.5 基础模型那样，词汇表中已包含特殊标记，这些标记也可能未经有效预训练。这些标记的嵌入值通常极小（近乎 0.0），容易在训练过程中引发 NaN（非数值）问题。
以 Llama 3.1 8B 和 Llama 3.2 3B 的基础模型为例：
注：Llama 3.2 3B 采用绑定嵌入技术，即词元嵌入层与语言建模头共享相同参数。
我们的目标是训练模型有效使用对话模板。为简化流程，我们将采用该模型指导版本中的对话模板。这是便捷之选，因为基础模型的词汇表已包含该对话模板使用的特殊标记：
{
"id": 128006,
"content": "<|start_header_id|>",
"single_word": false,
"lstrip": false,
"rstrip": false,
"normalized": false,
"special": true
},
{
"id": 128007,
"content": "<|end_header_id|>",
"single_word": false,
"lstrip": false,
"rstrip": false,
"normalized": false,
"special": true
},
{
"id": 128008,
"content": "<|eom_id|>",
"single_word": false,
"lstrip": false,
"rstrip": false,
"normalized": false,
"special": true
},
{
"id": 128009,
"content": "<|eot_id|>",
"single_word": false,
"lstrip": false,
"rstrip": false,
"normalized": false,
"special": true
}, 尽管这些特殊标记存在于词汇表中，但基础模型并未经过训练去使用它们。对于这些模型而言，这些特殊标记与词汇表中其他保留的特殊标记被同等对待，例如标记 128009 的嵌入向量可能与标记 128008 的嵌入向量过于接近，导致模型无法区分二者。
"可训练标记"策略完美契合这一场景。相较于重新训练整个嵌入矩阵，我们只需针对这 4 个标记的嵌入进行重新训练。
我们需要回答的关键问题是：
- 与完全重新训练嵌入矩阵（词元嵌入和语言建模头部）相比，能节省多少内存？
- 这种方法速度更快吗？
- 它的表现是否同样出色？
为了解答这些问题，我在 Ultrachat 数据集的 1 万个样本上对模型进行了两轮训练。通过应用聊天模板将"messages"列转换为文本序列。
我对比了四种 LoRA 微调配置方案：
- 冻结嵌入层
- 可训练标记：标记嵌入层和语言建模头部（LM 头部）
- 可训练标记：仅标记嵌入
- 全量重训练嵌入层
可训练参数与内存消耗
"可训练词元"策略的设计理念是仅在现有 LoRA 适配器参数基础上添加少量可训练参数，从而实现内存高效利用。相比之下，全量重训练嵌入层需要更新海量参数，将导致内存消耗显著增加。
正如预期，仅训练 4 个 token 在 LoRA 基础上增加的参数量微乎其微。我们以 Llama 3.1 8B 为例进行具体计算：
- 冻结嵌入层时：41,943,040 个可训练参数
- 仅训练词元嵌入时：额外增加 16,384 个可训练参数
这意味着每个被训练的词元对应 4,096 个可训练参数，因为 4,096 与模型的隐藏层维度一致。这个计算是正确的。添加一个可训练词元会引入 4,096 个可训练参数，如果同时训练语言模型头部则会增加 8,192 个参数。
在使用可训练标记重新训练嵌入层和语言模型头部时，我们观察到额外消耗了约 2GB 内存。对于仅 16k 参数而言，这看似消耗过大。其根本原因在于训练过程中，我们必须在某个时刻将整个嵌入矩阵存储并实体化在内存中以进行更新。
如果我们不对 Llama 3.1 8B 模型的 LM 头部进行词元训练（仅训练词元嵌入层），可训练词元仅比使用冻结嵌入层多消耗 16MB 内存。这一现象在 Llama 3.2 3B 模型中并未观察到。由于 LM 头部与词元嵌入层共享参数，我们需要实例化完整的语言建模头部。在此配置下，无论是否训练 LM 头部，仅训练词元嵌入层所消耗的内存是相同的。
虽然 Llama 3.1 8B 额外消耗的 2GB 内存相较于原始 LoRA 方法看似显著，但这仍比完全重新训练嵌入层和 LM 头部所需内存少 5GB。这是因为在使用可训练词元时，我们只需为少量可训练参数创建优化器状态和梯度。
我试图解决的另一个问题是：仅训练少量词元是否比完全重新训练所有嵌入层更快？
事实证明，仅训练少量词元并不会加快速度。实际上反而会稍慢一些，特别是当我们同时训练语言模型头部（LM head）和可训练词元时。这种减速现象很可能源于更新 LM 头部带来的额外计算复杂度，以及需要为重新训练的词元创建新矩阵。不过通过代码优化，这个过程应该能获得加速。
仅训练部分词元时我们会错过什么？
当内存充足时，完整训练词嵌入层和 LM 头部通常能提升 LoRA 微调的效果。我特别推荐在特定领域或专业任务微调时采用这种方法，因为基础词嵌入可能并不适配这些场景。
另一方面，完整训练词嵌入层也存在风险——这需要更新所有嵌入参数，若操作不当（例如学习率设置过高），反而可能损害模型性能。
让我们通过观察学习曲线来更好地理解这一点：
仅训练可训练词元的词嵌入时，模型表现往往欠佳，导致验证损失较高。这一现象在 Llama 3.1 8B 上尤为明显——完全冻结词嵌入的 LoRA 方法优于可训练词元方案。但若同时训练语言模型头部(LM head)，结果会显著提升，甚至超越 Llama 3.1 8B 中完整重训练词嵌入与语言模型头部的性能表现。
这些学习曲线表明，当我们对特殊标记的语言模型头部进行重训练时，模型能更高效地从数据集中学习。
结论
"可训练标记"策略在教导模型如何使用新标记方面既高效又有效。与完全重新训练嵌入矩阵相比，这种方法效果出色且节省内存。我强烈推荐此方法！
在这篇文章中，我特别将其应用于聊天模板中的特殊标记。但如果您在特定领域工作，可以调整这一策略，针对领域相关术语的标记进行训练。该方法同样适用于编程任务——您可以为目标编程语言特有的符号和关键词训练专属标记。
添加作者个人微信：wintersweet0flower（验证消息请注明CSDN）
分享LORA微调、QLoRA微调、RLHF以及RAG的代码实例
还可以就技术问题进行咨询。
代码
Installation
pip install --upgrade transformers bitsandbytes peft accelerate datasets trl flash_attnFine-Tuning Code
import torch, os, multiprocessing
from datasets import load_dataset
from peft import LoraConfig, prepare_model_for_kbit_training
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
BitsAndBytesConfig,
set_seed
)
from trl import SFTTrainer, SFTConfig
set_seed(1234)
compute_dtype = torch.bfloat16
attn_implementation = 'flash_attention_2'
def fine_tune(model_name, batch_size=1, gradient_accumulation_steps=32, LoRA=False, QLoRA=False, full_retrain_vocab=True, trainable_tokens=None):
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = "<|finetune_right_pad_id|>"
tokenizer.pad_token_id = 128004
tokenizer.padding_side = 'right'
tokenizer_name_chat_template = "meta-llama/Llama-3.1-8B-Instruct"
tokenizer_chat = AutoTokenizer.from_pretrained(tokenizer_name_chat_template)
ds_train = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft[:20000]")
ds_test = load_dataset("HuggingFaceH4/ultrachat_200k", split="test_sft[:1000]")
def process(row):
prompt_messages = tokenizer_chat.apply_chat_template(row["messages"], tokenize=False)
row["text"] = prompt_messages
return row
ds_train = ds_train.map(
process,
num_proc= multiprocessing.cpu_count(),
load_from_cache_file=False,
)
ds_test = ds_test.map(
process,
num_proc= multiprocessing.cpu_count(),
load_from_cache_file=False,
)
ds_train = ds_train.remove_columns(["messages", "prompt"])
ds_test = ds_test.remove_columns(["messages", "prompt"])
if QLoRA:
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=compute_dtype,
bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(
model_name, quantization_config=bnb_config, device_map={"": 0}, attn_implementation=attn_implementation
)
model = prepare_model_for_kbit_training(model, gradient_checkpointing_kwargs={'use_reentrant':True})
else:
model = AutoModelForCausalLM.from_pretrained(
model_name, device_map={"": 0}, torch_dtype=compute_dtype, attn_implementation=attn_implementation
)
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant':True})
if LoRA or QLoRA:
peft_config = LoraConfig(
lora_alpha=16,
lora_dropout=0.05,
r=16,
bias="none",
task_type="CAUSAL_LM",
target_modules= ['k_proj', 'q_proj', 'v_proj', 'o_proj', "gate_proj", "down_proj", "up_proj"],
modules_to_save = ["embed_tokens", "lm_head"] if full_retrain_vocab else None,
trainable_token_indices = trainable_tokens,
)
else:
peft_config = None
if LoRA:
output_dir = "./LoRA/"
elif QLoRA:
output_dir = "./QLoRA/"
else:
output_dir = "./FFT/"
training_arguments = SFTConfig(
output_dir=output_dir,
eval_strategy="steps",
do_eval=True,
optim="adamw_8bit",
per_device_train_batch_size=batch_size,
gradient_accumulation_steps=gradient_accumulation_steps,
per_device_eval_batch_size=batch_size,
log_level="debug",
save_strategy="epoch",
logging_steps=25,
learning_rate=1e-5,
bf16 = True,
eval_steps=25,
num_train_epochs=1,
warmup_ratio=0.1,
lr_scheduler_type="linear",
dataset_text_field="text",
max_seq_length=512,
report_to="none",
)
trainer = SFTTrainer(
model=model,
train_dataset=ds_train,
eval_dataset=ds_test,
peft_config=peft_config,
processing_class=tokenizer,
args=training_arguments,
)
#--code by Unsloth: https://colab.research.google.com/drive/1Ys44kVvmeZtnICzWz0xgpRnrIOjZAuxp?usp=sharing#scrollTo=pCqnaKmlO1U9
gpu_stats = torch.cuda.get_device_properties(0)
start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.")
print(f"{start_gpu_memory} GB of memory reserved.")
trainer_ = trainer.train()
used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
used_memory_for_trainer= round(used_memory - start_gpu_memory, 3)
used_percentage = round(used_memory /max_memory*100, 3)
trainer_percentage = round(used_memory_for_trainer/max_memory*100, 3)
print(f"{trainer_.metrics['train_runtime']} seconds used for training.")
print(f"{round(trainer_.metrics['train_runtime']/60, 2)} minutes used for training.")
print(f"Peak reserved memory = {used_memory} GB.")
print(f"Peak reserved memory for training = {used_memory_for_trainer} GB.")
print(f"Peak reserved memory % of max memory = {used_percentage} %.")
print(f"Peak reserved memory for training % of max memory = {trainer_percentage} %.")
print("-----")
#----Trainable Special Tokens from the Chat Template
{
"id": 128006,
"content": "<|start_header_id|>",
"single_word": false,
"lstrip": false,
"rstrip": false,
"normalized": false,
"special": true
},
{
"id": 128007,
"content": "<|end_header_id|>",
"single_word": false,
"lstrip": false,
"rstrip": false,
"normalized": false,
"special": true
},
{
"id": 128008,
"content": "<|eom_id|>",
"single_word": false,
"lstrip": false,
"rstrip": false,
"normalized": false,
"special": true
},
{
"id": 128009,
"content": "<|eot_id|>",
"single_word": false,
"lstrip": false,
"rstrip": false,
"normalized": false,
"special": true
},Trainable Tokens for The Token Embeddings and LM Head
fine_tune("meta-llama/Llama-3.2-3B", batch_size=1, gradient_accumulation_steps=32, LoRA=True, QLoRA=False, full_retrain_vocab=False, trainable_tokens={'embed_tokens':[128006, 128007, 128008, 128009], 'lm_head':[128006, 128007, 128008, 128009]})Trainable Tokens Only for The Token Embeddings
fine_tune("meta-llama/Llama-3.2-3B", batch_size=1, gradient_accumulation_steps=32, LoRA=True, QLoRA=False, full_retrain_vocab=False, trainable_tokens={'embed_tokens':[128006, 128007, 128008, 128009]})Full Retraining of the Token Embeddings and LM Head
fine_tune("meta-llama/Llama-3.2-3B", batch_size=1, gradient_accumulation_steps=32, LoRA=True, QLoRA=False, full_retrain_vocab=True, trainable_tokens=None)Standard LoRA with Frozen Embeddings
fine_tune("meta-llama/Llama-3.2-3B", batch_size=1, gradient_accumulation_steps=32, LoRA=True, QLoRA=False, full_retrain_vocab=False, trainable_tokens=None)Same but with Llama 3.1 8B
fine_tune("meta-llama/Llama-3.1-8B", batch_size=1, gradient_accumulation_steps=32, LoRA=True, QLoRA=False, full_retrain_vocab=False, trainable_tokens={'embed_tokens':[128006, 128007, 128008, 128009], 'lm_head':[128006, 128007, 128008, 128009]})
fine_tune("meta-llama/Llama-3.1-8B", batch_size=1, gradient_accumulation_steps=32, LoRA=True, QLoRA=False, full_retrain_vocab=False, trainable_tokens={'embed_tokens':[128006, 128007, 128008, 128009]})
fine_tune("meta-llama/Llama-3.1-8B", batch_size=1, gradient_accumulation_steps=32, LoRA=True, QLoRA=False, full_retrain_vocab=True, trainable_tokens=None)
fine_tune("meta-llama/Llama-3.1-8B", batch_size=1, gradient_accumulation_steps=32, LoRA=True, QLoRA=False, full_retrain_vocab=False, trainable_tokens=None)
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
