Use rng_seed param when creating custom dataset sampler

model/model_training/trainer_rl.py

@@ -3,6 +3,7 @@
 import os
 import random
 from argparse import Namespace
+from typing import Sequence
 
 import numpy as np
 import torch
@@ -21,7 +22,7 @@
 from utils.utils_rl import prepare_tensor
 
 
-def argument_parsing(notebook=False, notebook_args=None, **kwargs):
+def argument_parsing(notebook: bool = False, notebook_args: Sequence[str] | None = None, **kwargs):
     parser = argparse.ArgumentParser()
     parser.add_argument("--configs", nargs="+", required=True)
     parser.add_argument("--local_rank", type=int, default=-1)

model/model_training/trainer_rm.py

@@ -1,7 +1,7 @@
 import argparse
 import logging
 import os
-from typing import Callable, Literal, Optional, Union
+from typing import Callable, Literal, Optional, Sequence, Union
 
 import datasets
 import torch
@@ -128,7 +128,7 @@ def get_train_dataloader(self):
         return dataloader
 
 
-def argument_parsing(notebook=False, notebook_args=None):
+def argument_parsing(notebook: bool = False, notebook_args: Sequence[str] | None = None):
     parser = argparse.ArgumentParser()
     parser.add_argument("--configs", nargs="+", required=True)
     parser.add_argument("--local_rank", type=int, default=-1)

model/model_training/trainer_sft.py

@@ -3,7 +3,7 @@
 import logging
 import os
 from functools import partial
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 
 import datasets
 import torch
@@ -166,7 +166,7 @@ def get_train_dataloader(self):
         return dataloader
 
 
-def argument_parsing(notebook=False, notebook_args=None):
+def argument_parsing(notebook: bool = False, notebook_args: Sequence[str] | None = None):
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--configs",

model/model_training/utils/utils.py

@@ -81,6 +81,7 @@ def __init__(
         self.shuffle = shuffle
         self.rank = rank
         self.world_size = world_size
+        self.epoch = 0
 
         if world_size == 1:
             self.rank = 0
@@ -89,7 +90,7 @@ def __init__(
         self.seed = seed
         self.samples_length = samples_length
 
-    def set_epoch(self, epoch) -> None:
+    def set_epoch(self, epoch: int) -> None:
         self.epoch = epoch
 
     def __len__(self) -> int:
@@ -126,11 +127,12 @@ def __iter__(self):
         return iter(epoch_idx)
 
     @classmethod
-    def build_sampler_from_config(cls, training_conf, datasets: List[Dataset], verbose: bool = False, *args, **kwargs):
+    def build_sampler_from_config(cls, training_conf, datasets: List[Dataset], verbose: bool = False, **kwargs):
         dataset_sizes = [len(x) for x in datasets]
         fractions = get_dataset_fractions(training_conf.datasets, dataset_sizes, verbose)
         dataset_size_per_epoch = [int(size * frac) for size, frac in zip(dataset_sizes, fractions)]
-        return cls(dataset_sizes, dataset_size_per_epoch, *args, **kwargs)
+        seed = training_conf.rng_seed
+        return cls(dataset_sizes=dataset_sizes, dataset_size_per_epoch=dataset_size_per_epoch, seed=seed, **kwargs)
 
 
 def get_dataset_fractions(conf, dataset_sizes: List[int], verbose: bool = False):