Update train.py

train.py

@@ -3,7 +3,7 @@ import torch
 from transformers import AutoProcessor, AutoModelForVision2Seq, BitsAndBytesConfig
 
 
-def load_model(model_name="datalab-to/chandra", device_id=0):
+def load_model(model_name, device_id=0):
     bnb_config = BitsAndBytesConfig(
         load_in_4bit=True,
         bnb_4bit_compute_dtype=torch.bfloat16,
@@ -16,7 +16,7 @@ def load_model(model_name="datalab-to/chandra", device_id=0):
     model = AutoModelForVision2Seq.from_pretrained(
         model_name,
         quantization_config=bnb_config,
-        dtype=torch.bfloat16,
+        torch_dtype=torch.bfloat16,
         device_map={"": device_id},
     )
 
@@ -67,14 +67,16 @@ def caption_batch(batch, processor, model):
         generated = model.generate(
             input_ids=input_ids,
             attention_mask=attention_mask,
+            max_new_tokens=256,
         )
 
-    decoded = processor.batch_decode(generated)
+    decoded = processor.batch_decode(generated, skip_special_tokens=False)
 
     captions = []
     for d in decoded:
         if "<|im_start|>assistant" in d:
             d = d.split("<|im_start|>assistant")[-1].strip()
+        d = d.replace("<|im_end|>", "").strip()
         captions.append(d)
 
     return {
@@ -86,11 +88,38 @@ def caption_batch(batch, processor, model):
 import datasets
 from datasets import Dataset
 from typing import cast
-from concurrent.futures import ThreadPoolExecutor
+import multiprocessing as mp
 
 
+def process_shard_worker(args):
+    _, device_id, start, end, model_name, batch_size = args
+
+    torch.cuda.set_device(device_id)
+
+    processor, model = load_model(model_name, device_id)
+
+    input_dataset = "none-yet/anime-captions"
+    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)
+
+    result = shard.map(
+        lambda batch: caption_batch(batch, processor, model),
+        batched=True,
+        batch_size=batch_size,
+        remove_columns=shard.column_names,
+    )
+
+    return result
+
+mp.set_start_method("spawn", force=True)
+
 input_dataset = "none-yet/anime-captions"
 output_dataset = "nroggendorff/anime-captions"
+model_name = "datalab-to/chandra"
 
 loaded = datasets.load_dataset(input_dataset, split="train")
 
@@ -100,31 +129,18 @@ else:
     ds = cast(Dataset, loaded)
 
 num_gpus = torch.cuda.device_count()
-models = [load_model(device_id=i) for i in range(num_gpus)]
-
 batch_size = 32
-shard_size = len(ds) // num_gpus
-
-
-def process_shard(shard_idx, processor, model):
-    start = shard_idx * shard_size
-    end = start + shard_size if shard_idx < num_gpus - 1 else len(ds)
-    shard = ds.select(range(start, end))
-
-    return shard.map(
-        lambda batch: caption_batch(batch, processor, model),
-        batched=True,
-        batch_size=batch_size,
-        remove_columns=shard.column_names,
-    )
+total_size = len(ds)
+shard_size = total_size // num_gpus
 
+worker_args = []
+for i in range(num_gpus):
+    start = i * shard_size
+    end = start + shard_size if i < num_gpus - 1 else total_size
+    worker_args.append((i, i, start, end, model_name, batch_size))
 
-with ThreadPoolExecutor(max_workers=num_gpus) as executor:
-    futures = [
-        executor.submit(process_shard, i, proc, model)
-        for i, (proc, model) in enumerate(models)
-    ]
-    shards = [f.result() for f in futures]
+with mp.Pool(processes=num_gpus) as pool:
+    shards = pool.map(process_shard_worker, worker_args)
 
 ds = datasets.concatenate_datasets(shards)