enforce DropLast for the distributed sampler to work around set issues

RWKV · Feb 6, 2024 · bea9014 · bea9014
1 parent e02907a
commit bea9014
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/RWKV-v5/src/data.py b/RWKV-v5/src/data.py
@@ -1603,6 +1603,8 @@ def train_dataloader(self):
             shuffle=self.dataloader_shuffle_training and not self.sort_by_length,
             num_replicas=self.trainer.world_size,
             rank=self.trainer.global_rank,
+            ## This is required due to multi node alignment errors
+            drop_last=True
         )
 
         microbatch_size = 1
@@ -1638,6 +1640,8 @@ def val_dataloader(self):
             shuffle=False, 
             num_replicas=self.trainer.world_size,
             rank=self.trainer.global_rank,
+            ## This is required due to multi node alignment errors
+            drop_last=True
         )
 
         microbatch_size = 1