diff --git a/README.md b/README.md index 2316c33..8973fd8 100644 --- a/README.md +++ b/README.md @@ -1580,13 +1580,28 @@ DPO和在线PPO的区别在于reject和chosen都是离线准备的，和minimind ## Ⅳ RoPE长度外推 MiniMind支持通过YaRN算法进行RoPE位置编码的长度外推，使模型能够处理超出训练长度的文本序列。 -在使用`eval_llm.py`进行推理时，只需添加`--inference_rope_scaling`参数即可启用RoPE外推： + +原生torch模型在使用`eval_llm.py`进行推理时，只需添加`--inference_rope_scaling`参数即可启用RoPE外推： ```bash python eval_llm.py --weight full_sft --inference_rope_scaling ``` -下图展示了在不同文本「西游记」白话文小说长度下，使用RoPE scaling前后的困惑度(PPL)对比。可以看出，启用RoPE scaling后，模型在长文本上的表现显著提升： +对于Transformers格式的模型，可以在config.json中添加以下配置实现长度外推： + +```json +"rope_scaling": { + "type": "yarn", + "factor": 16.0, + "original_max_position_embeddings": 2048, + "beta_fast": 32.0, + "beta_slow": 1.0, + "attention_factor": 1.0 +} +``` + +在MiniMind-Small模型上，测试输入不同长度的「西游记」白话文小说，评估RoPE scaling前后的困惑度(PPL)对比。 +可以看出，启用YaRN外推后，模型在长文本上的PPL表现显著下降：
diff --git a/README_en.md b/README_en.md index d4e2d3d..8c492b8 100644 --- a/README_en.md +++ b/README_en.md @@ -1547,13 +1547,32 @@ Personal subjective evaluation basically aligns with DeepSeek-R1, where: ## Ⅳ RoPE Long-text Extrapolation MiniMind supports RoPE position encoding length extrapolation through YaRN algorithm, enabling models to handle text sequences exceeding training length. -When using `eval_llm.py` for inference, just add `--inference_rope_scaling` parameter to enable RoPE extrapolation: + +For native torch models, when using `eval_llm.py` for inference, just add `--inference_rope_scaling` parameter to enable RoPE extrapolation: ```bash python eval_llm.py --weight full_sft --inference_rope_scaling ``` -The chart below shows perplexity (PPL) comparison before and after RoPE scaling on different lengths of "Journey to the West" vernacular fiction text. You can see that after enabling RoPE scaling, model performance on long texts is significantly improved. +For Transformers format models, add the following configuration to config.json to enable length extrapolation: + +```json +"rope_scaling": { + "type": "yarn", + "factor": 16.0, + "original_max_position_embeddings": 2048, + "beta_fast": 32.0, + "beta_slow": 1.0, + "attention_factor": 1.0 +} +``` + +Testing on MiniMind-Small model with different lengths of "Journey to the West" vernacular fiction text to evaluate perplexity (PPL) comparison before and after RoPE scaling. +You can see that after enabling YaRN extrapolation, the model's PPL performance on long texts significantly decreases: + +
+ +
## Ⅴ Objective Benchmarks diff --git a/images/rope_ppl.png b/images/rope_ppl.png index 223292e..ea07260 100644 Binary files a/images/rope_ppl.png and b/images/rope_ppl.png differ diff --git a/model/model_minimind.py b/model/model_minimind.py index a522b03..78e243c 100755 --- a/model/model_minimind.py +++ b/model/model_minimind.py @@ -54,12 +54,13 @@ class MiniMindConfig(PretrainedConfig): self.rms_norm_eps = rms_norm_eps self.rope_theta = rope_theta self.inference_rope_scaling = inference_rope_scaling - # 外推长度 = factor * original_max_position_embeddings + # 外推长度 = factor * original_max_position_embeddings = 32768 self.rope_scaling = { - "beta_fast": 4, + "beta_fast": 32, "beta_slow": 1, - "factor": 4, + "factor": 16, "original_max_position_embeddings": 2048, + "attention_factor": 1.0, "type": "yarn" } if self.inference_rope_scaling else None self.flash_attn = flash_attn @@ -107,24 +108,23 @@ class RMSNorm(torch.nn.Module): def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), rope_base: float = 1e6, rope_scaling: Optional[dict] = None): - freqs = 1.0 / (rope_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) + freqs, attn_factor = 1.0 / (rope_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)), 1.0 if rope_scaling is not None: - orig_max, factor, beta_fast, beta_slow = ( - rope_scaling.get("original_max_position_embeddings", 2048), rope_scaling.get("factor", 4), - rope_scaling.get("beta_fast", 4.0), rope_scaling.get("beta_slow", 1.0) + orig_max, factor, beta_fast, beta_slow, attn_factor = ( + rope_scaling.get("original_max_position_embeddings", 2048), rope_scaling.get("factor", 16), + rope_scaling.get("beta_fast", 32.0), rope_scaling.get("beta_slow", 1.0), rope_scaling.get("attention_factor", 1.0) ) if end / orig_max > 1.0: - corr_dim = next((i for i in range(dim // 2) if 2 * math.pi / freqs[i] > orig_max), dim // 2) - power = torch.arange(0, dim // 2, device=freqs.device).float() / max(dim // 2 - 1, 1) - beta = beta_slow + (beta_fast - beta_slow) * power - # λ = (β·α - β + 1)/(β·α) YaRN标准公式 - scale = torch.where(torch.arange(dim // 2, device=freqs.device) < corr_dim, (beta * factor - beta + 1) / (beta * factor), 1.0 / factor) - freqs = freqs * scale + # YaRN: f'(i) = f(i)((1-γ) + γ/s), where γ∈[0,1] is linear ramp + inv_dim = lambda b: (dim * math.log(orig_max / (b * 2 * math.pi))) / (2 * math.log(rope_base)) + low, high = max(math.floor(inv_dim(beta_fast)), 0), min(math.ceil(inv_dim(beta_slow)), dim // 2 - 1) + ramp = torch.clamp((torch.arange(dim // 2, device=freqs.device).float() - low) / max(high - low, 0.001), 0, 1) + freqs = freqs * (1 - ramp + ramp / factor) t = torch.arange(end, device=freqs.device) freqs = torch.outer(t, freqs).float() - freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) - freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) + freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attn_factor + freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attn_factor return freqs_cos, freqs_sin diff --git a/trainer/train_distill_reason.py b/trainer/train_distill_reason.py index be04798..111fe92 100644 --- a/trainer/train_distill_reason.py +++ b/trainer/train_distill_reason.py @@ -85,10 +85,11 @@ def train_epoch(epoch, loader, iters, tokenizer, lm_config, start_step=0, wandb= state_dict = model.module.state_dict() else: state_dict = model.state_dict() - state_dict = {k: v.half() for k, v in state_dict.items()} # 半精度保存 + state_dict = {k: v.half().cpu() for k, v in state_dict.items()} torch.save(state_dict, ckp) lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints') model.train() + del state_dict del X, Y, loss_mask, res, loss diff --git a/trainer/train_distillation.py b/trainer/train_distillation.py index 8f5e3f9..58d2b87 100644 --- a/trainer/train_distillation.py +++ b/trainer/train_distillation.py @@ -123,10 +123,11 @@ def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_st state_dict = model.module.state_dict() else: state_dict = model.state_dict() - state_dict = {k: v.half() for k, v in state_dict.items()} # 半精度保存 + state_dict = {k: v.half().cpu() for k, v in state_dict.items()} torch.save(state_dict, ckp) lm_checkpoint(lm_config_student, weight=args.save_weight, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints') model.train() + del state_dict del X, Y, loss_mask, res, student_logits, teacher_logits, ce_loss, distill_loss, loss diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py index afa51f4..c48f32d 100644 --- a/trainer/train_dpo.py +++ b/trainer/train_dpo.py @@ -110,10 +110,11 @@ def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb= state_dict = model.module.state_dict() else: state_dict = model.state_dict() - state_dict = {k: v.half() for k, v in state_dict.items()} # 半精度保存 + state_dict = {k: v.half().cpu() for k, v in state_dict.items()} torch.save(state_dict, ckp) lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints') model.train() + del state_dict del x_chosen, x_rejected, y_chosen, y_rejected, mask_chosen, mask_rejected, x, y, mask del ref_outputs, ref_logits, ref_log_probs, outputs, logits, policy_log_probs, loss diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index 4d0a181..8b380c7 100644 --- a/trainer/train_full_sft.py +++ b/trainer/train_full_sft.py @@ -72,11 +72,12 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): state_dict = model.module.state_dict() else: state_dict = model.state_dict() - state_dict = {k: v.half() for k, v in state_dict.items()} # 半精度保存 + state_dict = {k: v.half().cpu() for k, v in state_dict.items()} torch.save(state_dict, 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 X, Y, loss_mask, res, loss diff --git a/trainer/train_grpo.py b/trainer/train_grpo.py index 52f6562..4e11f22 100755 --- a/trainer/train_grpo.py +++ b/trainer/train_grpo.py @@ -175,10 +175,11 @@ def grpo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_token moe_suffix = '_moe' if lm_config.use_moe else '' ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth' state_dict = model.module.state_dict() if isinstance(model, DistributedDataParallel) else model.state_dict() - torch.save({k: v.half() for k, v in state_dict.items()}, ckp) + 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', scheduler=scheduler) model.train() + del state_dict del prompt_inputs, outputs, completion_ids, per_token_logps, ref_per_token_logps del completions, rewards, grouped_rewards, mean_r, std_r, advantages, completion_mask diff --git a/trainer/train_ppo.py b/trainer/train_ppo.py index cb6d7ff..a5ca18b 100644 --- a/trainer/train_ppo.py +++ b/trainer/train_ppo.py @@ -224,7 +224,7 @@ def ppo_train_epoch(epoch, loader, iters, old_actor_model, ref_model, actor_sche moe_suffix = '_moe' if lm_config.use_moe else '' ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth' actor_state = actor_model.module.state_dict() if isinstance(actor_model, DistributedDataParallel) else actor_model.state_dict() - torch.save({k: v.half() for k, v in actor_state.items()}, ckp) + torch.save({k: v.half().cpu() for k, v in actor_state.items()}, ckp) # 使用 lm_checkpoint 保存完整状态（包括 critic） lm_checkpoint(lm_config, weight=args.save_weight, model=actor_model, optimizer=actor_optimizer, @@ -232,6 +232,7 @@ def ppo_train_epoch(epoch, loader, iters, old_actor_model, ref_model, actor_sche scheduler=actor_scheduler, critic_model=critic_model, critic_optimizer=critic_optimizer, critic_scheduler=critic_scheduler) actor_model.train() + del actor_state del enc, gen_out, responses_text, rewards, full_mask, values_seq, values, advantages del logits, labels, logp_tokens, final_mask, actor_logp, old_logits, old_logp, ref_logits, ref_logp diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index 5de170b..5bc8fc8 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -117,10 +117,11 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): state_dict = model.module.state_dict() else: state_dict = model.state_dict() - state_dict = {k: v.half() for k, v in state_dict.items()} # 半精度保存 + state_dict = {k: v.half().cpu() for k, v in state_dict.items()} torch.save(state_dict, ckp) lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints') model.train() + del state_dict del X, Y, loss_mask, res, loss diff --git a/trainer/train_spo.py b/trainer/train_spo.py index c250b14..bac7c14 100755 --- a/trainer/train_spo.py +++ b/trainer/train_spo.py @@ -223,10 +223,11 @@ def spo_train_epoch(epoch, loader, iters, ref_model, reward_model, reward_tokeni moe_suffix = '_moe' if lm_config.use_moe else '' ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth' state_dict = model.module.state_dict() if isinstance(model, DistributedDataParallel) else model.state_dict() - torch.save({k: v.half() for k, v in state_dict.items()}, ckp) + 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', scheduler=scheduler) model.train() + del state_dict del prompt_inputs, outputs, completion_ids, per_token_logps, ref_per_token_logps del completions, rewards, advantages, completion_mask, baselines, response_masks diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py index a7e5058..528868e 100644 --- a/trainer/trainer_utils.py +++ b/trainer/trainer_utils.py @@ -1,6 +1,7 @@ """ 训练工具函数集合 """ +import gc import os import random import math @@ -77,8 +78,9 @@ def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoc if model is not None: from torch.nn.parallel import DistributedDataParallel state_dict = model.module.state_dict() if isinstance(model, DistributedDataParallel) else model.state_dict() + state_dict = {k: v.half().cpu() for k, v in state_dict.items()} ckp_tmp = ckp_path + '.tmp' - torch.save({k: v.half() for k, v in state_dict.items()}, ckp_tmp) + torch.save(state_dict, ckp_tmp) os.replace(ckp_tmp, ckp_path) wandb_id = None if wandb: @@ -109,6 +111,9 @@ def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoc resume_tmp = resume_path + '.tmp' torch.save(resume_data, resume_tmp) os.replace(resume_tmp, resume_path) + del state_dict, resume_data + gc.collect() + torch.cuda.empty_cache() else: # 加载模式 if os.path.exists(resume_path): ckp_data = torch.load(resume_path, map_location='cpu')