diff --git a/model/model_minimind.py b/model/model_minimind.py index 4b43610..c0a970e 100644 --- a/model/model_minimind.py +++ b/model/model_minimind.py @@ -79,7 +79,9 @@ from typing import Optional, Tuple, List, Union import torch.nn.functional as F from transformers import PreTrainedModel, GenerationMixin, PretrainedConfig from transformers.modeling_outputs import CausalLMOutputWithPast - +import torch_npu +from torch_npu.npu import amp # 导入AMP模块 +from torch_npu.contrib import transfer_to_npu # 使能自动迁移 class RMSNorm(torch.nn.Module): def __init__(self, dim: int, eps: float = 1e-5): @@ -176,8 +178,22 @@ class Attention(nn.Module): if attention_mask is not None: attn_mask = attention_mask.view(bsz, 1, 1, -1).expand(bsz, self.n_local_heads, seq_len, -1) attn_mask = attn_mask.bool() if attention_mask is not None else None - - output = F.scaled_dot_product_attention(xq, xk, xv, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=True) + + try: + output = F.scaled_dot_product_attention(xq, xk, xv, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=True) + except: + atten_mask_npu= torch.triu(torch.ones([2048, 2048]), diagonal=1).bool().to("npu") + head_num = xq.shape[1] + output = torch_npu.npu_fusion_attention( + xq, xk, xv, head_num, input_layout="BNSD", + pse=None, + atten_mask=atten_mask_npu, + scale=1.0 / math.sqrt(xq.shape[-1]), + pre_tockens=2147483647, + next_tockens=2147483647, + keep_prob=1, + sparse_mode=2 + )[0] else: scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim) scores = scores + torch.triu( diff --git a/scripts/convert_model.py b/scripts/convert_model.py index a6e22f8..54a70cb 100644 --- a/scripts/convert_model.py +++ b/scripts/convert_model.py @@ -62,13 +62,13 @@ def convert_transformers2torch(transformers_path, torch_path): if __name__ == '__main__': - lm_config = MiniMindConfig(hidden_size=768, num_hidden_layers=16, max_seq_len=8192, use_moe=False) + lm_config = MiniMindConfig(hidden_size=512, num_hidden_layers=8, max_seq_len=8192, use_moe=False) torch_path = f"../out/full_sft_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth" transformers_path = '../MiniMind2' - convert_torch2transformers_minimind(torch_path, transformers_path) + convert_torch2transformers_llama(torch_path, transformers_path) # # # convert transformers to torch model # # convert_transformers2torch(transformers_path, torch_path) diff --git a/trainer/train_distill_reason.py b/trainer/train_distill_reason.py index 579fedb..aad96e4 100644 --- a/trainer/train_distill_reason.py +++ b/trainer/train_distill_reason.py @@ -17,7 +17,9 @@ from torch.utils.data import DataLoader, DistributedSampler from transformers import AutoTokenizer, AutoModelForCausalLM from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import SFTDataset - +import torch_npu +from torch_npu.npu import amp # 导入AMP模块 +from torch_npu.contrib import transfer_to_npu # 使能自动迁移 warnings.filterwarnings('ignore') diff --git a/trainer/train_distillation.py b/trainer/train_distillation.py index be790a4..1f3bdad 100644 --- a/trainer/train_distillation.py +++ b/trainer/train_distillation.py @@ -18,7 +18,9 @@ from torch.utils.data import DataLoader, DistributedSampler from transformers import AutoTokenizer, AutoModelForCausalLM from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import SFTDataset - +import torch_npu +from torch_npu.npu import amp # 导入AMP模块 +from torch_npu.contrib import transfer_to_npu # 使能自动迁移 warnings.filterwarnings('ignore') @@ -33,9 +35,9 @@ def get_lr(current_step, total_steps, lr): def distillation_loss_fn(student_logits, teacher_logits, temperature=1.0, reduction='batchmean'): with torch.no_grad(): - teacher_probs = F.softmax(teacher_logits / temperature, hidden_size=-1).detach() + teacher_probs = F.softmax(teacher_logits / temperature).detach() - student_log_probs = F.log_softmax(student_logits / temperature, hidden_size=-1) + student_log_probs = F.log_softmax(student_logits / temperature) kl = F.kl_div( student_log_probs, @@ -201,7 +203,7 @@ if __name__ == "__main__": parser.add_argument("--save_interval", type=int, default=100) parser.add_argument("--max_seq_len", type=int, default=512) parser.add_argument('--local_rank', type=int, default=-1) - parser.add_argument("--data_path", type=str, default="../dataset/sft_xxx.jsonl") + parser.add_argument("--data_path", type=str, default="../dataset/sft_mini_512.jsonl") args = parser.parse_args() # 定义学生模型和教师模型 diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py index c7701d0..b81a11b 100644 --- a/trainer/train_dpo.py +++ b/trainer/train_dpo.py @@ -17,7 +17,9 @@ from torch.utils.data import DataLoader, DistributedSampler from transformers import AutoTokenizer, AutoModelForCausalLM from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import DPODataset - +import torch_npu +from torch_npu.npu import amp # 导入AMP模块 +from torch_npu.contrib import transfer_to_npu # 使能自动迁移 warnings.filterwarnings('ignore') diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index e2a9a2f..884a6e8 100644 --- a/trainer/train_full_sft.py +++ b/trainer/train_full_sft.py @@ -17,7 +17,9 @@ from torch.utils.data import DataLoader, DistributedSampler from transformers import AutoTokenizer, AutoModelForCausalLM from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import SFTDataset - +import torch_npu +from torch_npu.npu import amp # 导入AMP模块 +from torch_npu.contrib import transfer_to_npu # 使能自动迁移 warnings.filterwarnings('ignore') diff --git a/trainer/train_lora.py b/trainer/train_lora.py index 4711062..40486d3 100644 --- a/trainer/train_lora.py +++ b/trainer/train_lora.py @@ -17,7 +17,9 @@ from transformers import AutoTokenizer, AutoModelForCausalLM from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import SFTDataset from model.model_lora import load_lora, save_lora, apply_lora - +import torch_npu +from torch_npu.npu import amp # 导入AMP模块 +from torch_npu.contrib import transfer_to_npu # 使能自动迁移 warnings.filterwarnings('ignore') @@ -80,7 +82,6 @@ def train_epoch(epoch, wandb): wandb.log({"loss": loss * args.accumulation_steps, "lr": optimizer.param_groups[-1]['lr'], "epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60}) - if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0): model.eval() lora_save_path = f'{args.save_dir}/lora/{args.lora_name}_{lm_config.hidden_size}.pth' @@ -128,7 +129,7 @@ if __name__ == "__main__": parser.add_argument("--grad_clip", type=float, default=1.0) parser.add_argument("--warmup_iters", type=int, default=0) parser.add_argument("--log_interval", type=int, default=100) - parser.add_argument("--save_interval", type=int, default=100) + parser.add_argument("--save_interval", type=int, default=40) parser.add_argument('--local_rank', type=int, default=-1) parser.add_argument('--hidden_size', default=512, type=int) parser.add_argument('--num_hidden_layers', default=8, type=int) diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index f0cfc58..2399112 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -16,7 +16,9 @@ from contextlib import nullcontext from transformers import AutoTokenizer from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import PretrainDataset - +import torch_npu +from torch_npu.npu import amp # 导入AMP模块 +from torch_npu.contrib import transfer_to_npu # 使能自动迁移 warnings.filterwarnings('ignore')