why my ‘prepare_model_for_int8_training’ is not defined?
Traceback (most recent call last):
File "/root/autodl-tmp/train_pt.py", line 67, in <module>
model = prepare_model_for_int8_training(model)
NameError: name 'prepare_model_for_int8_training' is not defined
from transformers import *
from peft import *
import torch
from datasets import load_dataset
import os
from torch.utils.data import DataLoader
from transformers import default_data_collator, get_linear_schedule_with_warmup
from tqdm import tqdm
from datasets import load_dataset
from tensorboard import *
device = "cuda"
tokenizer_name_or_path = "LLM4Binary/llm4decompile-1.3b-v1.5"
model_name_or_path = "LLM4Binary/llm4decompile-1.3b-v1.5"
dataset_name = "asm2c"
text_column = "asm text"
label_column = "text_label"
max_length = 64
lr = 3e-2
num_epochs = 50
batch_size = 8
from datasets import load_dataset
dataset = load_dataset("json", data_files="./traindata.jsonl")
dataset = dataset["train"].train_test_split(0.2)
tokenizer = AutoTokenizer.from_pretrained("LLM4Binary/llm4decompile-1.3b-v1.5")
def preprocess_function(examples):
inputs = examples["input"]
outputs = examples["output"]
# 合并input和output列
merged_texts = [f"{input} {output_text}" for input, output_text in zip(inputs, outputs)]
model_inputs = tokenizer(merged_texts, truncation=True, padding="max_length", max_length=512)
model_inputs["labels"] = model_inputs["input_ids"].copy() # 设置labels
return model_inputs
processed_datasets = dataset.map(
preprocess_function,
batched=True,
num_proc=1,
remove_columns=dataset["train"].column_names,
load_from_cache_file=False,
desc="Running tokenizer on dataset",
)
train_dataset = processed_datasets["train"]
eval_dataset = processed_datasets["test"]
peft_config = PromptTuningConfig(
task_type=TaskType.CAUSAL_LM,
prompt_tuning_init=PromptTuningInit.TEXT,
num_virtual_tokens=8,
prompt_tuning_init_text="What's the souce code of this asm?",
tokenizer_name_or_path=model_name_or_path,
)
checkpoint_name = f"{dataset_name}_{model_name_or_path}_{peft_config.peft_type}_{peft_config.task_type}_v1.pt".replace(
"/", "_"
)
# creating model
model = AutoModelForCausalLM.from_pretrained("LLM4Binary/llm4decompile-1.3b-v1.5", load_in_8bit=True, torch_dtype=torch.float16, device_map="auto")
model = prepare_model_for_int8_training(model)
peft_model = get_peft_model(model, peft_config)
training_args = TrainingArguments(
output_dir="./results3", # 保存模型的目录
evaluation_strategy="epoch", # 每个 epoch 进行评估
save_strategy="epoch", # 每个 epoch 结束时保存模型
learning_rate=2e-5,
per_device_train_batch_size=4, # 训练时的batch_size
per_device_eval_batch_size=8, # 验证时的batch_size
logging_steps=10, # log 打印的频率
num_train_epochs=3,
weight_decay=0.01,
load_best_model_at_end=False
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
#data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True)
)
trainer.train()
'''
trainer.evaluate(eval_dataset)
# 训练结束后手动保存模型
trainer.save_model(output_dir="./tuned_model") # 保存最终的模型到指定的目录
tokenizer.save_pretrained(save_directory="./tuned_tokenizer") # 保存tokenizer
'''
lora_adapter = "./lora_adapter"
peft_model.save_pretrained(lora_adapter, save_adapter=True, save_config=True)
model_to_merge = PeftModel.from_pretrained(AutoModelForCausalLM.from_pretrained(model_name_or_path).to("cuda"), lora_adapter)
merged_model = model_to_merge.merge_and_unload()
merged_model.save_pretrained(merged_model)