Update train.py

train.py

@@ -1,5 +1,5 @@
 import torch
-from transformers import AutoProcessor, AutoModelForVision2Seq, BitsAndBytesConfig
+from transformers import AutoProcessor, AutoModelForImageTextToText, BitsAndBytesConfig
 import datasets
 from datasets import Dataset
 from typing import cast
@@ -18,13 +18,13 @@ def load_model(model_name, device_id=0):
     )
 
     processor = AutoProcessor.from_pretrained(model_name)
+    processor.tokenizer.padding_side = "left"
 
-    model = AutoModelForVision2Seq.from_pretrained(
+    model = AutoModelForImageTextToText.from_pretrained(
         model_name,
         quantization_config=bnb_config,
         dtype=torch.bfloat16,
         device_map={"": device_id},
-        torch_dtype=torch.bfloat16,
         attn_implementation="flash_attention_2",
     )
 
@@ -63,7 +63,7 @@ def caption_batch(batch, processor, model):
 
     inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
-    with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.bfloat16):
+    with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.bfloat16):
         generated = model.generate(
             **inputs,
             max_new_tokens=128,
@@ -91,40 +91,44 @@ def caption_batch(batch, processor, model):
 
 
 def process_shard_worker(
-    gpu_id, start, end, model_name, batch_size, input_dataset, output_file
+    gpu_id, start, end, model_name, batch_size, input_dataset, output_file, error_queue
 ):
-    torch.cuda.set_device(gpu_id)
-
-    print(f"[GPU {gpu_id}] Loading model...", flush=True)
-    processor, model = load_model(model_name, gpu_id)
-
-    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"])
-    else:
-        shard = cast(Dataset, loaded)
-
-    print(f"[GPU {gpu_id}] Processing {len(shard)} examples...", flush=True)
-    result = shard.map(
-        lambda batch: caption_batch(batch, processor, model),
-        batched=True,
-        batch_size=batch_size,
-        remove_columns=[col for col in shard.column_names if col != "image"],
-        writer_batch_size=1000,
-    )
+    try:
+        torch.cuda.set_device(gpu_id)
+
+        print(f"[GPU {gpu_id}] Loading model...", flush=True)
+        processor, model = load_model(model_name, gpu_id)
+
+        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"])
+        else:
+            shard = cast(Dataset, loaded)
+
+        print(f"[GPU {gpu_id}] Processing {len(shard)} examples...", flush=True)
+        result = shard.map(
+            lambda batch: caption_batch(batch, processor, model),
+            batched=True,
+            batch_size=batch_size,
+            remove_columns=[col for col in shard.column_names if col != "image"],
+            writer_batch_size=1000,
+        )
 
-    print(f"[GPU {gpu_id}] Saving results to {output_file}...", flush=True)
-    result.save_to_disk(output_file)
+        print(f"[GPU {gpu_id}] Saving results to {output_file}...", flush=True)
+        result.save_to_disk(output_file)
 
-    print(f"[GPU {gpu_id}] Done!", flush=True)
-    return output_file
+        print(f"[GPU {gpu_id}] Done!", flush=True)
+        return output_file
+    except Exception as e:
+        error_queue.put((gpu_id, e))
+        raise
 
 
 def main():
-    input_dataset = "none-yet/wikiart"
-    output_dataset = "nroggendorff/wikiart"
+    input_dataset = "none-yet/anime-captions"
+    output_dataset = "nroggendorff/anime-captions"
     model_name = "datalab-to/chandra"
     batch_size = 16
 
@@ -148,6 +152,7 @@ def main():
 
     processes = []
     temp_files = []
+    error_queue = mp.Queue()
 
     for i in range(num_gpus):
         start = i * shard_size
@@ -157,13 +162,23 @@ def main():
 
         p = mp.Process(
             target=process_shard_worker,
-            args=(i, start, end, model_name, batch_size, input_dataset, output_file),
+            args=(i, start, end, model_name, batch_size, input_dataset, output_file, error_queue),
         )
         p.start()
         processes.append(p)
 
     for p in processes:
         p.join()
+        if not error_queue.empty():
+            gpu_id, error = error_queue.get()
+            print(f"\n[GPU {gpu_id}] Error occurred: {error}", 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"Process for GPU {gpu_id} failed with error: {error}")
 
     print("\nAll processes completed. Loading and concatenating results...")