大模型安装好后 如何训练是重中之重。本篇我就给大家详细说下在服务器下的如何进行chatglm2-6b的模型训练。
1、首先安装好模型 参考我之前发的两篇文章
unbuntu下安装chatglm2-6b记(一)-CSDN博客
2、准备数据集
这里给各位老铁科普下
“大模型数据集”通常指用于训练大规模机器学习模型,特别是深度学习模型(如大型语言模型、计算机视觉模型等)的大量数据集合。随着模型规模的增加,训练这些模型所需的数据量也相应增加,因此需要使用非常庞大的数据集。这里数据集需要就两种常用的格式csv、json.
这里代码里我用的是json格式。
3、编写用于训练的代码
上干货 这里我使用lora进行训练
import os
import torch
from transformers import AutoTokenizer, AutoModel, AutoConfig, Trainer, TrainingArguments
from peft import get_peft_model, LoraConfig, TaskType, prepare_model_for_kbit_training, PeftModel
import logging
# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# 设置环境变量和设备
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 设置本地模型路径和卸载文件夹
LOCAL_MODEL_PATH = "/home/sykj/anaconda3/envs/ChatGLM2/ChatGLM2-6B/model"
OFFLOAD_FOLDER = "./offload_folder"
# 确保卸载文件夹存在
os.makedirs(OFFLOAD_FOLDER, exist_ok=True)
# 计算可用 GPU 内存的 60%
if torch.cuda.is_available():
total_memory = torch.cuda.get_device_properties(0).total_memory
available_memory = int(total_memory * 0.6)
available_memory_gb = f"{available_memory / (1024**3):.2f}GB"
logger.info(f"Setting max memory to {available_memory_gb}")
else:
available_memory_gb = None
logger.warning("CUDA is not available