Tune training hyperparameters for L4 GPU

src/config.py

@@ -36,8 +36,8 @@ MAX_INPUT = 512
 MAX_TARGET = 256
 
 # Training Configuration
-BATCH_SIZE = 8
-GRADIENT_ACCUMULATION_STEPS = 8
+BATCH_SIZE = 4
+GRADIENT_ACCUMULATION_STEPS = 4
 LEARNING_RATE = 3e-4
 NUM_EPOCHS = 5
 EVAL_STEPS = 2000

src/train_forward.py

@@ -4,6 +4,7 @@ Trains T5 model to predict products from reactants.
 """
 import sacrebleu
 import numpy as np
+import torch
 from transformers import (
     AutoTokenizer,
     T5ForConditionalGeneration,
@@ -41,6 +42,8 @@ def main():
     model = T5ForConditionalGeneration.from_pretrained(BASE_MODEL)
     model.resize_token_embeddings(len(tokenizer))
     
+    use_bf16 = torch.cuda.is_available() and torch.cuda.get_device_capability(0)[0] >= 8
+
     # Setup training arguments
     print("\nSetting up training arguments...")
     args = Seq2SeqTrainingArguments(
@@ -60,7 +63,10 @@ def main():
         eval_steps=EVAL_STEPS,
         save_steps=SAVE_STEPS,
         report_to=[],
-        fp16=True if os.environ.get("CUDA_VISIBLE_DEVICES") else False,
+        bf16=use_bf16,
+        fp16=not use_bf16,
+        dataloader_num_workers=4,
+        dataloader_pin_memory=True,
         push_to_hub=True,
         hub_model_id=FORWARD_MODEL_NAME,
         hub_strategy="every_save",

src/train_retro.py

@@ -4,6 +4,7 @@ Trains T5 model to predict reactants from products.
 """
 import sacrebleu
 import numpy as np
+import torch
 from transformers import (
     AutoTokenizer,
     T5ForConditionalGeneration,
@@ -41,6 +42,8 @@ def main():
     model = T5ForConditionalGeneration.from_pretrained(BASE_MODEL)
     model.resize_token_embeddings(len(tokenizer))
     
+    use_bf16 = torch.cuda.is_available() and torch.cuda.get_device_capability(0)[0] >= 8
+
     # Setup training arguments
     print("\nSetting up training arguments...")
     args = Seq2SeqTrainingArguments(
@@ -60,7 +63,10 @@ def main():
         eval_steps=EVAL_STEPS,
         save_steps=SAVE_STEPS,
         report_to=[],
-        fp16=True if os.environ.get("CUDA_VISIBLE_DEVICES") else False,
+        bf16=use_bf16,
+        fp16=not use_bf16,
+        dataloader_num_workers=4,
+        dataloader_pin_memory=True,
         push_to_hub=True,
         hub_model_id=RETRO_MODEL_NAME,
         hub_strategy="every_save",