repalce set_random_seed with transformers's set_seed

finetune/sft/main.py

@@ -20,6 +20,7 @@
     os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir))
 )
 from constant import SFT
+from transformers import set_seed
 from utils.data.data_utils import create_prompt_dataset
 from utils.ds_utils import get_train_ds_config
 from utils.model.model_utils import create_hf_model
@@ -36,7 +37,6 @@
     load_hf_tokenizer,
     print_rank_0,
     save_hf_format,
-    set_random_seed,
     to_device,
 )
 
@@ -244,7 +244,7 @@ def main():
     )
 
     # If passed along, set the training seed now.
-    set_random_seed(args.seed)
+    set_seed(args.seed)
     torch.distributed.barrier()
 
     # load_hf_tokenizer will get the correct tokenizer and set padding tokens based on the model family

finetune/utils/utils.py

@@ -8,7 +8,7 @@
 import torch
 import torch.nn as nn
 from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
-from transformers import AutoTokenizer, set_seed
+from transformers import AutoTokenizer
 
 
 def print_rank_0(msg, rank=0):
@@ -99,15 +99,6 @@ def save_hf_format(model, tokenizer, args, sub_folder=""):
     copy(os.path.join(source, "tokenization_yi.py"), target)
 
 
-def set_random_seed(seed):
-    if seed is not None:
-        set_seed(seed)
-        random.seed(seed)
-        np.random.seed(seed)
-        torch.manual_seed(seed)
-        torch.cuda.manual_seed_all(seed)
-
-
 def get_all_reduce_mean(tensor):
     torch.distributed.all_reduce(tensor, op=torch.distributed.ReduceOp.SUM)
     tensor = tensor / torch.distributed.get_world_size()