diff --git a/model/model_minimind.py b/model/model_minimind.py index dc080e5..419c4b2 100755 --- a/model/model_minimind.py +++ b/model/model_minimind.py @@ -212,10 +212,14 @@ class MiniMindModel(nn.Module): past_key_values = past_key_values or [None] * len(self.layers) start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0 hidden_states = self.dropout(self.embed_tokens(input_ids)) - # Recompute RoPE buffers lost during meta-device init (transformers>=5.x) - if self.freqs_cos[0, 0] == 0: - freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.config.head_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling) - self.freqs_cos, self.freqs_sin = freqs_cos.to(hidden_states.device), freqs_sin.to(hidden_states.device) + # Lazy recompute of RoPE buffers if they were dropped by transformers>=5.x meta-device init. + # Hot path: skipped via cheap python bool (no GPU sync). The original `freqs_cos[0,0]==0` + # tensor compare forced a device-to-host sync on every forward. + if getattr(self, '_rope_ready', False) is False: + if self.freqs_cos.numel() == 0 or bool(self.freqs_cos[0, 0].eq(0).cpu()): + freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.config.head_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling) + self.freqs_cos, self.freqs_sin = freqs_cos.to(hidden_states.device), freqs_sin.to(hidden_states.device) + self._rope_ready = True position_embeddings = (self.freqs_cos[start_pos:start_pos + seq_length], self.freqs_sin[start_pos:start_pos + seq_length]) presents = [] for layer, past_key_value in zip(self.layers, past_key_values): diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index 50ba878..b3f71af 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -23,6 +23,7 @@ warnings.filterwarnings('ignore') def train_epoch(epoch, loader, iters, start_step=0, wandb=None): start_time = time.time() + last_log_time, last_log_step = start_time, start_step last_step = start_step for step, (input_ids, labels) in enumerate(loader, start=start_step + 1): input_ids = input_ids.to(args.device) @@ -49,13 +50,16 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): optimizer.zero_grad(set_to_none=True) if step % args.log_interval == 0 or step == iters: - spend_time = time.time() - start_time + now = time.time() + spend_time = now - start_time + window_its = (step - last_log_step) / max(now - last_log_time, 1e-6) + last_log_time, last_log_step = now, step current_loss = loss.item() * args.accumulation_steps current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0 current_logits_loss = current_loss - current_aux_loss current_lr = optimizer.param_groups[-1]['lr'] eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60 - Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min') + Logger(f'[{time.strftime("%H:%M:%S")}] Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), it/s: {window_its:.2f}, loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min') if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min}) if (step % args.save_interval == 0 or step == iters) and is_main_process(): @@ -103,7 +107,7 @@ if __name__ == "__main__": parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）") parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb") parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain", help="wandb项目名") - parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）") + parser.add_argument("--use_compile", default=1, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是，A100 上推荐默认 1）") args = parser.parse_args() # ========== 1. 初始化环境和随机种子 ========== diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py index b6211d5..96ad3dd 100644 --- a/trainer/trainer_utils.py +++ b/trainer/trainer_utils.py @@ -51,14 +51,18 @@ def init_distributed_mode(): return local_rank -def setup_seed(seed: int): +def setup_seed(seed: int, deterministic: bool = False): + # cudnn 决定性 + 禁 benchmark 是预训练吞吐杀手（A100 上 ~3-5x 慢）。 + # 默认走 benchmark=True 让 cudnn 选最快 kernel；要严格复现传 deterministic=True。 random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = deterministic + torch.backends.cudnn.benchmark = not deterministic + # A100 TF32 矩阵乘加速（autocast bf16 路径之外的 fp32 matmul 才生效，但无副作用） + torch.set_float32_matmul_precision('high') def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoch=0, step=0, wandb=None, save_dir='../checkpoints', **kwargs): os.makedirs(save_dir, exist_ok=True)