colossalai.context.random

colossalai.context.random.seed(parallel_mode)

A context for seed switch

Examples:

with seed(ParallelMode.DATA):
    output = F.dropout(input)
colossalai.context.random.set_mode(parallel_mode)

Sets the current mode of the seed manager.

Parameters

parallel_mode (colossalai.context.ParallelMode) – The chosen parallel mode

colossalai.context.random.with_seed(func, parallel_mode)

A function wrapper which executes the function with a specified seed.

Examples:

# use with decorator
@with_seed(ParallelMode.DATA)
def forward(input):
    return F.dropout(input)
out = forward(input)
# OR use it inline
def forward(input):
    return F.dropout(input)
wrapper_forward = with_seed(forward, ParallelMode.DATA)
out = wrapped_forward(input)
colossalai.context.random.add_seed(parallel_mode, seed, overwrite=False)

Adds a seed to the seed manager for parallel_mode.

Parameters

  • parallel_mode (colossalai.context.ParallelMode) – The chosen parallel mode

  • seed (int) – The seed to be added

Raises

AssertionError – Raises an AssertionError if parallel_mode is not an instance of colossalai.context.ParallelMode or the seed for parallel_mode has been added

colossalai.context.random.get_seeds()

Returns the seeds of the seed manager.

Returns

The seeds of the seed manager

Return type

dict

colossalai.context.random.get_states(copy=False)

Returns the seed states of the seed manager.

Returns

The seed states of the seed manager

Return type

dict

colossalai.context.random.get_current_mode()

Returns the current mode of the seed manager.

Returns

The current mode of the seed manager.

Return type

torch.ByteTensor

colossalai.context.random.set_seed_states(parallel_mode, state)

Sets the state of the seed manager for parallel_mode.

Parameters

  • parallel_mode (colossalai.context.ParallelMode) – The chosen parallel mode

  • state (torch.Tensor) – the state to be set

Raises

AssertionError – Raises an AssertionError if parallel_mode is not found in the seed manager