From 862dd151f24d4b5c699eeb0509793c777c01092c Mon Sep 17 00:00:00 2001 From: jingyaogong Date: Thu, 9 Apr 2026 16:36:48 +0800 Subject: [PATCH] [fix] lora moe --- scripts/convert_model.py | 3 ++- trainer/train_lora.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/scripts/convert_model.py b/scripts/convert_model.py index daccdb9..f409229 100644 --- a/scripts/convert_model.py +++ b/scripts/convert_model.py @@ -127,6 +127,7 @@ def convert_json_to_jinja(json_file_path, output_path): if __name__ == '__main__': lm_config = MiniMindConfig(hidden_size=768, num_hidden_layers=8, max_seq_len=8192, use_moe=False) + # convert torch to transformers torch_path = f"../out/full_sft_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth" transformers_path = '../minimind-3' @@ -134,7 +135,7 @@ if __name__ == '__main__': # # merge lora # base_torch_path = f"../out/full_sft_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth" - # lora_path = f"../out/lora_identity_{lm_config.hidden_size}.pth" + # lora_path = f"../out/lora_identity_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth" # merged_torch_path = f"../out/merge_identity_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth" # convert_merge_base_lora(base_torch_path, lora_path, merged_torch_path) diff --git a/trainer/train_lora.py b/trainer/train_lora.py index 8c51c08..271fecc 100644 --- a/trainer/train_lora.py +++ b/trainer/train_lora.py @@ -58,7 +58,8 @@ def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None): if (step % args.save_interval == 0 or step == iters) and is_main_process(): model.eval() - lora_save_path = f'{args.save_dir}/{args.lora_name}_{lm_config.hidden_size}.pth' + moe_suffix = '_moe' if lm_config.use_moe else '' + lora_save_path = f'{args.save_dir}/{args.lora_name}_{lm_config.hidden_size}{moe_suffix}.pth' # LoRA只保存LoRA权重 save_lora(model, lora_save_path) lm_checkpoint(lm_config, weight=args.lora_name, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints')