Update train.py

train.py

@@ -14,7 +14,7 @@ def load_model(model_name, device_id=0):
         load_in_4bit=True,
         bnb_4bit_compute_dtype=torch.bfloat16,
         bnb_4bit_quant_type="nf4",
-        bnb_4bit_use_double_quant=False,
+        bnb_4bit_use_double_quant=True,
     )
 
     processor = AutoProcessor.from_pretrained(model_name)
@@ -31,16 +31,7 @@ def load_model(model_name, device_id=0):
     return processor, model
 
 
-def caption_batch(batch, processor, model):
-    images = batch["image"]
-
-    pil_images = []
-    for image in images:
-        if isinstance(image, Image.Image):
-            if image.mode != "RGB":
-                image = image.convert("RGB")
-            pil_images.append(image)
-
+def getTemplate(processor):
     msg = [
         {
             "role": "user",
@@ -54,14 +45,26 @@ def caption_batch(batch, processor, model):
         }
     ]
 
-    text = processor.apply_chat_template(
+    return processor.apply_chat_template(
         msg, add_generation_prompt=True, tokenize=False
     )
+
+
+def caption_batch(batch, processor, model, text):
+    images = batch["image"]
+
+    pil_images = []
+    for image in images:
+        if isinstance(image, Image.Image):
+            if image.mode != "RGB":
+                image = image.convert("RGB")
+            pil_images.append(image)
+
     texts = [text] * len(pil_images)
 
     inputs = processor(text=texts, images=pil_images, return_tensors="pt", padding=True)
 
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+    inputs = {k: v.pin_memory().to(model.device, non_blocking=True) for k, v in inputs.items()}
 
     with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.bfloat16):
         generated = model.generate(
@@ -106,7 +109,7 @@ def process_shard(gpu_id, start, end, model_name, batch_size, input_dataset, out
 
         print(f"[GPU {gpu_id}] Processing {len(shard)} examples...", flush=True)
         result = shard.map(
-            lambda batch: caption_batch(batch, processor, model),
+            lambda batch: caption_batch(batch, processor, model, getTemplate(processor)),
             batched=True,
             batch_size=batch_size,
             remove_columns=[col for col in shard.column_names if col != "image"],
@@ -124,11 +127,11 @@ 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"
-    batch_size = 16
+    batch_size = 20
 
     print("Loading dataset info...")
     loaded = datasets.load_dataset(input_dataset, split="train")
@@ -192,4 +195,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()