2023-05-14 22:32:38 -07:00
|
|
|
"""Utils for model executor."""
|
2023-05-09 15:30:12 -07:00
|
|
|
import random
|
|
|
|
|
|
|
|
|
|
import numpy as np
|
|
|
|
|
import torch
|
|
|
|
|
|
|
|
|
|
from cacheflow.model_executor.parallel_utils.parallel_state import model_parallel_is_initialized
|
|
|
|
|
from cacheflow.model_executor.parallel_utils.tensor_parallel import model_parallel_cuda_manual_seed
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def set_random_seed(seed: int) -> None:
|
|
|
|
|
random.seed(seed)
|
|
|
|
|
np.random.seed(seed)
|
|
|
|
|
torch.manual_seed(seed)
|
|
|
|
|
if torch.cuda.is_available():
|
|
|
|
|
torch.cuda.manual_seed_all(seed)
|
|
|
|
|
|
|
|
|
|
if model_parallel_is_initialized():
|
|
|
|
|
model_parallel_cuda_manual_seed(seed)
|
