`core.tensor_parallel.random`#

Module Contents#

Classes#

`CudaRNGStatesTracker`	Tracker for the cuda RNG states.
`CheckpointFunction`	Checkpoint Function
`_RegionState`	Lifecycle of a recompute region. Transitions only move forward.
`_ArgLayout`	Positional layout of a call: tensors travel through autograd, other values stay here.
`_TERecomputeContext`	Transformer Engine state that the recomputation has to reproduce.
`_InputStash`	Holds the tensor inputs of a region in autograd’s saved-tensor storage.
`_RegionFunction`	Autograd node that stands in for the checkpointed function.
`_RecomputeRegion`	One call of a checkpointed function, plus what is needed to re-run it in backward.
`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`.
`_probe_cudagraph_needs_generator_registration`	Capture an RNG op on an unregistered generator and report whether capture rejects it.
`cudagraph_needs_generator_registration`	Whether generators must be registered with a `torch.cuda.CUDAGraph` before capture.
`prime_cuda_rng_states_for_graph_capture`	Create each CUDA generator’s capture state at the start of a graph capture.
`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.
`_as_tuple`	View a function result as a tuple of outputs.
`_register_rebuild_hook`	Call `rebuild()` once the gradient of `hook_tensor` is computed, before it is used.

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`
`_CUDAGRAPH_NEEDS_GENERATOR_REGISTRATION`
`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._CUDAGRAPH_NEEDS_GENERATOR_REGISTRATION#: None

core.tensor_parallel.random._probe_cudagraph_needs_generator_registration() → bool#: Capture an RNG op on an unregistered generator and report whether capture rejects it.

core.tensor_parallel.random.cudagraph_needs_generator_registration() → bool#

Whether generators must be registered with a torch.cuda.CUDAGraph before capture.

PyTorch with pytorch/pytorch#176753 lazily registers every generator whose Philox state is consumed during capture, and CUDAGraph.register_generator_state() became a deprecated no-op that prints a warning on every call. Skip the explicit registration there: it does nothing, and with one call per layer, per graph and per generator it floods stderr (tens of thousands of lines per rank for dynamic inference with CUDA graphs).

The version string cannot identify that change: 2.14.0a0 nightlies built before it landed still require explicit registration. Builds older than 2.14.0a0 always require it; for newer builds, a small capture probe decides, and its result is cached.

core.tensor_parallel.random.prime_cuda_rng_states_for_graph_capture() → None#

Create each CUDA generator’s capture state at the start of a graph capture.

Call as the first operation inside a torch.cuda.graph capture. PyTorch >= 2.14 (pytorch/pytorch#176753) creates this state lazily on the generator’s first Philox use in the capture and allocates it on the default stream. Mid-capture, that allocation makes the caching allocator record and query events for blocks freed earlier in the capture; if a block’s recorded stream has since joined the capture (e.g. an NCCL stream after a pipeline recv), the query raises cudaErrorCapturedEvent and invalidates the capture (pytorch/pytorch#193982). Priming here, while nothing is pending, restores the pre-2.14 behavior of capture_begin. Each primed generator advances by one tiny draw per replay. No-op outside a capture and on PyTorch < 2.14.

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.

class core.tensor_parallel.random._RegionState(*args, **kwds)#

Bases: enum.Enum

Lifecycle of a recompute region. Transitions only move forward.

Initialization

RECORDED#: ‘auto(…)’

RELEASED#: ‘auto(…)’

REBUILT#: ‘auto(…)’

DONE#: ‘auto(…)’

core.tensor_parallel.random._as_tuple(value: Any) → tuple#: View a function result as a tuple of outputs.

class core.tensor_parallel.random._ArgLayout(args: tuple)#

Positional layout of a call: tensors travel through autograd, other values stay here.

Initialization

tensors_of(args: tuple) → tuple#: Return the tensor arguments in positional order.

rebuild(tensors: tuple) → list#: Return the full argument list with tensors placed back at their positions.

class core.tensor_parallel.random._TERecomputeContext(enabled: bool)#

Transformer Engine state that the recomputation has to reproduce.

The first pass is marked recompute_phase=False so TE keeps its FP8 scaling state for the second pass, which runs under the forward-time fp8_autocast settings and is marked recompute_phase=True.

Initialization

first_pass()#: Context for the forward pass, which does not record a graph.

second_pass()#: Context for the recomputation in the backward pass.

class core.tensor_parallel.random._InputStash#

Bases: torch.autograd.Function

Holds the tensor inputs of a region in autograd’s saved-tensor storage.

Saving through save_for_backward keeps the inputs visible to active saved_tensors_hooks (e.g. fine-grained activation offloading) and to the CUDA graph saved-tensor observer, like inputs saved by any other op. The returned token never takes part in a loss, so backward never runs; the inputs are read back once, through token.grad_fn.saved_tensors, when the region is rebuilt. As for any custom Function, the node (and so the stash) only exists if some input requires grad.

static forward(ctx, *tensors)#: Save the inputs and return an empty token.

static backward(ctx, *grads)#: Never called: the token is not used by any loss.

class core.tensor_parallel.random._RegionFunction#

Bases: torch.autograd.Function

Autograd node that stands in for the checkpointed function.

It keeps its region in ctx.region so that backward can route the output gradients through the recomputed graph.

static forward(ctx, region, *tensors)#: Run the function once without recording a graph.

static backward(ctx, *grad_outputs)#: Back-propagate through the rebuilt graph and return the input gradients.

class core.tensor_parallel.random._RecomputeRegion( function: collections.abc.Callable, args: tuple, use_te_context: bool, retain_input_tensors: bool, )#

One call of a checkpointed function, plus what is needed to re-run it in backward.

The region keeps detach()-ed aliases of the outputs rather than the outputs themselves: an output reaches _RegionFunction through its grad_fn, and that node owns the region, so holding the output would close a reference cycle through the autograd graph.

Initialization

record(args: tuple)#: Run the function under no_grad and stash its tensor inputs for the rebuild.

run_without_graph(tensors: tuple)#: Body of the forward pass, called from _RegionFunction.forward.

release(outputs: tuple) → None#: Free the output storages; tensors and views over them stay valid handles.

rebuild() → None#: Re-run the function with grad enabled and give the released outputs the new data.

_unstash_inputs() → tuple#: Read the stashed inputs back (unpack hooks run here, once) as new graph leaves.

_share_rebuilt_storages(outputs: tuple) → None#: Point every released output storage at the rebuilt data.

backpropagate(grad_outputs: tuple) → tuple#: Route the output gradients through the rebuilt graph; return the input gradients.

core.tensor_parallel.random._register_rebuild_hook( hook_tensor: torch.Tensor, rebuild: collections.abc.Callable[[], None], ) → None#: Call rebuild() once the gradient of hook_tensor is computed, before it is used.

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.

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

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.

When grad is disabled, or no tensor argument requires grad, backward cannot need the outputs, so the function runs as a plain call and later discard calls do nothing.

_release_outputs() → Optional[core.tensor_parallel.random._RecomputeRegion]#

Release the outputs of the latest checkpoint() call and return its region.

From then on the region is owned by its autograd node and by its recompute hooks; the facade only keeps a weak reference (for further hooks), so a region whose backward never runs does not keep its recomputed graph alive through the facade.

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. Calling this again with another hook_tensor adds another trigger; the recomputation runs once, at the first trigger.

core.tensor_parallel.random#

Module Contents#

Classes#

Functions#

Data#

API#

`core.tensor_parallel.random`#