""" Fusion 模型全参数微调脚本 支持： - 本地 FusionModel（无需预训练权重） - 8B 模型：单卡 24GB（开启 ZeRO-3 offload） - 14B 模型：双卡 24GB 或单卡 48GB - DeepSpeed ZeRO-3 支持 - 混合精度训练（BF16/FP16） 使用方法： # 本地模型全参微调 python train/full_finetune.py --local_model --data_path data/example_data.json # 8B 模型 + DeepSpeed ZeRO-3 deepspeed train/full_finetune.py --local_model --model_size 8B --deepspeed configs/ds_zero3.json --data_path data/example_data.json 作者：zhan1206 项目：Fusion - 六边形开源大模型 许可证：Apache 2.0 """ import argparse import torch import torch.nn as nn import deepspeed from transformers import ( get_linear_schedule_with_warmup, ) from models.tokenizer import get_tokenizer, get_effective_vocab_size from torch.utils.data import Dataset, DataLoader import json import os import sys import logging # 添加项目根目录到路径 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from models import FusionModel, FusionConfig logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # ============================================================ # 数据格式说明 # ============================================================ """ 训练数据格式（JSON）： [ { "prompt": "解释量子纠缠", "response": "量子纠缠是...", "think_rank": 2 }, ... ] """ class FusionFullFinetuneDataset(Dataset): """ 全参数微调数据集 """ def __init__( self, data_path: str, tokenizer, max_length: int = 2048, ): self.tokenizer = tokenizer self.max_length = max_length with open(data_path, 'r', encoding='utf-8') as f: self.data = json.load(f) logger.info(f"[FusionFullFinetuneDataset] 加载数据集：{len(self.data)} 条样本") def __len__(self): return len(self.data) def __getitem__(self, idx): item = self.data[idx] prompt = item["prompt"] response = item["response"] think_rank = item.get("think_rank", 0) if think_rank > 0: thinking_token = f"<|think_depth_{think_rank}|>" full_text = f"{thinking_token}\n{prompt}\n{response}" else: full_text = f"{prompt}\n{response}" encoding = self.tokenizer( full_text, max_length=self.max_length, padding="max_length", truncation=True, return_tensors="pt", ) return { "input_ids": encoding["input_ids"].squeeze(0), "attention_mask": encoding["attention_mask"].squeeze(0), "labels": encoding["input_ids"].squeeze(0).clone(), } def create_local_model( model_size: str = "8B", torch_dtype: torch.dtype = torch.bfloat16, vocab_size_override: Optional[int] = None, ): """ 创建本地 FusionModel（无需预训练权重） """ model_configs = { "0.5B": dict(vocab_size=32000, hidden_size=2048, num_hidden_layers=16, num_attention_heads=16, num_key_value_heads=8, intermediate_size=5504), "1.5B": dict(vocab_size=32000, hidden_size=3072, num_hidden_layers=24, num_attention_heads=24, num_key_value_heads=8, intermediate_size=8192), "8B": dict(vocab_size=100000, hidden_size=4096, num_hidden_layers=32, num_attention_heads=32, num_key_value_heads=8, intermediate_size=11008), "14B": dict(vocab_size=100000, hidden_size=5120, num_hidden_layers=40, num_attention_heads=40, num_key_value_heads=8, intermediate_size=13824), } if model_size not in model_configs: raise ValueError(f"不支持的模型大小：{model_size}") config_dict = model_configs[model_size] # S3 fix: override vocab_size to match actual tokenizer if vocab_size_override is not None: config_dict['vocab_size'] = vocab_size_override common_config = dict( block_size=512, latent_dim=64, window_size=2048, sbla_mode="hybrid", rms_norm_eps=1e-6, rope_theta=10000.0, tie_word_embeddings=False, enable_thinking_dial=True, num_thinking_depths=4, ) config = FusionConfig(**config_dict, **common_config) logger.info(f"[create_local_model] 创建 Fusion-{model_size}（随机初始化）") logger.info(f" hidden_size={config.hidden_size}, layers={config.num_hidden_layers}, " f"heads={config.num_attention_heads}") model = FusionModel(config) total_params = sum(p.numel() for p in model.parameters()) logger.info(f"[create_local_model] 参数总量：{total_params / 1e9:.2f}B") return model, config def create_tokenizer(tokenizer_type: str = "fusion", vocab_size: int = 32000): """ Create tokenizer using the unified tokenizer module. """ logger.info(f"[create_tokenizer] Creating tokenizer: type={tokenizer_type}, vocab_size={vocab_size}") tokenizer = get_tokenizer(tokenizer_type, vocab_size=vocab_size) return tokenizer def train(args): """ 主训练函数 """ logger.info("=" * 60) logger.info("[train] 开始全参数微调") logger.info(f" 模型大小：{args.model_size}") logger.info(f" 使用 DeepSpeed：{args.deepspeed is not None}") logger.info(f" 数据路径：{args.data_path}") logger.info("=" * 60) # 1. 设备设置 if args.local_rank == -1: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") logger.info(f"[train] 单卡训练，设备：{device}") else: torch.cuda.set_device(args.local_rank) device = torch.device("cuda", args.local_rank) logger.info(f"[train] 分布式训练，local_rank：{args.local_rank}") # 2. 加载 tokenizer vocab_size_map = {"0.5B": 32000, "1.5B": 32000, "8B": 100000, "14B": 100000} tokenizer = create_tokenizer(vocab_size=vocab_size_map.get(args.model_size, 32000)) # Sync vocab_size to actual tokenizer size to prevent index-out-of-range (S3) actual_vocab_size = len(tokenizer) if actual_vocab_size != vocab_size_map.get(args.model_size, 32000): logger.warning(f"[S3-fix] Vocab size mismatch: config={vocab_size_map.get(args.model_size, 32000)}, tokenizer={actual_vocab_size}. Syncing to tokenizer.") vocab_size_map[args.model_size] = actual_vocab_size # 3. 创建模型（本地随机初始化） model, config = create_local_model(args.model_size, torch_dtype=args.torch_dtype, vocab_size_override=actual_vocab_size) # 4. 加载数据集 train_dataset = FusionFullFinetuneDataset( data_path=args.data_path, tokenizer=tokenizer, max_length=args.max_length, ) train_loader = DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True, ) # 5. 优化器 optimizer = torch.optim.AdamW( model.parameters(), lr=args.learning_rate, betas=(0.9, 0.95), weight_decay=args.weight_decay, ) # 6. 学习率调度器 total_steps = len(train_loader) * args.num_epochs // args.gradient_accumulation_steps warmup_steps = int(total_steps * args.warmup_ratio) scheduler = get_linear_schedule_with_warmup( optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps, ) # 7. DeepSpeed 初始化 if args.deepspeed: logger.info(f"[train] 使用 DeepSpeed：{args.deepspeed}") model_engine, optimizer, _, _ = deepspeed.initialize( model=model, optimizer=optimizer, config=args.deepspeed, ) else: model = model.to(device) model_engine = None # 8. 训练循环 logger.info("[train] 开始训练循环...") global_step = 0 for epoch in range(args.num_epochs): logger.info(f"[train] Epoch {epoch + 1}/{args.num_epochs}") model.train() for step, batch in enumerate(train_loader): input_ids = batch["input_ids"].to(device) attention_mask = batch["attention_mask"].to(device) labels = batch["labels"].to(device) if args.deepspeed: outputs = model_engine( input_ids=input_ids, attention_mask=attention_mask, labels=labels, ) loss = outputs.loss model_engine.backward(loss) model_engine.step() else: outputs = model( input_ids=input_ids, attention_mask=attention_mask, labels=labels, ) loss = outputs.loss loss = loss / args.gradient_accumulation_steps loss.backward() if (step + 1) % args.gradient_accumulation_steps == 0: torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm) optimizer.step() scheduler.step() optimizer.zero_grad() global_step += 1 if global_step % args.logging_steps == 0: logger.info(f"Step {global_step} | Loss: {loss.item():.4f} | " f"LR: {scheduler.get_last_lr()[0]:.2e}") # 保存检查点 if args.deepspeed: if model_engine.local_rank == 0: save_path = os.path.join(args.output_dir, f"epoch_{epoch + 1}") model_engine.save_checkpoint(save_path) else: if args.local_rank in [-1, 0]: save_path = os.path.join(args.output_dir, f"epoch_{epoch + 1}") model.save_pretrained(save_path) tokenizer.save_pretrained(save_path) config_path = os.path.join(save_path, "fusion_config.json") with open(config_path, 'w', encoding='utf-8') as f: json.dump(config.to_dict(), f, indent=2) logger.info(f"[train] Epoch {epoch + 1} 完成，保存到 {args.output_dir}") logger.info("[train] 全参数微调完成！") def main(): parser = argparse.ArgumentParser(description="Fusion 模型全参数微调") # 模型参数 parser.add_argument("--model_size", type=str, default="1.5B", choices=["0.5B", "1.5B", "8B", "14B"], help="模型大小") parser.add_argument("--local_model", action="store_true", default=True, help="使用本地 FusionModel（默认）") parser.add_argument("--torch_dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="模型精度") # 训练参数 parser.add_argument("--data_path", type=str, required=True, help="训练数据路径（JSON 格式）") parser.add_argument("--output_dir", type=str, default="./output/fusion-full", help="输出目录") parser.add_argument("--num_epochs", type=int, default=3, help="训练轮数") parser.add_argument("--batch_size", type=int, default=2, help="批次大小（根据显存调整）") parser.add_argument("--gradient_accumulation_steps", type=int, default=16, help="梯度累积步数") parser.add_argument("--learning_rate", type=float, default=1e-5, help="学习率") parser.add_argument("--weight_decay", type=float, default=0.01, help="权重衰减") parser.add_argument("--warmup_ratio", type=float, default=0.03, help="预热步数比例") parser.add_argument("--max_grad_norm", type=float, default=1.0, help="梯度裁剪") parser.add_argument("--max_length", type=int, default=2048, help="最大序列长度") # 硬件参数 parser.add_argument("--num_workers", type=int, default=4, help="数据加载线程数") parser.add_argument("--local_rank", type=int, default=-1, help="用于分布式训练（由 torchrun 自动设置）") # DeepSpeed parser.add_argument("--deepspeed", type=str, default=None, help="DeepSpeed 配置文件路径") # 日志 parser.add_argument("--logging_steps", type=int, default=10, help="日志打印间隔") args = parser.parse_args() # 设置 torch dtype dtype_map = {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16} args.torch_dtype = dtype_map.get(args.torch_dtype, torch.bfloat16) train(args) if __name__ == "__main__": main()