Update train.py

train.py

@@ -5,6 +5,7 @@ from datasets import Dataset
 from typing import cast
 import os
 import shutil
+import multiprocessing as mp
 from torch.utils.data import DataLoader
 from PIL import Image
 from functools import partial
@@ -71,8 +72,8 @@ def process_shard(gpu_id, start, end, model_name, batch_size, input_dataset, out
         print(f"[GPU {gpu_id}] Loading model...", flush=True)
         processor, model = load_model(model_name, gpu_id)
 
-        print(f"[GPU {gpu_id}] Streaming data shard [{start}:{end}]...", flush=True)
-        loaded = datasets.load_dataset(input_dataset, split=f"train[{start}:{end}]", streaming=False)
+        print(f"[GPU {gpu_id}] Loading data shard [{start}:{end}]...", flush=True)
+        loaded = datasets.load_dataset(input_dataset, split=f"train[{start}:{end}]")
 
         if isinstance(loaded, datasets.DatasetDict):
             shard = cast(Dataset, loaded["train"])
@@ -133,6 +134,8 @@ def process_shard(gpu_id, start, end, model_name, batch_size, input_dataset, out
 
 
 def main():
+    mp.set_start_method('spawn', force=True)
+    
     input_dataset = "none-yet/anime-captions"
     output_dataset = "nroggendorff/anime-captions"
     model_name = "datalab-to/chandra"
@@ -154,31 +157,33 @@ def main():
     print(f"Using {num_gpus} GPUs")
     print(f"Shard size: {shard_size}")
 
-    import concurrent.futures
-    
+    processes = []
     temp_files = []
-    
-    with concurrent.futures.ProcessPoolExecutor(max_workers=num_gpus) as executor:
-        futures = []
-        for i in range(num_gpus):
-            start = i * shard_size
-            end = start + shard_size if i < num_gpus - 1 else total_size
-            output_file = f"temp_shard_{i}"
-            temp_files.append(output_file)
-
-            future = executor.submit(
-                process_shard,
-                i, start, end, model_name, batch_size, input_dataset, output_file
-            )
-            futures.append(future)
-
-        for future in concurrent.futures.as_completed(futures):
-            try:
-                future.result()
-            except Exception as e:
-                print(f"Process failed with error: {e}", flush=True)
-                executor.shutdown(wait=False, cancel_futures=True)
-                raise
+
+    for i in range(num_gpus):
+        start = i * shard_size
+        end = start + shard_size if i < num_gpus - 1 else total_size
+        output_file = f"temp_shard_{i}"
+        temp_files.append(output_file)
+
+        p = mp.Process(
+            target=process_shard,
+            args=(i, start, end, model_name, batch_size, input_dataset, output_file),
+        )
+        p.start()
+        processes.append(p)
+
+    for p in processes:
+        p.join()
+        if p.exitcode != 0:
+            print(f"\nProcess failed with exit code {p.exitcode}", flush=True)
+            print("Terminating all processes...", flush=True)
+            for proc in processes:
+                if proc.is_alive():
+                    proc.terminate()
+            for proc in processes:
+                proc.join()
+            raise RuntimeError(f"At least one process failed")
 
     print("\nAll processes completed. Loading and concatenating results...")