Upload train_zimage_lora.py

train_zimage_lora.py

@@ -371,22 +371,51 @@ def main():
 
     if args.resume_from_checkpoint:
         logger.info(f"Resuming from checkpoint {args.resume_from_checkpoint}")
-        # accelerator.load_state(args.resume_from_checkpoint) # accelerate launch 会自动做，这里不用写
         
-        # 从文件夹名字里提取我们上次训练到了哪一步，例如 "checkpoint-500" -> 500
-        path = os.path.basename(args.resume_from_checkpoint)
-        global_step = int(path.split("-")[1])
-
+        # 【核心修复 1】总是调用 load_state，这样无论是否使用 accelerate launch 都能工作
+        accelerator.load_state(args.resume_from_checkpoint)
+        
+        # 【核心修复 2】更可靠的 global_step 恢复逻辑
+        checkpoint_path = args.resume_from_checkpoint.rstrip('/')
+        
+        # 方法 1：尝试从文件夹名提取
+        folder_name = os.path.basename(checkpoint_path)
+        if folder_name.startswith("checkpoint-") and "-" in folder_name:
+            try:
+                global_step = int(folder_name.split("-")[1])
+            except:
+                pass
+        
+        # 方法 2：如果无法提取，尝试从 saved_state.json 读取（accelerate 保存的）
+        if global_step == 0:
+            saved_state_path = os.path.join(checkpoint_path, "saved_state.json")
+            if os.path.exists(saved_state_path):
+                try:
+                    import json
+                    with open(saved_state_path, "r") as f:
+                        saved_state = json.load(f)
+                        global_step = saved_state.get("global_step", 0)
+                except:
+                    pass
+        
+        # 方法 3：从进度条保存的文件读取（如果有）
+        if global_step == 0:
+            progress_path = os.path.join(checkpoint_path, "progress.json")
+            if os.path.exists(progress_path):
+                try:
+                    import json
+                    with open(progress_path, "r") as f:
+                        progress_data = json.load(f)
+                        global_step = progress_data.get("step", 0)
+                except:
+                    pass
+        
+        logger.info(f"恢复的 Global step: {global_step}")
+        
         # 计算我们应该从哪个 epoch 和哪一步开始
         first_epoch = global_step // num_update_steps_per_epoch
         resume_step = global_step % num_update_steps_per_epoch
-        logger.info(f"Global step restored to {global_step}, resuming from epoch {first_epoch}, step {resume_step}")
-
-    logger.info("***** Running training (Final Fix) *****")
-
-    # 更新进度条，让它从正确的地方开始
-    # progress_bar = tqdm(range(global_step, args.max_train_steps), initial=global_step, desc="Steps", disable=not accelerator.is_local_main_process)
-    # progress_bar.update(global_step)
+        logger.info(f"Resuming from epoch {first_epoch}, step {resume_step}")
     
     for epoch in range(first_epoch, args.num_train_epochs):
         transformer.train()
@@ -474,7 +503,7 @@ def main():
             break
         # ============================
 
-# === 这里的缩进是 4 个空格 ===
+    # === 这里的缩进是 4 个空格 ===
     if accelerator.is_main_process:
         transformer = accelerator.unwrap_model(transformer)
         transformer = transformer.to(torch.float32)