`core.tensor_parallel.random`#

Module Contents#

Classes#

`CudaRNGStatesTracker`	Tracker for the cuda RNG states.
`CheckpointFunction`	Checkpoint Function
`CheckpointWithoutOutputFunction`	Checkpoint Function Helper for CheckpointWithoutOutput. Save context for recompute.
`CheckpointWithoutOutputManager`	Coordinates activation recomputation across multiple CheckpointWithoutOutput instances within a TransformerBlock, enabling unified recomputation during backward pass. This is particularly useful for scenarios where multiple checkpoint operations have sequential dependencies (i.e., the output of one checkpoint is the input of the next).
`CheckpointWithoutOutput`	Checkpoint a model or part of the model and release the output.

Functions#

`_get_share_storage`	Lazily compile & cache the share_storage extension.
`_get_cuda_rng_state`	Return the random number generator state of the specified GPU.
`_set_cuda_rng_state`	Sets the random number generator state of the current GPU.
`convert_cuda_rng_state`	Convert the cuda rng state tensor to the graphable version, or from the graphable version to the non-graphable tensor version.
`get_expert_parallel_rng_tracker_name`	Get the expert parallel rng tracker name
`get_data_parallel_rng_tracker_name`	Get the data parallel rng tracker name
`get_gtp_remat_rng_tracker_name`	Get the (E)GTP_remat weight-init rng tracker name (per-(E)GTP-rank distinct draws).
`initialize_rng_tracker`	Create the RNG tracker. ‘use_te_rng_tracker’ determines whether to use Megatron or TransformerEngine’s implementation. In particular, TransformerEngine’s implementation is cudagraphable and supports FP8.
`get_cuda_rng_tracker`	Get cuda rng tracker.
`get_all_rng_states`	Returns all generator states used by the current `CudaRNGStatesTracker`.
`model_parallel_cuda_manual_seed`	Initialize model parallel cuda seed.
`is_graph_safe_cuda_rng_tracker`	Check if the cuda rng tracker is graph safe version.
`_get_all_rng_states`	Get all the rng states.
`_set_all_rng_states`	Set all the rng states.
`_fork_rng`	Fork the rng state.
`_set_checkpointing`	Set state to checkpointing enabled.
`_unset_checkpointing`	Unset state to checkpointing enabled.
`is_checkpointing`	Check if currently in a checkpoint context.
`checkpoint`	Checkpoint a model or part of the model. This has been directly copied from torch.utils.checkpoint.
`_save_args_to_ctx`	Save mixed tensor/non-tensor arguments into autograd ctx.
`_load_args_from_ctx`	Load and reconstruct mixed tensor/non-tensor arguments from autograd ctx.

Data#

`_SHARE_STORAGE_SRC`
`_share_storage_ext`
`_MODEL_PARALLEL_RNG_TRACKER_NAME`
`_EXPERT_PARALLEL_RNG_TRACKER_NAME`
`_DATA_PARALLEL_RNG_TRACKER_NAME`
`_GTP_REMAT_RNG_TRACKER_NAME`
`_EXPERT_GTP_REMAT_RNG_TRACKER_NAME`
`_CUDA_RNG_STATE_TRACKER`
`_CUDA_RNG_STATE_TRACKER_INITIALIZED`
`IS_CHECKPOINTING`
`_R`
`_Ts`

API#

core.tensor_parallel.random._SHARE_STORAGE_SRC = <Multiline-String>#

core.tensor_parallel.random._share_storage_ext#: None

core.tensor_parallel.random._get_share_storage()#: Lazily compile & cache the share_storage extension.

core.tensor_parallel.random._MODEL_PARALLEL_RNG_TRACKER_NAME#: ‘model-parallel-rng’

core.tensor_parallel.random._EXPERT_PARALLEL_RNG_TRACKER_NAME#: ‘expert-parallel-rng’

core.tensor_parallel.random._DATA_PARALLEL_RNG_TRACKER_NAME#: ‘data-parallel-rng’

core.tensor_parallel.random._GTP_REMAT_RNG_TRACKER_NAME#: ‘gtp-remat-rng’

core.tensor_parallel.random._EXPERT_GTP_REMAT_RNG_TRACKER_NAME#: ‘egtp-remat-rng’

core.tensor_parallel.random._get_cuda_rng_state( device: Union[int, str, torch.device] = 'cuda', clone: bool = False, graph_safe: bool = False, ) → torch.Tensor#

Return the random number generator state of the specified GPU.

Parameters:

device (int) – The gpu to retrieve the rng state
clone (bool) – Whether to also clone the retrieved RNG state
graph_safe (bool) – Get the rng state in a graph safe manner.

This function is adapted from torch.cuda.random.get_rng_state()

core.tensor_parallel.random._set_cuda_rng_state( new_state: torch.Tensor, device: int = -1, graph_safe: bool = False, )#

Sets the random number generator state of the current GPU.

Parameters:

new_state (torch.ByteTensor) – The desired state
device (int) – The gpu to retrieve the rng state
graph_safe (bool) – Set the rng state in a graph safe manner.

This function is adapted from PyTorch repo (torch.cuda.set_rng_state) with a single change: the input state is not cloned. Cloning caused major performance issues for +4 GPU cases.

core.tensor_parallel.random.convert_cuda_rng_state( state: Union[torch.Tensor, torch.Generator], to_graphable: bool = False, ) → Union[torch.Tensor, torch.Generator]#: Convert the cuda rng state tensor to the graphable version, or from the graphable version to the non-graphable tensor version.

core.tensor_parallel.random.get_expert_parallel_rng_tracker_name()#: Get the expert parallel rng tracker name

core.tensor_parallel.random.get_data_parallel_rng_tracker_name()#: Get the data parallel rng tracker name

core.tensor_parallel.random.get_gtp_remat_rng_tracker_name(is_expert=False)#: Get the (E)GTP_remat weight-init rng tracker name (per-(E)GTP-rank distinct draws).

class core.tensor_parallel.random.CudaRNGStatesTracker( use_cudagraphable_rng=False, is_inference_rng_tracker=False, )#

Tracker for the cuda RNG states.

Using the add method, a cuda rng state is initialized based on the input seed and is assigned to name. Later, by forking the rng state, we can perform operations and return to our starting cuda state.

Initialization

is_initialized()#: Checks if the internal RNG state has been set wirth set_states().

reset()#: Set to the initial state (no tracker).

get_states()#: Get rng states. Copy the dictionary so we have direct pointers to the states, not just a pointer to the dictionary.

set_states(states)#: Set the rng states. For efficiency purposes, we do not check the size of seed for compatibility.

add(name, seed)#: Track the rng state.

fork(name=_MODEL_PARALLEL_RNG_TRACKER_NAME)#: Fork the cuda rng state, perform operations, and exit with the original state.

core.tensor_parallel.random._CUDA_RNG_STATE_TRACKER#: None

core.tensor_parallel.random._CUDA_RNG_STATE_TRACKER_INITIALIZED#: False

core.tensor_parallel.random.initialize_rng_tracker( use_te_rng_tracker: bool = False, inference_rng_tracker: bool = False, use_cudagraphable_rng: bool = False, force_reset: bool = False, )#: Create the RNG tracker. ‘use_te_rng_tracker’ determines whether to use Megatron or TransformerEngine’s implementation. In particular, TransformerEngine’s implementation is cudagraphable and supports FP8.

core.tensor_parallel.random.get_cuda_rng_tracker( use_te_rng_tracker: bool = False, inference_rng_tracker: bool = False, use_cudagraphable_rng: bool = False, )#: Get cuda rng tracker.

core.tensor_parallel.random.get_all_rng_states()#: Returns all generator states used by the current CudaRNGStatesTracker.

core.tensor_parallel.random.model_parallel_cuda_manual_seed( seed: int, te_rng_tracker: bool = False, inference_rng_tracker: bool = False, use_cudagraphable_rng: bool = False, tp_rank: Optional[int] = None, ep_rank: Optional[int] = None, etp_rank: Optional[int] = None, gtp_remat_rank: Optional[int] = None, egtp_remat_rank: Optional[int] = None, force_reset_rng: bool = False, gtp_remat_world_size: Optional[int] = None, egtp_remat_world_size: Optional[int] = None, )#

Initialize model parallel cuda seed.

This function should be called after the model parallel is initialized. Also, no torch.cuda.manual_seed should be called after this function. Basically, this is replacement for that function. Three set of RNG states are tracked: default state: This is for data parallelism and is the same among a set of model parallel GPUs but different across different model parallel groups. This is used for example for dropout in the non-tensor-model-parallel regions. tensor-model-parallel state: This state is different among a set of model parallel GPUs, but the same across data parallel groups. This is used for example for dropout in model parallel regions. expert-parallel-seed: This state is only used for the expert layer of MoE models. It is different among expert-tensor and expert-model parallel GPUs, and the same across expert-data parallel groups.

core.tensor_parallel.random.is_graph_safe_cuda_rng_tracker(cuda_rng_tracker)#: Check if the cuda rng tracker is graph safe version.

core.tensor_parallel.random._get_all_rng_states()#: Get all the rng states.

core.tensor_parallel.random._set_all_rng_states( cpu_rng_state, cuda_rng_state, cuda_rng_state_tracker, )#: Set all the rng states.

core.tensor_parallel.random._fork_rng()#: Fork the rng state.

core.tensor_parallel.random.IS_CHECKPOINTING#: False

core.tensor_parallel.random._set_checkpointing()#: Set state to checkpointing enabled.

core.tensor_parallel.random._unset_checkpointing()#: Unset state to checkpointing enabled.

core.tensor_parallel.random.is_checkpointing()#: Check if currently in a checkpoint context.

core.tensor_parallel.random._R#: ‘TypeVar(…)’

core.tensor_parallel.random._Ts#: ‘TypeVarTuple(…)’

class core.tensor_parallel.random.CheckpointFunction#

Bases: torch.autograd.Function

Checkpoint Function

This function is adapted from torch.utils.checkpoint with two main changes:

torch.cuda.set_rng_state is replaced with _set_cuda_rng_state
the states in the model parallel tracker are also properly tracked/set/reset.

static forward( ctx: Any, run_function: collections.abc.Callable[[typing_extensions.Unpack[core.tensor_parallel.random._Ts]], core.tensor_parallel.random._R], distribute_saved_activations: bool, *args: typing_extensions.Unpack[core.tensor_parallel.random._Ts], ) → core.tensor_parallel.random._R#: Forward pass.

static backward(ctx, *args)#: Backward pass.

core.tensor_parallel.random.checkpoint( function: collections.abc.Callable[[typing_extensions.Unpack[core.tensor_parallel.random._Ts]], core.tensor_parallel.random._R], distribute_saved_activations: bool, *args: typing_extensions.Unpack[core.tensor_parallel.random._Ts], ) → core.tensor_parallel.random._R#: Checkpoint a model or part of the model. This has been directly copied from torch.utils.checkpoint.

core.tensor_parallel.random._save_args_to_ctx(ctx, args)#

Save mixed tensor/non-tensor arguments into autograd ctx.

Since save_for_backward only supports tensors, this function separates tensor and non-tensor arguments, saving tensors via save_for_backward and storing non-tensor metadata (indices and values) as ctx attributes.

Use _load_args_from_ctx to reconstruct the original args.

core.tensor_parallel.random._load_args_from_ctx(ctx)#

Load and reconstruct mixed tensor/non-tensor arguments from autograd ctx.

This is the inverse of _save_args_to_ctx. It retrieves tensors from ctx.saved_tensors and merges them with stored non-tensor arguments to reconstruct the original args in their original order.

Returns:: tuple of reconstructed arguments in their original order.

class core.tensor_parallel.random.CheckpointWithoutOutputFunction#

Bases: torch.autograd.Function

Checkpoint Function Helper for CheckpointWithoutOutput. Save context for recompute.

Handles both tensor and non-tensor arguments:

Tensor arguments are saved via save_for_backward
Non-tensor arguments (int, float, bool, None, etc.) are stored separately in ctx attributes and reconstructed during recomputation

static forward( ctx: Any, run_function: collections.abc.Callable[[typing_extensions.Unpack[core.tensor_parallel.random._Ts]], core.tensor_parallel.random._R], checkpoint_without_output_obj: core.tensor_parallel.random.CheckpointWithoutOutput, *args: typing_extensions.Unpack[core.tensor_parallel.random._Ts], ) → core.tensor_parallel.random._R#: Forward pass.

static backward(ctx, *args)#: Backward pass.

class core.tensor_parallel.random.CheckpointWithoutOutputManager#

Coordinates activation recomputation across multiple CheckpointWithoutOutput instances within a TransformerBlock, enabling unified recomputation during backward pass. This is particularly useful for scenarios where multiple checkpoint operations have sequential dependencies (i.e., the output of one checkpoint is the input of the next).

Usage: manager = CheckpointWithoutOutputManager() ckpt_function = CheckpointWithoutOutput(ckpt_manager=manager) ckpt_function.checkpoint(run_function, *args) # other checkpointed operations manager.discard_all_outputs_and_register_unified_recompute(final_output)

Initialization

add_checkpoint(ckpt)#: Add a checkpoint to the manager.

discard_all_outputs_and_register_unified_recompute(hook_tensor)#: Discard all checkpoint outputs to save memory and register unified recompute hook.

_unified_recompute_hook(grad_output)#

class core.tensor_parallel.random.CheckpointWithoutOutput( fp8=False, ckpt_manager=None, retain_input_tensors=False, )#

Bases: object

Checkpoint a model or part of the model and release the output.

For the normal ‘checkpoint` function, the outputs of it may be saved by the following modules for their backward computation. However, the output of the checkpointed function is re-generated at recomputation, so the output store is not technically needed. This method can manually discard the output in the forward pass and restore it by recomputation in the backward pass to reduce the memory usage.

Due to the reason above, to save memory with this method, the caller should make sure that the discarded output tensors are directly saved in the following modules for backward computation.

Initialization

Initialize CheckpointWithoutOutput.

Parameters:

fp8 – Quantization recipe, or a bool. Note that the default fp8=False still evaluates to self.fp8 = True; every caller that constructs CheckpointWithoutOutput() with no arguments therefore takes the TE activation_recompute_forward path. That is long-standing behavior which several selective-recompute modules (“layernorm”, “moe_act”, “gdn_norm_out”) depend on for correct FP8 amax bookkeeping, so do NOT “fix” this to bool(fp8) here — tightening it changes FP8 numerics and needs its own PR with FP8 functional-test evidence.
ckpt_manager – Optional CheckpointWithoutOutputManager instance. When provided, checkpoint() will auto-register to the manager, and discard_output_and_register_recompute() will only discard output without registering individual hooks.
retain_input_tensors – Whether outputs sharing storage with checkpoint inputs should be retained when discarding outputs.

checkpoint( run_function: collections.abc.Callable[[typing_extensions.Unpack[core.tensor_parallel.random._Ts]], core.tensor_parallel.random._R], *args: typing_extensions.Unpack[core.tensor_parallel.random._Ts], ) → core.tensor_parallel.random._R#

Checkpoint function.

If ckpt_manager was provided during initialization, this checkpoint will be automatically registered to the manager after execution.

_recompute(_)#: Used as a hook to recompute the output.

_discard_outputs()#: Release output storage, preserving outputs that alias retained inputs.

discard_output_and_register_recompute(hook_tensor)#

Release the output tensor storages and register the recompute function as a grad hook of the hook_tensor.

Note: the caller should make sure that the output tensors are no longer used in the forward pass and the gradient of the hook_tensor is computed before the recomputed tensors are used.

core.tensor_parallel.random#

Module Contents#

Classes#

Functions#

Data#

API#

`core.tensor_parallel.random`#