polish

FlagOpen · Dec 4, 2024 · d5346d8 · d5346d8
1 parent 3014e91
commit d5346d8
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/flagscale/train/arguments.py b/flagscale/train/arguments.py
@@ -114,6 +114,8 @@ def pre_validate_args(self):
         if self.args.untie_embeddings_and_output_weights == False:
             assert all(hetero_process_meshes_tp[0] == hetero_process_meshes_tp[-1]), \
                 f"if untie_embeddings_and_output_weights is False, the first and last stage should have the same tp degree!"
+            assert self.args.hetero_use_cpu_communication == False, \
+                f"if untie_embeddings_and_output_weights is False, the hetero_use_cpu_communication should be False currently!"
 
         # Virtual parallel size.
         if self.args.enable_hetero: