Fix race condition in tokenizer save

deepdelirious · Jan 3, 2025 · 9ca96f8 · 9ca96f8
1 parent 042e743
commit 9ca96f8
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/flux_train.py b/flux_train.py
@@ -216,7 +216,14 @@ def train(args):
     else:
         t5xxl_max_token_length = args.t5xxl_max_token_length
 
-    flux_tokenize_strategy = strategy_flux.FluxTokenizeStrategy(t5xxl_max_token_length)
+    #Multiple writers makes things weird
+    if accelerator.is_local_main_process():
+        flux_tokenize_strategy = strategy_flux.FluxTokenizeStrategy(t5xxl_max_token_length)
+    accelerator.wait_for_everyone()
+    if not accelerator.is_local_main_process():
+        flux_tokenize_strategy = strategy_flux.FluxTokenizeStrategy(t5xxl_max_token_length)
+    accelerator.wait_for_everyone()
+
     strategy_base.TokenizeStrategy.set_strategy(flux_tokenize_strategy)
 
     # load clip_l, t5xxl for caching text encoder outputs