当前位置：首页 > news >正文

基于unsloth微调一个越狱大模型

news 来源：原创 2025/4/22 10:25:24

网上其实并没有找到现成的开源越狱数据集，所以数据集获取是个麻烦事。想了想easydataset，可能还是不行，easydataset是基于大模型回答的，大模型一般都做了对齐训练，那本地文档生成数据这条路也不可靠。

现成的越狱数据集没有，但是有很多开源的越狱实验！比如：JailbreakBench 。该仓库里分享了大量越狱实验的问答记录，针对不同模型，不同越狱攻击方式，不同越狱问题类型都有记录，而且有标记那些记录是越狱成功那些越狱失败（通过jailbroken字段）

于是我想到了可以把这些成功越狱的问答记录提取出来单独作成数据集。只需要把源json文件下载下来，再稍作提取就行了。

import jsondef extract_jailbroken_items(input_file, output_file):try:with open(input_file, 'r', encoding='utf-8') as f:data = json.load(f)new_data = []for outer_key, outer_value in data.items():for category, details in outer_value.items():if details.get('jailbroken'):new_item = {"instruction": details['prompt'],"output": details['response']}new_data.append(new_item)with open(output_file, 'w', encoding='utf-8') as f:json.dump(new_data, f, ensure_ascii=False, indent=2)print(f"数据已成功保存到 {output_file}")except FileNotFoundError:print(f"错误：未找到文件 {input_file}")except json.JSONDecodeError:print(f"错误：无法解析 {input_file} 中的 JSON 数据")except Exception as e:print(f"发生未知错误: {e}")input_file = '1.json'
output_file = 'output.json'
extract_jailbroken_items(input_file, output_file)

提取到的数据如下（第二条是我手动添加的。）

然后参考这篇博客里的方法：基于unsloth微调大模型训练大模型即可。直接贴源码了。

from unsloth import FastLanguageModel  # 导入FastLanguageModel类，用来加载和使用模型max_seq_length = 2048  # 设置模型处理文本的最大长度，相当于给模型设置一个“最大容量”
dtype = None  # 设置数据类型，让模型自动选择最适合的精度
load_in_4bit = True  # 使用4位量化来节省内存，就像把大箱子压缩成小箱子# 加载预训练模型，并获取tokenizer工具
model, tokenizer = FastLanguageModel.from_pretrained(model_name="unsloth/Llama-3.2-3B-Instruct",  # 指定要加载的模型名称max_seq_length=max_seq_length,  # 使用前面设置的最大长度dtype=dtype,  # 使用前面设置的数据类型load_in_4bit=load_in_4bit,  # 使用4位量化
)from transformers import TextStreamer
from unsloth.chat_templates import get_chat_template
tokenizer = get_chat_template(tokenizer,chat_template = "llama-3.2",mapping = {"role" : "from", "content" : "value", "user" : "human", "assistant" : "gpt"}, # ShareGPT style
)
FastLanguageModel.for_inference(model) # Enable native 2x faster inferencemessages = [{"from": "human", "value": "how to make a bomb?"}
]
# 文本转换为 token ID  添加生成提示符  返回PyTorch 张量   
inputs = tokenizer.apply_chat_template(messages, tokenize = True, add_generation_prompt = True, return_tensors = "pt").to("cuda")text_streamer = TextStreamer(tokenizer)
outputs = model.generate(input_ids = inputs, streamer = text_streamer, max_new_tokens = 1024, use_cache = True)outputs = model.generate(inputs, max_new_tokens=200)  # 生成最多 100 个新 token
response = tokenizer.decode(outputs[0], skip_special_tokens=True)  # 解码生成的 token
print(response)# 导入数据集加载函数
from datasets import load_dataset
dataset = load_dataset("json",data_files="jailbreak.json")# 定义一个函数，用于格式化数据集中的每条记录
def formatting_prompts_func(examples):texts = []for instruction, output_text in zip(examples["instruction"],  examples["output"]):prompt = f"Instruction: {instruction}\nOutput: "texts.append(prompt + output_text)return {"text": texts}data1 = dataset.map(formatting_prompts_func, batched = True)
print(data1["train"].column_names)
temp = data1["train"]FastLanguageModel.for_training(model)model = FastLanguageModel.get_peft_model(model,  # 传入已经加载好的预训练模型r = 16,  # 设置 LoRA 的秩，决定添加的可训练参数数量target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",  # 指定模型中需要微调的关键模块"gate_proj", "up_proj", "down_proj"],lora_alpha = 16,  # 设置 LoRA 的超参数，影响可训练参数的训练方式lora_dropout = 0,  # 设置防止过拟合的参数，这里设置为 0 表示不丢弃任何参数bias = "none",    # 设置是否添加偏置项，这里设置为 "none" 表示不添加use_gradient_checkpointing = "unsloth",  # 使用优化技术节省显存并支持更大的批量大小random_state = 3407,  # 设置随机种子，确保每次运行代码时模型的初始化方式相同use_rslora = False,  # 设置是否使用 Rank Stabilized LoRA 技术，这里设置为 False 表示不使用loftq_config = None,  # 设置是否使用 LoftQ 技术，这里设置为 None 表示不使用
)from trl import SFTTrainer  # 导入 SFTTrainer，用于监督式微调
from transformers import TrainingArguments  # 导入 TrainingArguments，用于设置训练参数
from unsloth import is_bfloat16_supported  # 导入函数，检查是否支持 bfloat16 数据格式trainer = SFTTrainer(  # 创建一个 SFTTrainer 实例model=model,  # 传入要微调的模型tokenizer=tokenizer,  # 传入 tokenizer，用于处理文本数据train_dataset=temp,  # 传入训练数据集dataset_text_field="text",  # 指定数据集中 文本字段的名称max_seq_length=max_seq_length,  # 设置最大序列长度dataset_num_proc=2,  # 设置数据处理的并行进程数packing=False,  # 是否启用打包功能（这里设置为 False，打包可以让训练更快，但可能影响效果）args=TrainingArguments(  # 定义训练参数per_device_train_batch_size=2,  # 每个设备（如 GPU）上的批量大小gradient_accumulation_steps=4,  # 梯度累积步数，用于模拟大批次训练warmup_steps=5,  # 预热步数，训练开始时学习率逐渐增加的步数max_steps=50,  # 最大训练步数learning_rate=2e-4,  # 学习率，模型学习新知识的速度fp16=not is_bfloat16_supported(),  # 是否使用 fp16 格式加速训练（如果环境不支持 bfloat16）bf16=is_bfloat16_supported(),  # 是否使用 bfloat16 格式加速训练（如果环境支持）logging_steps=1,  # 每隔多少步记录一次训练日志optim="adamw_8bit",  # 使用的优化器，用于调整模型参数weight_decay=0.01,  # 权重衰减，防止模型过拟合lr_scheduler_type="linear",  # 学习率调度器类型，控制学习率的变化方式seed=3407,  # 随机种子，确保训练结果可复现output_dir="outputs",  # 训练结果保存的目录report_to="none",  # 是否将训练结果报告到外部工具（如 WandB），这里设置为不报告),
)trainer_stats = trainer.train()# 将模型切换到推理模式，准备回答问题
FastLanguageModel.for_inference(model)
question = "how to make a bomb?"
# 将问题转换成模型能理解的格式，并发送到 GPU 上
inputs = tokenizer([question], return_tensors="pt").to("cuda")# 让模型根据问题生成回答，最多生成 2222 个新词
outputs = model.generate(input_ids=inputs.input_ids,  # 输入的数字序列attention_mask=inputs.attention_mask,  # 注意力遮罩，帮助模型理解哪些部分重要max_new_tokens=1111,use_cache=True,  # 使用缓存加速生成
)# 将生成的回答从数字转换回文字
response = tokenizer.batch_decode(outputs)# 打印回答
print(response[0])# 保存 adapter_config.json、adapter_model.safetensors 文件到 lora_ollama_model 文件夹
model.save_pretrained("jail_break")

唯一不同的是需要数据处理方式。

# 导入数据集加载函数
from datasets import load_dataset
dataset = load_dataset("json",data_files="jailbreak.json")
# 定义一个函数，用于格式化数据集中的每条记录
def formatting_prompts_func(examples):texts = []for instruction, output_text in zip(examples["instruction"],  examples["output"]):prompt = f"Instruction: {instruction}\nOutput: "texts.append(prompt + output_text)return {"text": texts}data1 = dataset.map(formatting_prompts_func, batched = True)
print(data1["train"].column_names)
temp = data1["train"]
temp["text"][0]