已合并
add llm script glm4 and llama3s #37370
HandsoemLemon创建于 6月1日
add llm script glm4 and llama3s #37370
已合并
共 6 个文件变更+1078-0
| @@ -0,0 +1,106 @@ | |||
| 1 | +# glm4-9B-chat 微调训练 | ||
| 2 | + | ||
| 3 | +本 README 说明如何使用 **glm4-9B-chat** 模型权重，结合 **LlamaFactory** 提供的 `c4_demo.jsonl` 示例数据，完成数据下载、预处理与训练启动（含 eager / torch.compile 两种模式）。 | ||
| 4 | + | ||
| 5 | +--- | ||
| 6 | + | ||
| 7 | +## 目录 | ||
| 8 | + | ||
| 9 | +- [1. 模型权重](#1-模型权重) | ||
| 10 | +- [2. 数据获取](#2-数据获取) | ||
| 11 | +- [3. 数据预处理](#3-数据预处理) | ||
| 12 | +- [4. 模型训练](#4-模型训练) | ||
| 13 | + - [4.1 eager mode（默认）](#41-eager-mode默认) | ||
| 14 | + - [4.2 启用 torchcompile（可选）](#42-启用-torchcompile可选) | ||
| 15 | + - [4.3 采集profile文件（可选）](#43-采集profile文件可选) | ||
| 16 | + | ||
| 17 | +--- | ||
| 18 | + | ||
| 19 | +## 1. 模型权重 | ||
| 20 | + | ||
| 21 | +- Hugging Face 模型：**glm4-9B-chat** | ||
| 22 | + https://huggingface.co/zai-org/glm-4-9b-chat | ||
| 23 | + | ||
| 24 | +可在链接页面中 `Files and versions` 一栏直接下载。 | ||
| 25 | + | ||
| 26 | +### 环境提示 | ||
| 27 | + | ||
| 28 | +本项目依赖已整理到 `requirements.txt`，可直接安装： | ||
| 29 | + | ||
| 30 | +```bash | ||
| 31 | +pip install -r ../utils/requirements.txt | ||
| 32 | +``` | ||
| 33 | + | ||
| 34 | +## 2. 数据获取 | ||
| 35 | + | ||
| 36 | +训练数据集来自 LlamaFactory 仓库示例数据，可直接在如下连接获取： | ||
| 37 | + | ||
| 38 | +- `c4_demo.json` | ||
| 39 | + https://github.com/hiyouga/LlamaFactory/blob/main/data/c4_demo.jsonl | ||
| 40 | + | ||
| 41 | +或直接下载 raw 文件到本地：： | ||
| 42 | + | ||
| 43 | +```bash | ||
| 44 | +wget -O c4_demo.json \ | ||
| 45 | +https://github.com/hiyouga/LlamaFactory/blob/main/data/c4_demo.jsonl | ||
| 46 | +``` | ||
| 47 | + | ||
| 48 | +## 3. 数据预处理 | ||
| 49 | + | ||
| 50 | +在`train_glm4_9B.py`中已经对数据进行了处理 | ||
| 51 | + | ||
| 52 | +## 4. 模型训练 | ||
| 53 | + | ||
| 54 | +训练脚本：`train_glm4_9B.sh`，脚本支持在 GPU 和 NPU 训练 | ||
| 55 | + | ||
| 56 | +开始训练前，请修改脚本中的路径参数： | ||
| 57 | + | ||
| 58 | +- 模型权重路径（本地目录） | ||
| 59 | +- 训练数据路径（预处理后的 `train.json`） | ||
| 60 | + | ||
| 61 | +### 4.1 eager mode（默认） | ||
| 62 | + | ||
| 63 | +```bash | ||
| 64 | +bash train_glm4_9B.sh | ||
| 65 | +``` | ||
| 66 | + | ||
| 67 | +### 4.2 启用 `torch.compile`（可选） | ||
| 68 | + | ||
| 69 | +可通过添加 `--enable_compile` 选项运行图模式 | ||
| 70 | +当在 GPU 上训练时，默认后端使用triton；当在 NPU 上训练时，可进一步指定后端为 mlir 或 dvm，默认使用 mlir。 | ||
| 71 | + | ||
| 72 | +**默认后端（mlir，可不写 --npu-backend）：** | ||
| 73 | + | ||
| 74 | +```bash | ||
| 75 | +bash train_glm4_9B.sh \ | ||
| 76 | + --enable_compile | ||
| 77 | +``` | ||
| 78 | + | ||
| 79 | +**显式指定后端为 mlir：** | ||
| 80 | + | ||
| 81 | +```bash | ||
| 82 | +bash train_glm4_9B.sh \ | ||
| 83 | + --enable_compile \ | ||
| 84 | + --npu-backend mlir | ||
| 85 | +``` | ||
| 86 | + | ||
| 87 | +**切换后端为 dvm：：** | ||
| 88 | + | ||
| 89 | +```bash | ||
| 90 | +bash train_glm4_9B.sh \ | ||
| 91 | + --enable_compile \ | ||
| 92 | + --npu-backend dvm | ||
| 93 | +``` | ||
| 94 | + | ||
| 95 | +### 4.3 采集profile文件（可选） | ||
| 96 | + | ||
| 97 | +脚本已支持 `--enable_profiler` 这样的开关，开启方式为： | ||
| 98 | + | ||
| 99 | +```bash | ||
| 100 | +bash train_glm4_9B.sh \ | ||
| 101 | + --enable_profiler \ | ||
| 102 | + --profiler_start_step 5 \ | ||
| 103 | + --profiler_end_step 6 \ | ||
| 104 | +``` | ||
| 105 | + | ||
| 106 | +可以通过 `--profiler_start_step` 和 `--profiler_end_step` 分别设置profile开始和结束步数。结束步数。 | ||
| @@ -0,0 +1,20 @@ | |||
| 1 | +export TORCHINDUCTOR_CACHE_DIR='./cache' | ||
| 2 | +export ASCEND_RT_VISIBLE_DEVICES=0 | ||
| 3 | +export TORCH_COMPILE_DEBUG=1 | ||
| 4 | +export TORCH_NPU_USE_COMPATIBLE_IMPL=1 | ||
| 5 | + | ||
| 6 | +rm -rf ./cache/* | ||
| 7 | +mkdir -p ./cache logs | ||
| 8 | + | ||
| 9 | +python train_glm4_9B.py \ | ||
| 10 | + --model_path "your model path" \ | ||
| 11 | + --data_path "dataset path" \ | ||
| 12 | + --output_dir "./glm4-finetuned" \ | ||
| 13 | + --num_epochs 3 \ | ||
| 14 | + --max_steps 200 \ | ||
| 15 | + --batch_size 1 \ | ||
| 16 | + --learning_rate 2e-5 \ | ||
| 17 | + --max_length 512 \ | ||
| 18 | + --use_lora \ | ||
| 19 | + --use_bf16 \ | ||
| 20 | + --pad_to_max_length > logs/train_glm4.log 2>&1 | ||
| @@ -0,0 +1,424 @@ | |||
| 1 | +import torch | ||
| 2 | +import os | ||
| 3 | +import sys | ||
| 4 | +from pathlib import Path | ||
| 5 | +import torch.nn as nn | ||
| 6 | +from datasets import Dataset, load_dataset | ||
| 7 | +from transformers import ( | ||
| 8 | + AutoModelForCausalLM, | ||
| 9 | + AutoTokenizer, | ||
| 10 | + TrainingArguments, | ||
| 11 | + Trainer, | ||
| 12 | + DataCollatorForLanguageModeling, | ||
| 13 | + BitsAndBytesConfig, | ||
| 14 | +) | ||
| 15 | +from peft import ( | ||
| 16 | + LoraConfig, | ||
| 17 | + get_peft_model, | ||
| 18 | + prepare_model_for_kbit_training, | ||
| 19 | + TaskType | ||
| 20 | +) | ||
| 21 | +import json | ||
| 22 | +from typing import Dict, Union, Any, Optional | ||
| 23 | +import argparse | ||
| 24 | +import logging | ||
| 25 | + | ||
| 26 | + | ||
| 27 | +sys.path.append(str(Path(__file__).parent.parent)) | ||
| 28 | +from utils.utils import ( | ||
| 29 | + TimingCallback, | ||
| 30 | + get_profile, | ||
| 31 | + detect_device_type | ||
| 32 | +) | ||
| 33 | + | ||
| 34 | +logging.basicConfig(level=logging.INFO) | ||
| 35 | +logger = logging.getLogger(__name__) | ||
| 36 | + | ||
| 37 | +model_name = 'GLM4-9B' | ||
| 38 | + | ||
| 39 | +class CustomTrainer(Trainer): | ||
| 40 | + def compute_loss( | ||
| 41 | + self, | ||
| 42 | + model: nn.Module, | ||
| 43 | + inputs: Dict[str, Union[torch.Tensor, Any]], | ||
| 44 | + return_outputs: bool = False, | ||
| 45 | + num_items_in_batch: Optional[torch.Tensor] = None, | ||
| 46 | + ): | ||
| 47 | + return super().compute_loss( | ||
| 48 | + model=model, | ||
| 49 | + inputs=inputs, | ||
| 50 | + return_outputs=return_outputs, | ||
| 51 | + num_items_in_batch=None | ||
| 52 | + ) | ||
| 53 | + | ||
| 54 | +class GLM4Trainer: | ||
| 55 | + def __init__(self, args): | ||
| 56 | + self.args = args | ||
| 57 | + self.setup_training() | ||
| 58 | + | ||
| 59 | + def setup_training(self): | ||
| 60 | + logger.info(f"Loading GLM-4 model from {self.args.model_path}") | ||
| 61 | + | ||
| 62 | + bnb_config = None | ||
| 63 | + if self.args.use_4bit: | ||
| 64 | + bnb_config = BitsAndBytesConfig( | ||
| 65 | + load_in_4bit=True, | ||
| 66 | + bnb_4bit_quant_type="nf4", | ||
| 67 | + bnb_4bit_compute_dtype=torch.float16, | ||
| 68 | + bnb_4bit_use_double_quant=True | ||
| 69 | + ) | ||
| 70 | + | ||
| 71 | + self.model = AutoModelForCausalLM.from_pretrained( | ||
| 72 | + self.args.model_path, | ||
| 73 | + quantization_config=bnb_config if self.args.use_4bit else None, | ||
| 74 | + torch_dtype=torch.bfloat16 if self.args.use_bf16 else torch.float16, | ||
| 75 | + trust_remote_code=True, | ||
| 76 | + ) | ||
| 77 | + | ||
| 78 | + logger.info(f"Moving model to {self.args.device_type}...") | ||
| 79 | + self.model = self.model.to(self.args.device_type) | ||
| 80 | + | ||
| 81 | + self.tokenizer = AutoTokenizer.from_pretrained( | ||
| 82 | + self.args.model_path, | ||
| 83 | + trust_remote_code=True | ||
| 84 | + ) | ||
| 85 | + | ||
| 86 | + if self.tokenizer.pad_token is None: | ||
| 87 | + self.tokenizer.pad_token = self.tokenizer.eos_token | ||
| 88 | + self.model.config.pad_token_id = self.tokenizer.eos_token_id | ||
| 89 | + | ||
| 90 | + if self.args.gradient_checkpointing: | ||
| 91 | + self.model.gradient_checkpointing_enable() | ||
| 92 | + self.model.config.use_cache = False | ||
| 93 | + | ||
| 94 | + | ||
| 95 | + def apply_lora(self): | ||
| 96 | + if not self.args.use_lora: | ||
| 97 | + return | ||
| 98 | + | ||
| 99 | + logger.info("Applying LoRA configuration...") | ||
| 100 | + | ||
| 101 | + if self.args.use_4bit: | ||
| 102 | + self.model = prepare_model_for_kbit_training(self.model) | ||
| 103 | + | ||
| 104 | + lora_config = LoraConfig( | ||
| 105 | + task_type=TaskType.CAUSAL_LM, | ||
| 106 | + r=self.args.lora_r, | ||
| 107 | + lora_alpha=self.args.lora_alpha, | ||
| 108 | + lora_dropout=self.args.lora_dropout, | ||
| 109 | + target_modules=self.get_lora_target_modules(), | ||
| 110 | + bias="none", | ||
| 111 | + ) | ||
| 112 | + | ||
| 113 | + self.model = get_peft_model(self.model, lora_config) | ||
| 114 | + self.model.print_trainable_parameters() | ||
| 115 | + | ||
| 116 | + if self.args.enable_compile: | ||
| 117 | + self.model = torch.compile(self.model, dynamic=False) | ||
| 118 | + | ||
| 119 | + | ||
| 120 | + def get_lora_target_modules(self): | ||
| 121 | + target_modules = [ | ||
| 122 | + "query_key_value", | ||
| 123 | + "dense", | ||
| 124 | + "dense_h_to_4h", | ||
| 125 | + "dense_4h_to_h", | ||
| 126 | + ] | ||
| 127 | + | ||
| 128 | + model_modules = sorted([name for name, _ in self.model.named_modules()]) | ||
| 129 | + available_modules = [] | ||
| 130 | + | ||
| 131 | + for target in target_modules: | ||
| 132 | + found = False | ||
| 133 | + for name in model_modules: | ||
| 134 | + if target in name: | ||
| 135 | + available_modules.append(name) | ||
| 136 | + found = True | ||
| 137 | + break | ||
| 138 | + if not found: | ||
| 139 | + logger.warning(f"Target module {target} not found in model") | ||
| 140 | + | ||
| 141 | + available_modules = list(dict.fromkeys(available_modules)) | ||
| 142 | + | ||
| 143 | + if not available_modules: | ||
| 144 | + available_modules = ["qkv", "proj", "fc1", "fc2"] | ||
| 145 | + logger.warning(f"No avaliable, GLM-4 using LoRA target modules: {available_modules}") | ||
| 146 | + else: | ||
| 147 | + logger.info(f"GLM-4 LoRA target modules: {available_modules}") | ||
| 148 | + | ||
| 149 | + return available_modules | ||
| 150 | + | ||
| 151 | + | ||
| 152 | + def load_and_preprocess_data(self) -> Dataset: | ||
| 153 | + logger.info(f"Loading dataset from {self.args.data_path}") | ||
| 154 | + | ||
| 155 | + if self.args.data_path.endswith('.json') or self.args.data_path.endswith('.jsonl'): | ||
| 156 | + with open(self.args.data_path, 'r', encoding='utf-8') as f: | ||
| 157 | + if self.args.data_path.endswith('.jsonl'): | ||
| 158 | + data = [json.loads(line) for line in f] | ||
| 159 | + else: | ||
| 160 | + data = json.load(f) | ||
| 161 | + | ||
| 162 | + formatted_data = [] | ||
| 163 | + for item in data: | ||
| 164 | + if "conversations" in item or "messages" in item: | ||
| 165 | + messages = item.get("messages", item.get("conversations", [])) | ||
| 166 | + text = self.tokenizer.apply_chat_template( | ||
| 167 | + messages, | ||
| 168 | + tokenize=False, | ||
| 169 | + add_generation_prompt=False | ||
| 170 | + ) | ||
| 171 | + else: | ||
| 172 | + text = item.get("text", "") | ||
| 173 | + | ||
| 174 | + formatted_data.append({"text": text}) | ||
| 175 | + | ||
| 176 | + dataset = Dataset.from_list(formatted_data) | ||
| 177 | + else: | ||
| 178 | + try: | ||
| 179 | + dataset = load_dataset( | ||
| 180 | + self.args.data_path, | ||
| 181 | + split=self.args.split | ||
| 182 | + ) | ||
| 183 | + except: | ||
| 184 | + dataset = load_dataset( | ||
| 185 | + "json", | ||
| 186 | + data_files=self.args.data_path, | ||
| 187 | + split="train" | ||
| 188 | + ) | ||
| 189 | + | ||
| 190 | + tokenizer_ref = self.tokenizer | ||
| 191 | + max_length_ref = self.args.max_length | ||
| 192 | + | ||
| 193 | + def preprocess_function(examples): | ||
| 194 | + tokenized = tokenizer_ref( | ||
| 195 | + examples["text"], | ||
| 196 | + truncation=True, | ||
| 197 | + padding="max_length" if self.args.pad_to_max_length else False, | ||
| 198 | + max_length=max_length_ref, | ||
| 199 | + return_tensors=None, | ||
| 200 | + return_attention_mask=True, | ||
| 201 | + ) | ||
| 202 | + import copy | ||
| 203 | + tokenized["labels"] = copy.deepcopy(tokenized["input_ids"]) | ||
| 204 | + return tokenized | ||
| 205 | + | ||
| 206 | + | ||
| 207 | + num_proc = self.args.num_proc if self.args.num_proc > 0 else None | ||
| 208 | + processed_dataset = dataset.map( | ||
| 209 | + preprocess_function, | ||
| 210 | + batched=True, | ||
| 211 | + remove_columns=dataset.column_names, | ||
| 212 | + num_proc=num_proc, | ||
| 213 | + load_from_cache_file=not self.args.overwrite_cache | ||
| 214 | + ) | ||
| 215 | + logger.info(f"Dataset size: {len(processed_dataset)}") | ||
| 216 | + return processed_dataset | ||
| 217 | + | ||
| 218 | + | ||
| 219 | + def create_trainer(self, train_dataset, eval_dataset=None): | ||
| 220 | + has_eval = eval_dataset is not None | ||
| 221 | + eval_strategy = "steps" if has_eval else "no" | ||
| 222 | + save_strategy = "steps" | ||
| 223 | + | ||
| 224 | + training_args = TrainingArguments( | ||
| 225 | + output_dir=self.args.output_dir, | ||
| 226 | + overwrite_output_dir=True, | ||
| 227 | + num_train_epochs=self.args.num_epochs, | ||
| 228 | + max_steps=self.args.max_steps if self.args.max_steps > 0 else 30, | ||
| 229 | + per_device_train_batch_size=self.args.batch_size, | ||
| 230 | + per_device_eval_batch_size=self.args.batch_size, | ||
| 231 | + gradient_accumulation_steps=self.args.gradient_accumulation_steps, | ||
| 232 | + weight_decay=self.args.weight_decay, | ||
| 233 | + logging_dir=f"{self.args.output_dir}/logs", | ||
| 234 | + logging_steps=self.args.logging_steps, | ||
| 235 | + save_steps=self.args.save_steps, | ||
| 236 | + save_total_limit=self.args.save_total_limit, | ||
| 237 | + eval_strategy=eval_strategy, | ||
| 238 | + eval_steps=self.args.eval_steps if has_eval else None, | ||
| 239 | + save_strategy=save_strategy, | ||
| 240 | + load_best_model_at_end=has_eval, | ||
| 241 | + metric_for_best_model="eval_loss", | ||
| 242 | + greater_is_better=False, | ||
| 243 | + learning_rate=self.args.learning_rate, | ||
| 244 | + lr_scheduler_type=self.args.lr_scheduler_type, | ||
| 245 | + fp16=self.args.use_fp16, | ||
| 246 | + bf16=self.args.use_bf16, | ||
| 247 | + gradient_checkpointing=self.args.gradient_checkpointing, | ||
| 248 | + report_to="none", | ||
| 249 | + ddp_find_unused_parameters=False if torch.cuda.device_count() > 1 else None, | ||
| 250 | + remove_unused_columns=False, | ||
| 251 | + dataloader_num_workers=self.args.dataloader_num_workers, | ||
| 252 | + group_by_length=self.args.group_by_length, | ||
| 253 | + length_column_name="length", | ||
| 254 | + prediction_loss_only=True, | ||
| 255 | + ) | ||
| 256 | + | ||
| 257 | + data_collator = DataCollatorForLanguageModeling( | ||
| 258 | + tokenizer=self.tokenizer, | ||
| 259 | + mlm=False, | ||
| 260 | + ) | ||
| 261 | + | ||
| 262 | + mod='compile' if self.args.enable_compile else 'eager' | ||
| 263 | + prof=None | ||
| 264 | + | ||
| 265 | + if self.args.enable_profiler: | ||
| 266 | + profiling_save_path = self.args.profiler_save_path + '/' + model_name + '/' + mod | ||
| 267 | + prof = get_profile(self.args.profiler_start_step, self.args.profiler_end_step, profiling_save_path) | ||
| 268 | + | ||
| 269 | + timing_callback = TimingCallback(prof, mod) | ||
| 270 | + | ||
| 271 | + trainer = CustomTrainer( | ||
| 272 | + model=self.model, | ||
| 273 | + args=training_args, | ||
| 274 | + train_dataset=train_dataset, | ||
| 275 | + eval_dataset=eval_dataset, | ||
| 276 | + tokenizer=self.tokenizer, | ||
| 277 | + data_collator=data_collator, | ||
| 278 | + callbacks=[timing_callback], | ||
| 279 | + ) | ||
| 280 | + return trainer | ||
| 281 | + | ||
| 282 | + | ||
| 283 | + def train(self): | ||
| 284 | + logger.info("Starting GLM-4 training...") | ||
| 285 | + | ||
| 286 | + if self.args.use_lora: | ||
| 287 | + self.apply_lora() | ||
| 288 | + | ||
| 289 | + dataset = self.load_and_preprocess_data() | ||
| 290 | + | ||
| 291 | + if self.args.validation_split > 0: | ||
| 292 | + split_dataset = dataset.train_test_split( | ||
| 293 | + test_size=self.args.validation_split, | ||
| 294 | + seed=self.args.seed | ||
| 295 | + ) | ||
| 296 | + train_dataset = split_dataset["train"] | ||
| 297 | + eval_dataset = split_dataset["test"] | ||
| 298 | + else: | ||
| 299 | + train_dataset = dataset | ||
| 300 | + eval_dataset = None | ||
| 301 | + | ||
| 302 | + trainer = self.create_trainer(train_dataset, eval_dataset) | ||
| 303 | + train_result = trainer.train() | ||
| 304 | + if self.args.enable_compile: | ||
| 305 | + headers, values = torch._dynamo.utils.compile_times("csv") | ||
| 306 | + for header, value in zip(headers, values): | ||
| 307 | + if header == "PyCodeCache.load_by_key_path": | ||
| 308 | + numbers = [float(num.strip()) for num in value.split(',') if num.strip()] | ||
| 309 | + op_compile_time = sum(numbers) | ||
| 310 | + print(f"op_compile_time:{op_compile_time * 1e3} ms", ) | ||
| 311 | + | ||
| 312 | + trainer.save_model() | ||
| 313 | + self.tokenizer.save_pretrained(self.args.output_dir) | ||
| 314 | + metrics = train_result.metrics | ||
| 315 | + trainer.log_metrics("train", metrics) | ||
| 316 | + trainer.save_metrics("train", metrics) | ||
| 317 | + trainer.save_state() | ||
| 318 | + logger.info(f"Training completed! Model saved to {self.args.output_dir}") | ||
| 319 | + return metrics | ||
| 320 | + | ||
| 321 | + | ||
| 322 | +def main(): | ||
| 323 | + parser = argparse.ArgumentParser(description="Train GLM-4 model") | ||
| 324 | + parser.add_argument("--model_path", type=str, default="ZhipuAI/glm-4-9b-chat", | ||
| 325 | + help="Path to the pretrained GLM-4 model") | ||
| 326 | + parser.add_argument("--data_path", type=str, required=True, | ||
| 327 | + help="Path to training data (json/jsonl file or dataset name)") | ||
| 328 | + parser.add_argument("--output_dir", type=str, default="./glm4-finetuned", | ||
| 329 | + help="Output directory for trained model") | ||
| 330 | + parser.add_argument("--num_epochs", type=int, default=3, | ||
| 331 | + help="Number of training epochs") | ||
| 332 | + parser.add_argument("--batch_size", type=int, default=2, | ||
| 333 | + help="Batch size per device") | ||
| 334 | + parser.add_argument("--gradient_accumulation_steps", type=int, default=1, | ||
| 335 | + help="Gradient accumulation steps") | ||
| 336 | + parser.add_argument("--learning_rate", type=float, default=1e-4, | ||
| 337 | + help="Learning rate") | ||
| 338 | + parser.add_argument("--warmup_steps", type=int, default=50, | ||
| 339 | + help="Warmup steps") | ||
| 340 | + parser.add_argument("--max_steps", type=int, default=-1, | ||
| 341 | + help="Total training steps ") | ||
| 342 | + parser.add_argument("--weight_decay", type=float, default=0.01, | ||
| 343 | + help="Weight decay") | ||
| 344 | + parser.add_argument("--max_length", type=int, default=2048, | ||
| 345 | + help="Maximum sequence length") | ||
| 346 | + parser.add_argument("--lr_scheduler_type", type=str, default="cosine", | ||
| 347 | + choices=["linear", "cosine", "cosine_with_restarts", "constant"], | ||
| 348 | + help="Learning rate scheduler type") | ||
| 349 | + parser.add_argument("--use_lora", action="store_true", | ||
| 350 | + help="Use LoRA for parameter-efficient fine-tuning") | ||
| 351 | + parser.add_argument("--lora_r", type=int, default=8, | ||
| 352 | + help="LoRA rank") | ||
| 353 | + parser.add_argument("--lora_alpha", type=int, default=32, | ||
| 354 | + help="LoRA alpha") | ||
| 355 | + parser.add_argument("--lora_dropout", type=float, default=0.1, | ||
| 356 | + help="LoRA dropout") | ||
| 357 | + parser.add_argument("--use_4bit", action="store_true", | ||
| 358 | + help="Use 4-bit quantization") | ||
| 359 | + parser.add_argument("--use_fp16", action="store_true", | ||
| 360 | + help="Use FP16 precision") | ||
| 361 | + parser.add_argument("--use_bf16", action="store_true", | ||
| 362 | + help="Use BF16 precision") | ||
| 363 | + parser.add_argument("--gradient_checkpointing", action="store_true", | ||
| 364 | + help="Enable gradient checkpointing") | ||
| 365 | + parser.add_argument("--validation_split", type=float, default=0.1, | ||
| 366 | + help="Validation split ratio") | ||
| 367 | + parser.add_argument("--split", type=str, default="train", | ||
| 368 | + help="Dataset split to use") | ||
| 369 | + parser.add_argument("--num_proc", type=int, default=0, | ||
| 370 | + help="Number of processes for data preprocessing (0 = single process)") | ||
| 371 | + parser.add_argument("--pad_to_max_length", action="store_true", | ||
| 372 | + help="Pad sequences to max_length") | ||
| 373 | + parser.add_argument("--overwrite_cache", action="store_true", | ||
| 374 | + help="Overwrite cached features") | ||
| 375 | + parser.add_argument("--group_by_length", action="store_true", | ||
| 376 | + help="Group sequences by length for efficient training") | ||
| 377 | + parser.add_argument("--dataloader_num_workers", type=int, default=4, | ||
| 378 | + help="Number of workers for data loading") | ||
| 379 | + parser.add_argument("--seed", type=int, default=42, | ||
| 380 | + help="Random seed") | ||
| 381 | + parser.add_argument("--logging_steps", type=int, default=1, | ||
| 382 | + help="Log every X updates steps") | ||
| 383 | + parser.add_argument("--save_steps", type=int, default=500, | ||
| 384 | + help="Save checkpoint every X updates steps") | ||
| 385 | + parser.add_argument("--eval_steps", type=int, default=500, | ||
| 386 | + help="Evaluate every X updates steps") | ||
| 387 | + parser.add_argument("--save_total_limit", type=int, default=3, | ||
| 388 | + help="Limit the total amount of checkpoints") | ||
| 389 | + parser.add_argument("--report_to_tensorboard", action="store_true", | ||
| 390 | + help="Report metrics to TensorBoard") | ||
| 391 | + parser.add_argument("--enable_compile", action="store_true", | ||
| 392 | + help="Enable torch.compile and Inductor backend") | ||
| 393 | + parser.add_argument("--enable_profiler", action="store_true", | ||
| 394 | + help="Enable profiler for performance analysis") | ||
| 395 | + parser.add_argument("--profiler_start_step", type=int, default=5, | ||
| 396 | + help="Output directory for trained model") | ||
| 397 | + parser.add_argument("--profiler_end_step", type=int, default=8, | ||
| 398 | + help="Output directory for trained model") | ||
| 399 | + parser.add_argument("--profiler_save_path", type=str, default="./profile", | ||
| 400 | + help="Output directory for trained model") | ||
| 401 | + parser.add_argument("--npu-backend", type=str, default="mlir") | ||
| 402 | + parser.add_argument("--mfusion", action="store_true", help="Enable MFusion for graph fusion optimization") | ||
| 403 | + args = parser.parse_args() | ||
| 404 | + torch.manual_seed(args.seed) | ||
| 405 | + args.device_type = detect_device_type() | ||
| 406 | + os.environ['TORCHINDUCTOR_NPU_BACKEND']=args.npu_backend | ||
| 407 | + if args.npu_backend == "akg": | ||
| 408 | + os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir' | ||
| 409 | + os.environ['TORCHINDUCTOR_USE_AKG'] = '1' | ||
| 410 | + if args.mfusion: | ||
| 411 | + os.environ['TORCHINDUCTOR_ENABLE_MFUSION']='1' | ||
| 412 | + print(f"{args.device_type} train {model_name}") | ||
| 413 | + | ||
| 414 | + trainer = GLM4Trainer(args) | ||
| 415 | + metrics = trainer.train() | ||
| 416 | + | ||
| 417 | + print("\n" + "="*50) | ||
| 418 | + print("GLM-4 Training completed successfully!") | ||
| 419 | + print(f"Model saved to: {args.output_dir}") | ||
| 420 | + print(f"Final training loss: {metrics.get('train_loss', 'N/A')}") | ||
| 421 | + print("="*50) | ||
| 422 | + | ||
| 423 | +if __name__ == "__main__": | ||
| 424 | + main() | ||
| @@ -0,0 +1,120 @@ | |||
| 1 | +# llama3&llama3.2 微调训练 | ||
| 2 | + | ||
| 3 | +本 README 说明如何使用 **llama3-8B & llama3.2-3B** 模型权重，结合 **LlamaFactory** 提供的 `c4_demo.jsonl` 示例数据，完成数据下载、预处理与训练启动（含 eager / torch.compile 两种模式）。 | ||
| 4 | + | ||
| 5 | +--- | ||
| 6 | + | ||
| 7 | +## 目录 | ||
| 8 | + | ||
| 9 | +- [1. 模型权重](#1-模型权重) | ||
| 10 | +- [2. 数据获取](#2-数据获取) | ||
| 11 | +- [3. 数据预处理](#3-数据预处理) | ||
| 12 | +- [4. 模型训练](#4-模型训练) | ||
| 13 | + - [4.1 eager mode（默认）](#41-eager-mode默认) | ||
| 14 | + - [4.2 启用 torchcompile（可选）](#42-启用-torchcompile可选) | ||
| 15 | + - [4.3 采集profile文件（可选）](#43-采集profile文件可选) | ||
| 16 | + | ||
| 17 | +--- | ||
| 18 | + | ||
| 19 | +## 1. 模型权重 | ||
| 20 | + | ||
| 21 | +- ModelScope 模型：**llama3-8B** | ||
| 22 | + https://modelscope.cn/models/LLM-Research/Meta-Llama-3-8B | ||
| 23 | + | ||
| 24 | +- ModelScope 模型：**llama3.2-3B** | ||
| 25 | + https://modelscope.cn/models/LLM-Research/Llama-3.2-3B | ||
| 26 | + | ||
| 27 | +可在链接页面中模型文件一栏直接下载模型，或使用ModelScope官方推荐SDK下载方式： | ||
| 28 | + | ||
| 29 | +```bash | ||
| 30 | +#llama3-8B | ||
| 31 | +from modelscope import snapshot_download | ||
| 32 | +model_dir = snapshot_download('LLM-Research/Meta-Llama-3-8B') | ||
| 33 | + | ||
| 34 | + | ||
| 35 | +#llama3.2-3B | ||
| 36 | +from modelscope import snapshot_download | ||
| 37 | +model_dir = snapshot_download('LLM-Research/Llama-3.2-3B') | ||
| 38 | +``` | ||
| 39 | + | ||
| 40 | +### 环境提示 | ||
| 41 | + | ||
| 42 | +本项目依赖已整理到 `requirements.txt`，可直接安装： | ||
| 43 | + | ||
| 44 | +```bash | ||
| 45 | +pip install -r ../utils/requirements.txt | ||
| 46 | +``` | ||
| 47 | + | ||
| 48 | +## 2. 数据获取 | ||
| 49 | + | ||
| 50 | +训练数据集来自 LlamaFactory 仓库示例数据，可直接在如下连接获取： | ||
| 51 | + | ||
| 52 | +- `c4_demo.json` | ||
| 53 | + https://github.com/hiyouga/LlamaFactory/blob/main/data/c4_demo.jsonl | ||
| 54 | + | ||
| 55 | +或直接下载 raw 文件到本地： | ||
| 56 | + | ||
| 57 | +```bash | ||
| 58 | +wget -O c4_demo.json \ | ||
| 59 | +https://github.com/hiyouga/LlamaFactory/blob/main/data/c4_demo.jsonl | ||
| 60 | +``` | ||
| 61 | + | ||
| 62 | +## 3. 数据预处理 | ||
| 63 | + | ||
| 64 | +在`train_llama3.py`中已经对数据进行了处理, 本脚本支持对 `llama3-8B` 和`Llama-3.2-3B`预训练，在训练时通过`--model_path`修改模型权重路径即可 | ||
| 65 | + | ||
| 66 | +## 4. 模型训练 | ||
| 67 | + | ||
| 68 | +训练脚本：`run_llama3.sh`，脚本支持在 GPU 和 NPU 训练 | ||
| 69 | + | ||
| 70 | +开始训练前，请修改脚本中的路径参数： | ||
| 71 | + | ||
| 72 | +- 模型权重路径 | ||
| 73 | +- 训练数据路径 | ||
| 74 | + | ||
| 75 | +### 4.1 eager mode（默认） | ||
| 76 | + | ||
| 77 | +```bash | ||
| 78 | +bash run_llama3.sh | ||
| 79 | +``` | ||
| 80 | + | ||
| 81 | +### 4.2 启用 `torch.compile`（可选） | ||
| 82 | + | ||
| 83 | +可通过添加 `--enable_compile` 选项运行图模式 | ||
| 84 | +当在 GPU 上训练时，默认后端使用triton；当在 NPU 上训练时，可进一步指定后端为 mlir 或 dvm，默认使用 mlir。 | ||
| 85 | + | ||
| 86 | +**默认后端（mlir，可不写 --npu-backend）：** | ||
| 87 | + | ||
| 88 | +```bash | ||
| 89 | +bash run_llama3.sh \ | ||
| 90 | + --enable_compile | ||
| 91 | +``` | ||
| 92 | + | ||
| 93 | +**显式指定后端为 mlir：** | ||
| 94 | + | ||
| 95 | +```bash | ||
| 96 | +bash run_llama3.sh \ | ||
| 97 | + --enable_compile \ | ||
| 98 | + --npu-backend mlir | ||
| 99 | +``` | ||
| 100 | + | ||
| 101 | +**切换后端为 dvm：：** | ||
| 102 | + | ||
| 103 | +```bash | ||
| 104 | +bash run_llama3.sh \ | ||
| 105 | + --enable_compile \ | ||
| 106 | + --npu-backend dvm | ||
| 107 | +``` | ||
| 108 | + | ||
| 109 | +### 4.3 采集profile文件（可选） | ||
| 110 | + | ||
| 111 | +脚本已支持 `--enable_profiler` 这样的开关，开启方式为： | ||
| 112 | + | ||
| 113 | +```bash | ||
| 114 | +bash run_llama3.sh \ | ||
| 115 | + --enable_profiler \ | ||
| 116 | + --profiler_start_step 5 \ | ||
| 117 | + --profiler_end_step 6 \ | ||
| 118 | +``` | ||
| 119 | + | ||
| 120 | +可以通过 `--profiler_start_step` 和 `--profiler_end_step` 分别设置profile开始和结束步数。 | ||
| @@ -0,0 +1,21 @@ | |||
| 1 | +export TORCHINDUCTOR_CACHE_DIR='./cache' | ||
| 2 | +export ASCEND_RT_VISIBLE_DEVICES=3 | ||
| 3 | +export TORCH_COMPILE_DEBUG=1 | ||
| 4 | +export TORCH_NPU_USE_COMPATIBLE_IMPL=1 | ||
| 5 | + | ||
| 6 | +rm -rf ./cache/* | ||
| 7 | +mkdir -p ./cache logs | ||
| 8 | + | ||
| 9 | +python train_llama3.py \ | ||
| 10 | + --model_path "your model path" \ | ||
| 11 | + --model_name "your model name" \ | ||
| 12 | + --data_path "dataset path" \ | ||
| 13 | + --output_dir "./llama3-finetuned" \ | ||
| 14 | + --num_epochs 3 \ | ||
| 15 | + --max_steps 200 \ | ||
| 16 | + --batch_size 1 \ | ||
| 17 | + --learning_rate 2e-5 \ | ||
| 18 | + --max_length 512 \ | ||
| 19 | + --use_lora \ | ||
| 20 | + --use_bf16 \ | ||
| 21 | + --pad_to_max_length > logs/train_llama3.log 2>&1 | ||
| @@ -0,0 +1,387 @@ | |||
| 1 | +import torch | ||
| 2 | +import sys | ||
| 3 | +import os | ||
| 4 | +from pathlib import Path | ||
| 5 | +import torch.nn as nn | ||
| 6 | +from datasets import Dataset, load_dataset | ||
| 7 | +from transformers import ( | ||
| 8 | + AutoModelForCausalLM, | ||
| 9 | + AutoTokenizer, | ||
| 10 | + TrainingArguments, | ||
| 11 | + Trainer, | ||
| 12 | + DataCollatorForLanguageModeling, | ||
| 13 | + get_linear_schedule_with_warmup, | ||
| 14 | + BitsAndBytesConfig, | ||
| 15 | + TrainerCallback | ||
| 16 | +) | ||
| 17 | +from peft import ( | ||
| 18 | + LoraConfig, | ||
| 19 | + get_peft_model, | ||
| 20 | + prepare_model_for_kbit_training, | ||
| 21 | + TaskType | ||
| 22 | +) | ||
| 23 | +import json | ||
| 24 | +from typing import Dict, List, Optional | ||
| 25 | +import argparse | ||
| 26 | +import logging | ||
| 27 | + | ||
| 28 | +sys.path.append(str(Path(__file__).parent.parent)) | ||
| 29 | +from utils.utils import ( | ||
| 30 | + TimingCallback, | ||
| 31 | + get_profile, | ||
| 32 | + detect_device_type | ||
| 33 | +) | ||
| 34 | + | ||
| 35 | +logging.basicConfig(level=logging.INFO) | ||
| 36 | +logger = logging.getLogger(__name__) | ||
| 37 | + | ||
| 38 | +class LLama3Trainer: | ||
| 39 | + def __init__(self, args): | ||
| 40 | + self.args = args | ||
| 41 | + self.setup_training() | ||
| 42 | + | ||
| 43 | + def setup_training(self): | ||
| 44 | + logger.info(f"Loading model from {self.args.model_path}") | ||
| 45 | + | ||
| 46 | + bnb_config = None | ||
| 47 | + if self.args.use_4bit: | ||
| 48 | + bnb_config = BitsAndBytesConfig( | ||
| 49 | + load_in_4bit=True, | ||
| 50 | + bnb_4bit_quant_type="nf4", | ||
| 51 | + bnb_4bit_compute_dtype=torch.float16, | ||
| 52 | + bnb_4bit_use_double_quant=True | ||
| 53 | + ) | ||
| 54 | + | ||
| 55 | + self.model = AutoModelForCausalLM.from_pretrained( | ||
| 56 | + self.args.model_path, | ||
| 57 | + quantization_config=bnb_config if self.args.use_4bit else None, | ||
| 58 | + torch_dtype=torch.bfloat16 if self.args.use_bf16 else torch.float32, | ||
| 59 | + trust_remote_code=True | ||
| 60 | + ) | ||
| 61 | + | ||
| 62 | + logger.info(f"Moving model to {self.args.device_type}...") | ||
| 63 | + self.model = self.model.to(self.args.device_type) | ||
| 64 | + self.tokenizer = AutoTokenizer.from_pretrained( | ||
| 65 | + self.args.model_path, | ||
| 66 | + trust_remote_code=True | ||
| 67 | + ) | ||
| 68 | + | ||
| 69 | + if self.tokenizer.pad_token is None: | ||
| 70 | + self.tokenizer.pad_token = self.tokenizer.eos_token | ||
| 71 | + self.model.config.pad_token_id = self.tokenizer.eos_token_id | ||
| 72 | + | ||
| 73 | + if self.args.gradient_checkpointing: | ||
| 74 | + self.model.gradient_checkpointing_enable() | ||
| 75 | + self.model.config.use_cache = False | ||
| 76 | + | ||
| 77 | + | ||
| 78 | + def apply_lora(self): | ||
| 79 | + if not self.args.use_lora: | ||
| 80 | + return | ||
| 81 | + | ||
| 82 | + logger.info("Applying LoRA configuration...") | ||
| 83 | + if self.args.use_4bit: | ||
| 84 | + self.model = prepare_model_for_kbit_training(self.model) | ||
| 85 | + | ||
| 86 | + lora_config = LoraConfig( | ||
| 87 | + task_type=TaskType.CAUSAL_LM, | ||
| 88 | + r=self.args.lora_r, | ||
| 89 | + lora_alpha=self.args.lora_alpha, | ||
| 90 | + lora_dropout=self.args.lora_dropout, | ||
| 91 | + target_modules=self.get_lora_target_modules(), | ||
| 92 | + bias="none", | ||
| 93 | + ) | ||
| 94 | + | ||
| 95 | + self.model = get_peft_model(self.model, lora_config) | ||
| 96 | + self.model.print_trainable_parameters() | ||
| 97 | + | ||
| 98 | + if self.args.enable_compile: | ||
| 99 | + self.model = torch.compile(self.model, dynamic=False) | ||
| 100 | + | ||
| 101 | + def get_lora_target_modules(self): | ||
| 102 | + target_modules = [ | ||
| 103 | + "q_proj", "k_proj", "v_proj", "o_proj", | ||
| 104 | + "gate_proj", "up_proj", "down_proj", | ||
| 105 | + ] | ||
| 106 | + | ||
| 107 | + model_modules = set([name for name, _ in self.model.named_modules()]) | ||
| 108 | + available_modules = [m for m in target_modules if any(m in name for name in model_modules)] | ||
| 109 | + | ||
| 110 | + if not available_modules: | ||
| 111 | + available_modules = ["qkv_proj", "dense", "fc1", "fc2"] | ||
| 112 | + | ||
| 113 | + logger.info(f"Using LoRA target modules: {available_modules}") | ||
| 114 | + return available_modules | ||
| 115 | + | ||
| 116 | + def load_and_preprocess_data(self) -> Dataset: | ||
| 117 | + logger.info(f"Loading dataset from {self.args.data_path}") | ||
| 118 | + | ||
| 119 | + if self.args.data_path.endswith('.json') or self.args.data_path.endswith('.jsonl'): | ||
| 120 | + with open(self.args.data_path, 'r', encoding='utf-8') as f: | ||
| 121 | + if self.args.data_path.endswith('.jsonl'): | ||
| 122 | + data = [json.loads(line) for line in f] | ||
| 123 | + else: | ||
| 124 | + data = json.load(f) | ||
| 125 | + | ||
| 126 | + formatted_data = [] | ||
| 127 | + for item in data: | ||
| 128 | + if "conversations" in item: | ||
| 129 | + conversations = item["conversations"] | ||
| 130 | + text = self.format_conversation(conversations) | ||
| 131 | + elif "messages" in item: | ||
| 132 | + messages = item["messages"] | ||
| 133 | + text = self.tokenizer.apply_chat_template( | ||
| 134 | + messages, | ||
| 135 | + tokenize=False, | ||
| 136 | + add_generation_prompt=False | ||
| 137 | + ) | ||
| 138 | + else: | ||
| 139 | + text = item.get("text", "") | ||
| 140 | + | ||
| 141 | + formatted_data.append({"text": text}) | ||
| 142 | + | ||
| 143 | + dataset = Dataset.from_list(formatted_data) | ||
| 144 | + else: | ||
| 145 | + try: | ||
| 146 | + dataset = load_dataset( | ||
| 147 | + self.args.data_path, | ||
| 148 | + split=self.args.split | ||
| 149 | + ) | ||
| 150 | + except: | ||
| 151 | + dataset = load_dataset( | ||
| 152 | + "json", | ||
| 153 | + data_files=self.args.data_path, | ||
| 154 | + split="train" | ||
| 155 | + ) | ||
| 156 | + | ||
| 157 | + def preprocess_function(examples, tokenizer=self.tokenizer, max_length=self.args.max_length): | ||
| 158 | + tokenized = tokenizer( | ||
| 159 | + examples["text"], | ||
| 160 | + truncation=True, | ||
| 161 | + padding="max_length", | ||
| 162 | + max_length=max_length, | ||
| 163 | + return_tensors=None, | ||
| 164 | + return_attention_mask=True, | ||
| 165 | + ) | ||
| 166 | + import copy | ||
| 167 | + tokenized["labels"] = copy.deepcopy(tokenized["input_ids"]) | ||
| 168 | + return tokenized | ||
| 169 | + | ||
| 170 | + num_proc = self.args.num_proc if self.args.num_proc > 0 else None | ||
| 171 | + processed_dataset = dataset.map( | ||
| 172 | + preprocess_function, | ||
| 173 | + batched=True, | ||
| 174 | + remove_columns=dataset.column_names, | ||
| 175 | + num_proc=num_proc, | ||
| 176 | + load_from_cache_file=not self.args.overwrite_cache | ||
| 177 | + ) | ||
| 178 | + | ||
| 179 | + logger.info(f"Dataset size: {len(processed_dataset)}") | ||
| 180 | + return processed_dataset | ||
| 181 | + | ||
| 182 | + def format_conversation(self, conversations: List[Dict]) -> str: | ||
| 183 | + formatted_text = "" | ||
| 184 | + | ||
| 185 | + for turn in conversations: | ||
| 186 | + role = turn.get("from", "").lower() | ||
| 187 | + content = turn.get("value", "") | ||
| 188 | + if role == "human" or role == "user": | ||
| 189 | + formatted_text += f"<|im_start|>user\n{content}<|im_end|>\n" | ||
| 190 | + elif role == "gpt" or role == "assistant": | ||
| 191 | + formatted_text += f"<|im_start|>assistant\n{content}<|im_end|>\n" | ||
| 192 | + elif role == "system": | ||
| 193 | + formatted_text += f"<|im_start|>system\n{content}<|im_end|>\n" | ||
| 194 | + | ||
| 195 | + return formatted_text | ||
| 196 | + | ||
| 197 | + def create_trainer(self, train_dataset, eval_dataset=None): | ||
| 198 | + has_eval = eval_dataset is not None | ||
| 199 | + eval_strategy = "steps" if has_eval else "no" | ||
| 200 | + save_strategy = "steps" if has_eval else "steps" | ||
| 201 | + | ||
| 202 | + training_args = TrainingArguments( | ||
| 203 | + output_dir=self.args.output_dir, | ||
| 204 | + overwrite_output_dir=True, | ||
| 205 | + num_train_epochs=self.args.num_epochs, | ||
| 206 | + max_steps=self.args.max_steps, | ||
| 207 | + per_device_train_batch_size=self.args.batch_size, | ||
| 208 | + per_device_eval_batch_size=self.args.batch_size, | ||
| 209 | + gradient_accumulation_steps=self.args.gradient_accumulation_steps, | ||
| 210 | + weight_decay=self.args.weight_decay, | ||
| 211 | + logging_dir=f"{self.args.output_dir}/logs", | ||
| 212 | + logging_steps=self.args.logging_steps, | ||
| 213 | + save_steps=self.args.save_steps, | ||
| 214 | + save_total_limit=self.args.save_total_limit, | ||
| 215 | + eval_strategy=eval_strategy, | ||
| 216 | + eval_steps=self.args.eval_steps if has_eval else None, | ||
| 217 | + save_strategy=save_strategy, | ||
| 218 | + load_best_model_at_end=has_eval, | ||
| 219 | + metric_for_best_model="eval_loss", | ||
| 220 | + greater_is_better=False, | ||
| 221 | + learning_rate=self.args.learning_rate, | ||
| 222 | + lr_scheduler_type="cosine", | ||
| 223 | + fp16=self.args.use_fp16, | ||
| 224 | + bf16=self.args.use_bf16, | ||
| 225 | + gradient_checkpointing=self.args.gradient_checkpointing, | ||
| 226 | + report_to="none", | ||
| 227 | + ddp_find_unused_parameters=False if torch.cuda.device_count() > 1 else None, | ||
| 228 | + remove_unused_columns=False, | ||
| 229 | + ) | ||
| 230 | + data_collator = DataCollatorForLanguageModeling( | ||
| 231 | + tokenizer=self.tokenizer, | ||
| 232 | + mlm=False, | ||
| 233 | + ) | ||
| 234 | + mod='compile' if self.args.enable_compile else 'eager' | ||
| 235 | + prof=None | ||
| 236 | + | ||
| 237 | + if self.args.enable_profiler: | ||
| 238 | + profiling_save_path = self.args.profiler_save_path + '/' + self.args.model_name + '/' + mod | ||
| 239 | + prof = get_profile(self.args.profiler_start_step, self.args.profiler_end_step, profiling_save_path) | ||
| 240 | + | ||
| 241 | + timing_callback = TimingCallback(prof, mod) | ||
| 242 | + | ||
| 243 | + trainer = Trainer( | ||
| 244 | + model=self.model, | ||
| 245 | + args=training_args, | ||
| 246 | + train_dataset=train_dataset, | ||
| 247 | + eval_dataset=eval_dataset, | ||
| 248 | + tokenizer=self.tokenizer, | ||
| 249 | + data_collator=data_collator, | ||
| 250 | + callbacks=[timing_callback], | ||
| 251 | + ) | ||
| 252 | + return trainer | ||
| 253 | + | ||
| 254 | + def train(self): | ||
| 255 | + logger.info("Starting training...") | ||
| 256 | + | ||
| 257 | + if self.args.use_lora: | ||
| 258 | + self.apply_lora() | ||
| 259 | + | ||
| 260 | + dataset = self.load_and_preprocess_data() | ||
| 261 | + if self.args.validation_split > 0: | ||
| 262 | + split_dataset = dataset.train_test_split( | ||
| 263 | + test_size=self.args.validation_split, | ||
| 264 | + seed=self.args.seed | ||
| 265 | + ) | ||
| 266 | + train_dataset = split_dataset["train"] | ||
| 267 | + eval_dataset = split_dataset["test"] | ||
| 268 | + else: | ||
| 269 | + train_dataset = dataset | ||
| 270 | + eval_dataset = None | ||
| 271 | + | ||
| 272 | + trainer = self.create_trainer(train_dataset, eval_dataset) | ||
| 273 | + train_result = trainer.train() | ||
| 274 | + if self.args.enable_compile: | ||
| 275 | + headers, values = torch._dynamo.utils.compile_times("csv") | ||
| 276 | + for header, value in zip(headers, values): | ||
| 277 | + if header == "PyCodeCache.load_by_key_path": | ||
| 278 | + numbers = [float(num.strip()) for num in value.split(',') if num.strip()] | ||
| 279 | + op_compile_time = sum(numbers) | ||
| 280 | + print(f"op_compile_time:{op_compile_time * 1e3} ms", ) | ||
| 281 | + | ||
| 282 | + trainer.save_model() | ||
| 283 | + self.tokenizer.save_pretrained(self.args.output_dir) | ||
| 284 | + | ||
| 285 | + metrics = train_result.metrics | ||
| 286 | + trainer.log_metrics("train", metrics) | ||
| 287 | + trainer.save_metrics("train", metrics) | ||
| 288 | + trainer.save_state() | ||
| 289 | + logger.info(f"Training completed! Model saved to {self.args.output_dir}") | ||
| 290 | + return metrics | ||
| 291 | + | ||
| 292 | + | ||
| 293 | +def main(): | ||
| 294 | + parser = argparse.ArgumentParser(description="Train Qwen3 model") | ||
| 295 | + parser.add_argument("--model_path", type=str, required=True, | ||
| 296 | + help="Path to the pretrained model") | ||
| 297 | + parser.add_argument("--data_path", type=str, required=True, | ||
| 298 | + help="Path to training data (json/jsonl file or dataset name)") | ||
| 299 | + parser.add_argument("--output_dir", type=str, default="./llama3-finetuned", | ||
| 300 | + help="Output directory for trained model") | ||
| 301 | + parser.add_argument("--num_epochs", type=int, default=3, | ||
| 302 | + help="Number of training epochs") | ||
| 303 | + parser.add_argument("--batch_size", type=int, default=4, | ||
| 304 | + help="Batch size per device") | ||
| 305 | + parser.add_argument("--gradient_accumulation_steps", type=int, default=2, | ||
| 306 | + help="Gradient accumulation steps") | ||
| 307 | + parser.add_argument("--learning_rate", type=float, default=2e-4, | ||
| 308 | + help="Learning rate") | ||
| 309 | + parser.add_argument("--warmup_steps", type=int, default=5, | ||
| 310 | + help="Warmup steps") | ||
| 311 | + parser.add_argument("--max_steps", type=int, default=1, | ||
| 312 | + help="Total training steps") | ||
| 313 | + parser.add_argument("--weight_decay", type=float, default=0.01, | ||
| 314 | + help="Weight decay") | ||
| 315 | + parser.add_argument("--max_length", type=int, default=2048, | ||
| 316 | + help="Maximum sequence length") | ||
| 317 | + parser.add_argument("--use_lora", action="store_true", | ||
| 318 | + help="Use LoRA for parameter-efficient fine-tuning") | ||
| 319 | + parser.add_argument("--lora_r", type=int, default=8, | ||
| 320 | + help="LoRA rank") | ||
| 321 | + parser.add_argument("--lora_alpha", type=int, default=32, | ||
| 322 | + help="LoRA alpha") | ||
| 323 | + parser.add_argument("--lora_dropout", type=float, default=0.1, | ||
| 324 | + help="LoRA dropout") | ||
| 325 | + parser.add_argument("--use_4bit", action="store_true", | ||
| 326 | + help="Use 4-bit quantization") | ||
| 327 | + parser.add_argument("--use_fp16", action="store_true", | ||
| 328 | + help="Use FP16 precision") | ||
| 329 | + parser.add_argument("--use_bf16", action="store_true", | ||
| 330 | + help="Use BF16 precision") | ||
| 331 | + parser.add_argument("--gradient_checkpointing", action="store_true", | ||
| 332 | + help="Enable gradient checkpointing") | ||
| 333 | + parser.add_argument("--validation_split", type=float, default=0.1, | ||
| 334 | + help="Validation split ratio") | ||
| 335 | + parser.add_argument("--split", type=str, default="train", | ||
| 336 | + help="Dataset split to use") | ||
| 337 | + parser.add_argument("--num_proc", type=int, default=0, | ||
| 338 | + help="Number of processes for data preprocessing (0 = single process, avoid CUDA conflict)") | ||
| 339 | + parser.add_argument("--pad_to_max_length", action="store_true", | ||
| 340 | + help="Pad sequences to max_length") | ||
| 341 | + parser.add_argument("--overwrite_cache", action="store_true", | ||
| 342 | + help="Overwrite cached features") | ||
| 343 | + parser.add_argument("--seed", type=int, default=42, | ||
| 344 | + help="Random seed") | ||
| 345 | + parser.add_argument("--logging_steps", type=int, default=1, | ||
| 346 | + help="Log every X updates steps") | ||
| 347 | + parser.add_argument("--save_steps", type=int, default=500, | ||
| 348 | + help="Save checkpoint every X updates steps") | ||
| 349 | + parser.add_argument("--eval_steps", type=int, default=500, | ||
| 350 | + help="Evaluate every X updates steps") | ||
| 351 | + parser.add_argument("--save_total_limit", type=int, default=3, | ||
| 352 | + help="Limit the total amount of checkpoints") | ||
| 353 | + parser.add_argument("--enable_compile", action="store_true", | ||
| 354 | + help="Enable torch.compile and Inductor backend") | ||
| 355 | + parser.add_argument("--enable_profiler", action="store_true", | ||
| 356 | + help="Enable profiler for performance analysis") | ||
| 357 | + parser.add_argument("--profiler_start_step", type=int, default=5, | ||
| 358 | + help="Output directory for trained model") | ||
| 359 | + parser.add_argument("--profiler_end_step", type=int, default=8, | ||
| 360 | + help="Output directory for trained model") | ||
| 361 | + parser.add_argument("--profiler_save_path", type=str, default="./profile", | ||
| 362 | + help="Output directory for trained model") | ||
| 363 | + parser.add_argument("--npu-backend", type=str, default="mlir") | ||
| 364 | + parser.add_argument("--model_name", type=str, default="llama3-8B", help="your model name") | ||
| 365 | + parser.add_argument("--mfusion", action="store_true", help="Enable MFusion for graph fusion optimization") | ||
| 366 | + | ||
| 367 | + args = parser.parse_args() | ||
| 368 | + torch.manual_seed(args.seed) | ||
| 369 | + | ||
| 370 | + args.device_type = detect_device_type() | ||
| 371 | + os.environ['TORCHINDUCTOR_NPU_BACKEND']=args.npu_backend | ||
| 372 | + if args.npu_backend == "akg": | ||
| 373 | + os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir' | ||
| 374 | + os.environ['TORCHINDUCTOR_USE_AKG'] = '1' | ||
| 375 | + if args.mfusion: | ||
| 376 | + os.environ['TORCHINDUCTOR_ENABLE_MFUSION']='1' | ||
| 377 | + print(f"{args.device_type} train {args.model_name}") | ||
| 378 | + trainer = LLama3Trainer(args) | ||
| 379 | + metrics = trainer.train() | ||
| 380 | + print("\n" + "="*50) | ||
| 381 | + print("Training completed successfully!") | ||
| 382 | + print(f"Model saved to: {args.output_dir}") | ||
| 383 | + print(f"Final training loss: {metrics.get('train_loss', 'N/A')}") | ||
| 384 | + print("="*50) | ||
| 385 | + | ||
| 386 | +if __name__ == "__main__": | ||
| 387 | + main() | ||