diff --git a/.gitignore b/.gitignore index 7956e1d..cfaa448 100644 --- a/.gitignore +++ b/.gitignore @@ -3,4 +3,5 @@ __pycache__/ .DS_Store out website/ -docs-minimind/ \ No newline at end of file +docs-minimind/ +.venv/ diff --git a/model/model_minimind.py b/model/model_minimind.py index dc080e5..419c4b2 100755 --- a/model/model_minimind.py +++ b/model/model_minimind.py @@ -212,10 +212,14 @@ class MiniMindModel(nn.Module): past_key_values = past_key_values or [None] * len(self.layers) start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0 hidden_states = self.dropout(self.embed_tokens(input_ids)) - # Recompute RoPE buffers lost during meta-device init (transformers>=5.x) - if self.freqs_cos[0, 0] == 0: - freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.config.head_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling) - self.freqs_cos, self.freqs_sin = freqs_cos.to(hidden_states.device), freqs_sin.to(hidden_states.device) + # Lazy recompute of RoPE buffers if they were dropped by transformers>=5.x meta-device init. + # Hot path: skipped via cheap python bool (no GPU sync). The original `freqs_cos[0,0]==0` + # tensor compare forced a device-to-host sync on every forward. + if getattr(self, '_rope_ready', False) is False: + if self.freqs_cos.numel() == 0 or bool(self.freqs_cos[0, 0].eq(0).cpu()): + freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.config.head_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling) + self.freqs_cos, self.freqs_sin = freqs_cos.to(hidden_states.device), freqs_sin.to(hidden_states.device) + self._rope_ready = True position_embeddings = (self.freqs_cos[start_pos:start_pos + seq_length], self.freqs_sin[start_pos:start_pos + seq_length]) presents = [] for layer, past_key_value in zip(self.layers, past_key_values): diff --git a/model/quant.py b/model/quant.py new file mode 100644 index 0000000..b276346 --- /dev/null +++ b/model/quant.py @@ -0,0 +1,177 @@ +""" +W8A8 QAT fake-quantization utilities (pure PyTorch, STE). + +- WeightFakeQuant: per-output-channel symmetric int8 on Linear weights. +- ActFakeQuant: per-tensor symmetric int8 on activations, with EMA observer + and a warmup window (collect stats only, no quant in forward). +- QATLinear: drop-in nn.Linear replacement that fake-quantizes its input + activation and its weight on every forward. +- prepare_qat: in-place swap every nn.Linear in a model with QATLinear, + skipping any module whose qualified name matches `skip_patterns` + (default: lm_head, MoE router `mlp.gate`). + +The STE used here zeros gradients for values that saturate (clip-aware STE), +which is more stable than plain pass-through for low-bit training. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class _FakeQuantize(torch.autograd.Function): + """round(x/scale + zp) -> clamp -> dequant. Backward: STE inside clip, 0 outside.""" + + @staticmethod + def forward(ctx, x, scale, zero_point, qmin, qmax): + q = torch.round(x / scale + zero_point) + q_clamped = q.clamp(qmin, qmax) + ctx.save_for_backward((q >= qmin) & (q <= qmax)) + return (q_clamped - zero_point) * scale + + @staticmethod + def backward(ctx, grad_output): + (mask,) = ctx.saved_tensors + return grad_output * mask.to(grad_output.dtype), None, None, None, None + + +def fake_quantize(x, scale, zero_point, qmin, qmax): + return _FakeQuantize.apply(x, scale, zero_point, qmin, qmax) + + +class WeightFakeQuant(nn.Module): + """Per-output-channel symmetric int8 fake-quant for Linear weight [out, in].""" + + def __init__(self, bits: int = 8): + super().__init__() + self.bits = bits + self.qmax = (1 << (bits - 1)) - 1 + self.qmin = -(1 << (bits - 1)) + + def forward(self, w: torch.Tensor) -> torch.Tensor: + max_abs = w.detach().abs().amax(dim=1, keepdim=True).clamp_min(1e-8) + scale = (max_abs / self.qmax).to(w.dtype) + zp = torch.zeros_like(scale) + return fake_quantize(w, scale, zp, self.qmin, self.qmax) + + +class ActFakeQuant(nn.Module): + """Per-tensor symmetric int8 fake-quant with EMA min/max observer + warmup.""" + + def __init__(self, bits: int = 8, momentum: float = 0.99, observer_steps: int = 200): + super().__init__() + self.bits = bits + self.qmax = (1 << (bits - 1)) - 1 + self.qmin = -(1 << (bits - 1)) + self.momentum = momentum + self.observer_steps = observer_steps + self.register_buffer("running_max", torch.zeros(1)) + self.register_buffer("step", torch.zeros(1, dtype=torch.long)) + self.enabled = True + + def extra_repr(self) -> str: + return f"bits={self.bits}, observer_steps={self.observer_steps}, momentum={self.momentum}" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if not self.enabled: + return x + if self.training: + cur_max = x.detach().abs().amax().to(self.running_max.dtype) + if self.step.item() == 0: + self.running_max.copy_(cur_max.unsqueeze(0)) + else: + self.running_max.mul_(self.momentum).add_(cur_max.unsqueeze(0) * (1.0 - self.momentum)) + self.step += 1 + # Warmup: collect stats only; passing x through unquantized lets the + # observer stabilize before fake-quant starts perturbing gradients. + if self.step.item() < self.observer_steps: + return x + max_abs = self.running_max.clamp_min(1e-8) + scale = (max_abs / self.qmax).to(x.dtype) + zp = torch.zeros_like(scale) + return fake_quantize(x, scale, zp, self.qmin, self.qmax) + + +class QATLinear(nn.Linear): + """nn.Linear with fake-quant on input activation and on weight.""" + + def __init__(self, in_features, out_features, bias=True, *, w_bits=8, a_bits=8, + a_momentum=0.99, a_observer_steps=200): + super().__init__(in_features, out_features, bias=bias) + self.weight_fq = WeightFakeQuant(bits=w_bits) + self.act_fq = ActFakeQuant(bits=a_bits, momentum=a_momentum, observer_steps=a_observer_steps) + + @classmethod + def from_float(cls, linear: nn.Linear, **kw) -> "QATLinear": + qm = cls(linear.in_features, linear.out_features, bias=linear.bias is not None, **kw) + # Share existing tensors so optimizer state stays attached after swap. + qm.weight = linear.weight + if linear.bias is not None: + qm.bias = linear.bias + return qm + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x_q = self.act_fq(x) + w_q = self.weight_fq(self.weight) + return F.linear(x_q, w_q, self.bias) + + +DEFAULT_SKIP_PATTERNS = ("lm_head", "mlp.gate") + + +def _matches_skip(full_name: str, pattern: str) -> bool: + """Pattern matches at a dot boundary, so 'mlp.gate' does NOT skip 'mlp.gate_proj'.""" + return full_name == pattern or full_name.endswith("." + pattern) + + +def prepare_qat(model: nn.Module, w_bits: int = 8, a_bits: int = 8, + skip_patterns=DEFAULT_SKIP_PATTERNS, + a_momentum: float = 0.99, a_observer_steps: int = 200) -> int: + """Swap every nn.Linear in `model` with QATLinear in place. + + A Linear is skipped iff its qualified name equals a skip pattern or ends + with `'.' + pattern` (segment-aware, not substring). Defaults skip the LM + head (precision-sensitive, often tied to embeddings) and the MoE router + `mlp.gate` (small, routing-sensitive) without touching the dense + FeedForward's `mlp.gate_proj`. + + Returns the number of linears replaced. + """ + replaced = 0 + for parent_name, parent in model.named_modules(): + for child_name, child in list(parent.named_children()): + if not isinstance(child, nn.Linear) or isinstance(child, QATLinear): + continue + full_name = f"{parent_name}.{child_name}" if parent_name else child_name + if any(_matches_skip(full_name, p) for p in skip_patterns): + continue + new_mod = QATLinear.from_float( + child, w_bits=w_bits, a_bits=a_bits, + a_momentum=a_momentum, a_observer_steps=a_observer_steps, + ) + # Move buffers to weight's device. dtype on buffers stays fp32 on purpose. + new_mod.act_fq.running_max = new_mod.act_fq.running_max.to(child.weight.device) + new_mod.act_fq.step = new_mod.act_fq.step.to(child.weight.device) + setattr(parent, child_name, new_mod) + replaced += 1 + return replaced + + +def set_act_quant_enabled(model: nn.Module, enabled: bool) -> None: + """Toggle activation fake-quant (e.g. disable for eval-on-fp comparison).""" + for m in model.modules(): + if isinstance(m, ActFakeQuant): + m.enabled = enabled + + +@torch.no_grad() +def bake_quantized_weights(model: nn.Module) -> None: + """Replace each QATLinear's weight with its fake-quantized value, in place. + + After this, the model can run as plain QATLinears without further weight + updates and the saved fp weights already reflect the int8 grid; useful for + final eval / export. + """ + for m in model.modules(): + if isinstance(m, QATLinear): + m.weight.copy_(m.weight_fq(m.weight)) diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index 50ba878..b3f71af 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -23,6 +23,7 @@ warnings.filterwarnings('ignore') def train_epoch(epoch, loader, iters, start_step=0, wandb=None): start_time = time.time() + last_log_time, last_log_step = start_time, start_step last_step = start_step for step, (input_ids, labels) in enumerate(loader, start=start_step + 1): input_ids = input_ids.to(args.device) @@ -49,13 +50,16 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): optimizer.zero_grad(set_to_none=True) if step % args.log_interval == 0 or step == iters: - spend_time = time.time() - start_time + now = time.time() + spend_time = now - start_time + window_its = (step - last_log_step) / max(now - last_log_time, 1e-6) + last_log_time, last_log_step = now, step current_loss = loss.item() * args.accumulation_steps current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0 current_logits_loss = current_loss - current_aux_loss current_lr = optimizer.param_groups[-1]['lr'] eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60 - Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min') + Logger(f'[{time.strftime("%H:%M:%S")}] Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), it/s: {window_its:.2f}, loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min') if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min}) if (step % args.save_interval == 0 or step == iters) and is_main_process(): @@ -103,7 +107,7 @@ if __name__ == "__main__": parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）") parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb") parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain", help="wandb项目名") - parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）") + parser.add_argument("--use_compile", default=1, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是，A100 上推荐默认 1）") args = parser.parse_args() # ========== 1. 初始化环境和随机种子 ========== diff --git a/trainer/train_qat.py b/trainer/train_qat.py new file mode 100644 index 0000000..9d9c227 --- /dev/null +++ b/trainer/train_qat.py @@ -0,0 +1,221 @@ +import os +import sys + +__package__ = "trainer" +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +import datasets # noqa: F401 # Windows pyarrow/torch DLL conflict workaround (issue #771) +import argparse +import time +import warnings +import torch +import torch.distributed as dist +from contextlib import nullcontext +from torch import optim +from torch.nn.parallel import DistributedDataParallel +from torch.utils.data import DataLoader, DistributedSampler +from model.model_minimind import MiniMindConfig +from model.quant import prepare_qat, DEFAULT_SKIP_PATTERNS +from dataset.lm_dataset import SFTDataset +from trainer.trainer_utils import ( + get_lr, Logger, is_main_process, lm_checkpoint, + init_distributed_mode, setup_seed, init_model, SkipBatchSampler, +) + +warnings.filterwarnings('ignore') + + +def train_epoch(epoch, loader, iters, start_step=0, wandb=None): + start_time = time.time() + last_step = start_step + for step, (input_ids, labels) in enumerate(loader, start=start_step + 1): + input_ids = input_ids.to(args.device) + labels = labels.to(args.device) + last_step = step + lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate) + for param_group in optimizer.param_groups: + param_group['lr'] = lr + + with autocast_ctx: + res = model(input_ids, labels=labels) + loss = res.loss + res.aux_loss + loss = loss / args.accumulation_steps + + scaler.scale(loss).backward() + + if step % args.accumulation_steps == 0: + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad(set_to_none=True) + + if step % args.log_interval == 0 or step == iters: + spend_time = time.time() - start_time + current_loss = loss.item() * args.accumulation_steps + current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0 + current_logits_loss = current_loss - current_aux_loss + current_lr = optimizer.param_groups[-1]['lr'] + eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60 + Logger( + f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), ' + f'loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, ' + f'aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, ' + f'epoch_time: {eta_min:.1f}min' + ) + if wandb: + wandb.log({ + "loss": current_loss, "logits_loss": current_logits_loss, + "aux_loss": current_aux_loss, "learning_rate": current_lr, + "epoch_time": eta_min, + }) + + if (step % args.save_interval == 0 or step == iters) and is_main_process(): + model.eval() + moe_suffix = '_moe' if lm_config.use_moe else '' + ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth' + raw_model = model.module if isinstance(model, DistributedDataParallel) else model + raw_model = getattr(raw_model, '_orig_mod', raw_model) + state_dict = raw_model.state_dict() + torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp) + lm_checkpoint( + lm_config, weight=args.save_weight, model=model, optimizer=optimizer, + epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler, + ) + model.train() + del state_dict + + del input_ids, labels, res, loss + + if last_step > start_step and last_step % args.accumulation_steps != 0: + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad(set_to_none=True) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="MiniMind QAT (W8A8 fake-quant fine-tune)") + parser.add_argument("--save_dir", type=str, default="../out") + parser.add_argument('--save_weight', default='qat', type=str) + parser.add_argument("--epochs", type=int, default=1, help="QAT 通常 1~2 epoch 即可") + parser.add_argument("--batch_size", type=int, default=16) + parser.add_argument("--learning_rate", type=float, default=2e-6, + help="QAT 用比 SFT 更小的 lr（默认 1/5），防止打乱已收敛的权重分布") + parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu") + parser.add_argument("--dtype", type=str, default="bfloat16") + parser.add_argument("--num_workers", type=int, default=8) + parser.add_argument("--accumulation_steps", type=int, default=1) + parser.add_argument("--grad_clip", type=float, default=1.0) + parser.add_argument("--log_interval", type=int, default=100) + parser.add_argument("--save_interval", type=int, default=1000) + parser.add_argument('--hidden_size', default=768, type=int) + parser.add_argument('--num_hidden_layers', default=8, type=int) + parser.add_argument('--max_seq_len', default=768, type=int) + parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1]) + parser.add_argument("--data_path", type=str, default="../dataset/sft_t2t_mini.jsonl") + parser.add_argument('--from_weight', default='full_sft', type=str, + help="QAT 起点权重前缀（默认基于 full_sft 微调）") + parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1]) + parser.add_argument("--use_wandb", action="store_true") + parser.add_argument("--wandb_project", type=str, default="MiniMind-QAT") + parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1]) + + # ---- QAT-specific ---- + parser.add_argument("--w_bits", type=int, default=8, help="weight 量化位宽") + parser.add_argument("--a_bits", type=int, default=8, help="activation 量化位宽") + parser.add_argument("--quant_observer_steps", type=int, default=200, + help="activation observer 预热步数（仅采集 EMA，不做 fake-quant）") + parser.add_argument("--quant_skip", type=str, default=",".join(DEFAULT_SKIP_PATTERNS), + help="名称中包含这些子串的 Linear 跳过量化（逗号分隔）" + "；默认跳过 lm_head 和 MoE 路由 mlp.gate") + args = parser.parse_args() + + # 1. dist + seed + local_rank = init_distributed_mode() + if dist.is_initialized(): + args.device = f"cuda:{local_rank}" + setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) + + # 2. config + checkpoint probe + os.makedirs(args.save_dir, exist_ok=True) + lm_config = MiniMindConfig( + hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, + use_moe=bool(args.use_moe), + ) + ckp_data = lm_checkpoint( + lm_config, weight=args.save_weight, save_dir='../checkpoints', + ) if args.from_resume == 1 else None + + # 3. autocast + device_type = "cuda" if "cuda" in args.device else "cpu" + dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 + autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) + + # 4. wandb + wandb = None + if args.use_wandb and is_main_process(): + import swanlab as wandb + wandb_id = ckp_data.get('wandb_id') if ckp_data else None + resume = 'must' if wandb_id else None + wandb_run_name = ( + f"MiniMind-QAT-W{args.w_bits}A{args.a_bits}-Epoch-{args.epochs}-" + f"BS-{args.batch_size}-LR-{args.learning_rate}" + ) + wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume) + + # 5. model: load fp weights -> wrap with QATLinear -> dataset/optimizer + if args.from_weight == 'none': + raise SystemExit("--from_weight=none is not supported for QAT; load a pretrained/SFT ckpt.") + model, tokenizer = init_model(lm_config, args.from_weight, device=args.device) + + skip_patterns = tuple(s.strip() for s in args.quant_skip.split(',') if s.strip()) + replaced = prepare_qat( + model, w_bits=args.w_bits, a_bits=args.a_bits, + skip_patterns=skip_patterns, a_observer_steps=args.quant_observer_steps, + ) + Logger(f'QAT: wrapped {replaced} Linear modules ' + f'(W{args.w_bits}A{args.a_bits}, skip={skip_patterns}, ' + f'observer_warmup={args.quant_observer_steps} steps)') + + train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len) + train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None + scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) + optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) + + # 6. resume + start_epoch, start_step = 0, 0 + if ckp_data: + model.load_state_dict(ckp_data['model']) # state dict already contains QAT buffers + optimizer.load_state_dict(ckp_data['optimizer']) + scaler.load_state_dict(ckp_data['scaler']) + start_epoch = ckp_data['epoch'] + start_step = ckp_data.get('step', 0) + + # 7. compile + DDP + if args.use_compile == 1: + model = torch.compile(model) + Logger('torch.compile enabled') + if dist.is_initialized(): + model = DistributedDataParallel(model, device_ids=[local_rank]) + + # 8. train + for epoch in range(start_epoch, args.epochs): + train_sampler and train_sampler.set_epoch(epoch) + setup_seed(42 + epoch) + indices = torch.randperm(len(train_ds)).tolist() + skip = start_step if (epoch == start_epoch and start_step > 0) else 0 + batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) + loader = DataLoader(train_ds, batch_sampler=batch_sampler, + num_workers=args.num_workers, pin_memory=True) + if skip > 0: + Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') + train_epoch(epoch, loader, len(loader) + skip, start_step, wandb) + else: + train_epoch(epoch, loader, len(loader), 0, wandb) + + # 9. cleanup + if dist.is_initialized(): + dist.barrier() + dist.destroy_process_group() diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py index b6211d5..96ad3dd 100644 --- a/trainer/trainer_utils.py +++ b/trainer/trainer_utils.py @@ -51,14 +51,18 @@ def init_distributed_mode(): return local_rank -def setup_seed(seed: int): +def setup_seed(seed: int, deterministic: bool = False): + # cudnn 决定性 + 禁 benchmark 是预训练吞吐杀手（A100 上 ~3-5x 慢）。 + # 默认走 benchmark=True 让 cudnn 选最快 kernel；要严格复现传 deterministic=True。 random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = deterministic + torch.backends.cudnn.benchmark = not deterministic + # A100 TF32 矩阵乘加速（autocast bf16 路径之外的 fp32 matmul 才生效，但无副作用） + torch.set_float32_matmul_precision('high') def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoch=0, step=0, wandb=None, save_dir='../checkpoints', **kwargs): os.makedirs(save_dir, exist_ok=True)