"vscode:/vscode.git/clone" did not exist on "887b6f1f3172c7f2fa88315f9896d8384b07a97a"
utils.py 556 Bytes
Newer Older
1
"""Utils for model executor."""
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
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)