[feat] shuffle data

trainer/train_distill_reason.py

@@ -199,7 +199,7 @@ if __name__ == "__main__":
         batch_size=args.batch_size,
         pin_memory=True,
         drop_last=False,
-        shuffle=False,
+        shuffle=(train_sampler is None),
         num_workers=args.num_workers,
         sampler=train_sampler
     )

trainer/train_distillation.py

@@ -248,7 +248,7 @@ if __name__ == "__main__":
         batch_size=args.batch_size,
         pin_memory=True,
         drop_last=False,
-        shuffle=False,
+        shuffle=(train_sampler is None),
         num_workers=args.num_workers,
         sampler=train_sampler
     )

trainer/train_dpo.py

@@ -231,7 +231,7 @@ if __name__ == "__main__":
         batch_size=args.batch_size,
         pin_memory=True,
         drop_last=False,
-        shuffle=False,
+        shuffle=(train_sampler is None),
         num_workers=args.num_workers,
         sampler=train_sampler
     )

trainer/train_full_sft.py

@@ -185,7 +185,7 @@ if __name__ == "__main__":
         batch_size=args.batch_size,
         pin_memory=True,
         drop_last=False,
-        shuffle=False,
+        shuffle=(train_sampler is None),
         num_workers=args.num_workers,
         sampler=train_sampler
     )

trainer/train_grpo.py

@@ -299,7 +299,7 @@ if __name__ == "__main__":
     train_ds = RLAIFDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
     train_sampler = DistributedSampler(train_ds) if ddp else None
     train_loader = DataLoader(train_ds, batch_size=args.batch_size, pin_memory=True,
-                              drop_last=False, shuffle=False,
+                              drop_last=False, shuffle=(train_sampler is None),
                               num_workers=args.num_workers, sampler=train_sampler)
 
     optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)

trainer/train_lora.py

@@ -196,7 +196,7 @@ if __name__ == "__main__":
         batch_size=args.batch_size,
         pin_memory=True,
         drop_last=False,
-        shuffle=False,
+        shuffle=(train_sampler is None),
         num_workers=args.num_workers,
         sampler=train_sampler
     )

trainer/train_ppo.py

@@ -344,7 +344,7 @@ if __name__ == "__main__":
     train_ds = RLAIFDataset(args.data_path, tokenizer, max_length=(args.max_seq_len + args.max_gen_len))
     train_sampler = DistributedSampler(train_ds) if ddp else None
     train_loader = DataLoader(train_ds, batch_size=args.batch_size, pin_memory=True,
-                              drop_last=False, shuffle=False,
+                              drop_last=False, shuffle=(train_sampler is None),
                               num_workers=args.num_workers, sampler=train_sampler)
 
     # 初始化优化器

trainer/train_pretrain.py

@@ -183,7 +183,7 @@ if __name__ == "__main__":
         batch_size=args.batch_size,
         pin_memory=True,
         drop_last=False,
-        shuffle=False,
+        shuffle=(train_sampler is None),
         num_workers=args.num_workers,
         sampler=train_sampler
     )

trainer/train_spo.py

@@ -348,7 +348,7 @@ if __name__ == "__main__":
     train_ds = RLAIFDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
     train_sampler = DistributedSampler(train_ds) if ddp else None
     train_loader = DataLoader(train_ds, batch_size=args.batch_size, pin_memory=True,
-                              drop_last=False, shuffle=False,
+                              drop_last=False, shuffle=(train_sampler is None),
                               num_workers=args.num_workers, sampler=train_sampler)
 
     optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)