From ccdf6c576fecb90eeab5d0e351f3738d6e82a6e6 Mon Sep 17 00:00:00 2001 From: Gouryella Date: Wed, 21 May 2025 12:16:02 +0800 Subject: [PATCH] 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 | 12 +++++++++++- trainer/train_distill_reason.py | 12 +++++++++--- trainer/train_distillation.py | 12 ++++++++++-- trainer/train_dpo.py | 13 ++++++++++--- trainer/train_full_sft.py | 12 +++++++++--- trainer/train_lora.py | 10 ++++++++-- trainer/train_pretrain.py | 10 ++++++++-- 7 files changed, 65 insertions(+), 16 deletions(-) diff --git a/model/model_minimind.py b/model/model_minimind.py index 4b43610..0ed1080 100644 --- a/model/model_minimind.py +++ b/model/model_minimind.py @@ -422,7 +422,17 @@ class MiniMindForCausalLM(PreTrainedModel, GenerationMixin): self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False) 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, diff --git a/trainer/train_distill_reason.py b/trainer/train_distill_reason.py index 579fedb..a774615 100644 --- a/trainer/train_distill_reason.py +++ b/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,7 +196,8 @@ 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 train_loader = DataLoader( @@ -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: diff --git a/trainer/train_distillation.py b/trainer/train_distillation.py index be790a4..679049e 100644 --- a/trainer/train_distillation.py +++ b/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 @@ -240,6 +245,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 @@ -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: diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py index c7701d0..0939993 100644 --- a/trainer/train_dpo.py +++ b/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,7 +226,9 @@ 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 train_loader = DataLoader( @@ -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: diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index e2a9a2f..ee08fac 100644 --- a/trainer/train_full_sft.py +++ b/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,7 +182,8 @@ 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 train_loader = DataLoader( @@ -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: diff --git a/trainer/train_lora.py b/trainer/train_lora.py index 4711062..b279e3b 100644 --- a/trainer/train_lora.py +++ b/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): diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index f0cfc58..8f291b0 100644 --- a/trainer/train_pretrain.py +++ b/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: