升腾NPU适配

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):
@@ -177,7 +179,21 @@ class Attention(nn.Module):
                 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
             
+            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(

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)

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')
 
 

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()
     # 定义学生模型和教师模型

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')
 
 

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')
 
 

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)

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')