Set drop last to ensure modulo16 restriction for fp8 (#1189)

* set drop last to ensure modulo16 restriction for fp8 * fix quality * Use all eval samples for non-FP8 case
huggingface · Mar 14, 2023 · 41479fe · 41479fe
1 parent eac5d13
commit 41479fe
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/examples/nlp_example.py b/examples/nlp_example.py
@@ -98,10 +98,14 @@ def collate_fn(examples):
 
     # Instantiate dataloaders.
     train_dataloader = DataLoader(
-        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size, drop_last=True
     )
     eval_dataloader = DataLoader(
-        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+        tokenized_datasets["validation"],
+        shuffle=False,
+        collate_fn=collate_fn,
+        batch_size=EVAL_BATCH_SIZE,
+        drop_last=(accelerator.mixed_precision == "fp8"),
     )
 
     return train_dataloader, eval_dataloader