mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
90 lines
3.3 KiB
90 lines
3.3 KiB
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
import torch
|
|
from torch import Tensor
|
|
|
|
from colossalai.context.parallel_mode import ParallelMode
|
|
|
|
|
|
class SeedManager:
|
|
"""This class is a manager of all random seeds involved in the system.
|
|
|
|
Note:
|
|
The parallel_mode should be concluded in ``ParallelMode``. More details about ``ParallelMode`` could be found
|
|
in `parallel_mode <https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/context/parallel_mode.py>`_.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._current_mode = None
|
|
self._seeds = dict()
|
|
self._seed_states = dict()
|
|
|
|
@property
|
|
def current_mode(self):
|
|
return self._current_mode
|
|
|
|
@property
|
|
def seeds(self):
|
|
return self._seeds
|
|
|
|
@property
|
|
def seed_states(self):
|
|
return self._seed_states
|
|
|
|
def set_state(self, parallel_mode: ParallelMode, state: Tensor):
|
|
"""Sets the state of the seed manager for `parallel_mode`.
|
|
|
|
Args:
|
|
parallel_mode (:class:`colossalai.context.ParallelMode`): The chosen parallel mode.
|
|
state (:class:`torch.Tensor`): the state to be set.
|
|
|
|
Raises:
|
|
AssertionError: Raises an AssertionError if `parallel_mode` is not found in the seed manager.
|
|
"""
|
|
assert parallel_mode in self._seed_states, f'Parallel mode {parallel_mode} is not found in the seed manager'
|
|
self._seed_states[parallel_mode] = state
|
|
|
|
def set_mode(self, parallel_mode: ParallelMode):
|
|
"""Sets the current mode of the seed manager.
|
|
|
|
Args:
|
|
parallel_mode (:class:`colossalai.context.ParallelMode`): The chosen parallel mode.
|
|
"""
|
|
if self.current_mode:
|
|
# save the current state for current mode
|
|
self._seed_states[self._current_mode] = torch.cuda.get_rng_state()
|
|
|
|
# set the new state for new mode
|
|
self._current_mode = parallel_mode
|
|
torch.cuda.set_rng_state(self._seed_states[parallel_mode])
|
|
|
|
def add_seed(self, parallel_mode: ParallelMode, seed: int, overwrtie: bool = False):
|
|
"""Adds a seed to the seed manager for `parallel_mode`.
|
|
|
|
Args:
|
|
parallel_mode (:class:`colossalai.context.ParallelMode`): The chosen parallel mode.
|
|
seed (int): The seed to be added.
|
|
overwrtie (bool, optional): Whether allows to overwrite the seed that has been set already
|
|
|
|
Raises:
|
|
AssertionError: Raises an AssertionError if `parallel_mode` is not an instance of :class:`colossalai.context.ParallelMode`
|
|
or the seed for `parallel_mode` has been added.
|
|
"""
|
|
assert isinstance(parallel_mode, ParallelMode), 'A valid ParallelMode must be provided'
|
|
if overwrtie is False:
|
|
assert parallel_mode not in self._seed_states, f'The seed for {parallel_mode} has been added'
|
|
elif parallel_mode in self._seed_states:
|
|
print(f"Warnning: {parallel_mode} seed has been overwritten.", flush=True)
|
|
|
|
current_state = torch.cuda.get_rng_state()
|
|
torch.cuda.manual_seed(seed)
|
|
self._seed_states[parallel_mode] = torch.cuda.get_rng_state()
|
|
self._seeds[parallel_mode] = seed
|
|
torch.cuda.set_rng_state(current_state)
|
|
|
|
def reset(self):
|
|
self._current_mode = None
|
|
self._seeds = dict()
|
|
self._seed_states = dict()
|