Improve training performance with torch.compile and torch.amp

- Update AMP usage from `torch.cuda.amp` to `torch.amp` for broader device compatibility.
- Add `torch.compile` to models in various training scripts (SFT, LoRA, DPO, Distillation, Pretrain, Distill Reason) to leverage potential performance gains.
- Implement a custom `load_state_dict` in `MiniMindForCausalLM` to handle state dict keys potentially modified by `torch.compile`.

model/model_minimind.py

@@ -423,6 +423,16 @@ class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
         self.model.embed_tokens.weight = self.lm_head.weight
         self.OUT = CausalLMOutputWithPast()
         
+    def load_state_dict(self, state_dict, strict: bool = True):
+        new_state = {}
+        for k, v in state_dict.items():
+            if k.startswith("_orig_mod."):
+                new_key = k[len("_orig_mod."):]
+            else:
+                new_key = k
+            new_state[new_key] = v
+        return super().load_state_dict(new_state, strict=strict)
+    
     def forward(self,
                 input_ids: Optional[torch.Tensor] = None,
                 attention_mask: Optional[torch.Tensor] = None,

trainer/train_distill_reason.py

@@ -168,7 +168,12 @@ if __name__ == "__main__":
 
     args.wandb_run_name = f"MiniMind-Distill-Reasoning-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
 
-    ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
+    if device_type == "cpu":
+        ctx = nullcontext()
+    else:
+        amp_dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+        ctx = torch.amp.autocast(device_type=device_type, dtype=amp_dtype)
+        
     ddp = int(os.environ.get("RANK", -1)) != -1  # is this a ddp run?
     ddp_local_rank, DEVICE = 0, "cuda:0"
     base_seed = 1337
@@ -191,6 +196,7 @@ if __name__ == "__main__":
         wandb = None
 
     model, tokenizer = init_model(lm_config)
+    model = torch.compile(model)
     
     train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
     train_sampler = DistributedSampler(train_ds) if ddp else None
@@ -204,7 +210,7 @@ if __name__ == "__main__":
         sampler=train_sampler
     )
 
-    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
+    scaler = torch.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
     optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
 
     if ddp:

trainer/train_distillation.py

@@ -215,7 +215,12 @@ if __name__ == "__main__":
 
     args.wandb_run_name = f"MiniMind-Dist-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
 
-    ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
+    if device_type == "cpu":
+        ctx = nullcontext()
+    else:
+        amp_dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+        ctx = torch.amp.autocast(device_type=device_type, dtype=amp_dtype)
+        
     ddp = int(os.environ.get("RANK", -1)) != -1  # is this a ddp run?
     ddp_local_rank, DEVICE = 0, "cuda:0"
     base_seed = 1337
@@ -241,6 +246,9 @@ if __name__ == "__main__":
     model, tokenizer = init_student_model(lm_config_student)
     teacher_model = init_teacher_model(lm_config_teacher)
     
+    model = torch.compile(model)
+    teacher_model = torch.compile(teacher_model)
+
     train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
     train_sampler = DistributedSampler(train_ds) if ddp else None
     train_loader = DataLoader(
@@ -253,7 +261,7 @@ if __name__ == "__main__":
         sampler=train_sampler
     )
 
-    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
+    scaler = torch.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
     optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
 
     if ddp:

trainer/train_dpo.py

@@ -198,7 +198,12 @@ if __name__ == "__main__":
 
     args.wandb_run_name = f"MiniMind-Full-DPO-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
 
-    ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
+    if device_type == "cpu":
+        ctx = nullcontext()
+    else:
+        amp_dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+        ctx = torch.amp.autocast(device_type=device_type, dtype=amp_dtype)
+        
     ddp = int(os.environ.get("RANK", -1)) != -1  # is this a ddp run?
     ddp_local_rank, DEVICE = 0, "cuda:0"
     base_seed = 1337
@@ -221,6 +226,8 @@ if __name__ == "__main__":
         wandb = None
 
     model, ref_model, tokenizer = init_model(lm_config)
+    model = torch.compile(model)
+    ref_model = torch.compile(ref_model)
     
     train_ds = DPODataset(args.data_path, tokenizer, max_length=args.max_seq_len)
     train_sampler = DistributedSampler(train_ds) if ddp else None
@@ -234,7 +241,7 @@ if __name__ == "__main__":
         sampler=train_sampler
     )
 
-    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
+    scaler = torch.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
     optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
 
     if ddp:

trainer/train_full_sft.py

@@ -154,7 +154,12 @@ if __name__ == "__main__":
 
     args.wandb_run_name = f"MiniMind-Full-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
 
-    ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
+    if device_type == "cpu":
+        ctx = nullcontext()
+    else:
+        amp_dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+        ctx = torch.amp.autocast(device_type=device_type, dtype=amp_dtype)
+        
     ddp = int(os.environ.get("RANK", -1)) != -1  # is this a ddp run?
     ddp_local_rank, DEVICE = 0, "cuda:0"
     base_seed = 1337
@@ -177,6 +182,7 @@ if __name__ == "__main__":
         wandb = None
 
     model, tokenizer = init_model(lm_config)
+    model = torch.compile(model)
     
     train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
     train_sampler = DistributedSampler(train_ds) if ddp else None
@@ -190,7 +196,7 @@ if __name__ == "__main__":
         sampler=train_sampler
     )
 
-    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
+    scaler = torch.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
     optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
 
     if ddp:

trainer/train_lora.py

@@ -146,7 +146,12 @@ if __name__ == "__main__":
     tokens_per_iter = args.batch_size * args.max_seq_len
     device_type = "cuda" if "cuda" in args.device else "cpu"
 
-    ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
+    if device_type == "cpu":
+        ctx = nullcontext()
+    else:
+        amp_dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+        ctx = torch.amp.autocast(device_type=device_type, dtype=amp_dtype)
+        
     ddp = int(os.environ.get("RANK", -1)) != -1  # is this a ddp run?
     ddp_local_rank, DEVICE = 0, "cuda:0"
     base_seed = 1337
@@ -170,6 +175,7 @@ if __name__ == "__main__":
         wandb = None
 
     model, tokenizer = init_model(lm_config)
+    model = torch.compile(model)
     apply_lora(model)
 
     total_params = sum(p.numel() for p in model.parameters())  # 总参数数量
@@ -201,7 +207,7 @@ if __name__ == "__main__":
         sampler=train_sampler
     )
 
-    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
+    scaler = torch.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
     iter_per_epoch = len(train_loader)
 
     for epoch in range(args.epochs):

trainer/train_pretrain.py

@@ -149,7 +149,11 @@ if __name__ == "__main__":
 
     args.wandb_run_name = f"MiniMind-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
 
-    ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
+    if device_type == "cpu":
+        ctx = nullcontext()
+    else:
+        amp_dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+        ctx = torch.amp.autocast(device_type=device_type, dtype=amp_dtype)
 
     ddp = int(os.environ.get("RANK", -1)) != -1  # is this a ddp run?
     ddp_local_rank, DEVICE = 0, "cuda:0"
@@ -174,6 +178,8 @@ if __name__ == "__main__":
         wandb = None
 
     model, tokenizer = init_model(lm_config)
+    model = torch.compile(model)
+    
     train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
     train_sampler = DistributedSampler(train_ds) if ddp else None
     train_loader = DataLoader(
@@ -186,7 +192,7 @@ if __name__ == "__main__":
         sampler=train_sampler
     )
 
-    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
+    scaler = torch.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
     optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
 
     if ddp: