Revert "Use quantized 8-bit Adam optimizer"

This reverts commit a8df9db.

dalm/training/rag_e2e/train_rage2e.py

@@ -48,8 +48,6 @@
     save_model_hook,
 )
 from dalm.utils import load_dataset
-import bitsandbytes as bnb
-
 
 logger = get_logger(__name__)
 
@@ -348,11 +346,7 @@ def train_e2e(
 
 
     logger.info("create optimizer")
-
-    # Avoid OOM errors by using the 8-bit quantized Adam optimizer
-    # optimizer = torch.optim.Adam(rag_model.parameters(), lr=learning_rate)
-    optimizer = bnb.optim.Adam8bit(rag_model.parameters(), lr=learning_rate)
-
+    optimizer = torch.optim.Adam(rag_model.parameters(), lr=learning_rate)
     logger.info("/create optimizer")
 
     # Scheduler and math around the number of training steps.