From b389fa72f5458c36c1dfe040269e9aa922041d3f Mon Sep 17 00:00:00 2001
From: Jintao Huang <huangjintao.hjt@alibaba-inc.com>
Date: Tue, 18 Nov 2025 20:01:06 +0800
Subject: [PATCH] fix save

---
 swift/megatron/train/sft.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/swift/megatron/train/sft.py b/swift/megatron/train/sft.py
index b369261482..5d9cb7a423 100644
--- a/swift/megatron/train/sft.py
+++ b/swift/megatron/train/sft.py
@@ -4,6 +4,7 @@
 from typing import List, Optional, Union
 
 import torch
+import torch.distributed as dist
 
 from swift.llm import TEMPLATE_MAPPING
 from swift.llm.train import SwiftSft
@@ -60,6 +61,7 @@ def run(self):
 
         try:
             self.trainer.train(train_dataset, val_dataset, data_collator)
+            dist.barrier()  # Ensure all weights are saved completely
         finally:
             # Visualization
             if is_last_rank():