diff --git a/astra-sim-alibabacloud/astra-sim/system/collective/NcclTreeFlowModel.cc b/astra-sim-alibabacloud/astra-sim/system/collective/NcclTreeFlowModel.cc
index 4c84a873..ad5092f7 100644
--- a/astra-sim-alibabacloud/astra-sim/system/collective/NcclTreeFlowModel.cc
+++ b/astra-sim-alibabacloud/astra-sim/system/collective/NcclTreeFlowModel.cc
@@ -433,9 +433,14 @@ bool NcclTreeFlowModel::iteratable(int channel_id) {
       break;
     }
   }
+  bool should_exit = false;
+  if (all_channel_finished && all_packets_freed) {
+    if (!judge_exit_flag.exchange(true)) {
+      should_exit = true;
+    }
+  }
   cs.ExitSection();
-  if (all_channel_finished == true &&
-      all_packets_freed == true) {
+  if (should_exit) {
     exit();
     return false;
   }