Store real inputs Args: real_inputs: The real inputs to store (can be tensors, lists, tuples, etc.)

Recursively get the attribute of a module. Args: model (`torch.nn.Module`) The model to get the attribute from. m

Sets a flag within a module in `model` to instruct ZeRO3 to stop setting hooks recursively when it encounters a module class listed in `leaf_module_cl

Will construct and populate a batch and KVCache with the given sequence parameters. Arguments: seq_params (List[Tuple[int, int]]): A

Construct a simple batch with the given sequence lengths. This method should not be used for for testing scenarios that require information a

Returns nvcc compute capability compile flags. 1. Under ``jit_mode``, the precedence is: a. preserved ``TORCH_CUDA_ARCH_L

Convert HF tp_plan to DeepSpeed layer specs. Returns None if the plan contains any unsupported partition styles, allowing the caller

Compute-capability major of CUDA device 0, or ``None`` when it cannot be read without side effects. ``torch.cuda.get_device_propertie

Return the coordinate owned by a process rank. The axes of the returned namedtuple can be directly accessed as members. For example:

The group of ranks within the same stage of all pipelines.

Return True if the given activation function is gated.

Drop any empty strings from the list of compile and link flags

Assert no recompute params / active consumers remain and params are partitioned.

Inject sequence-parallel wrappers into *model* in-place. Scans the model's named modules and replaces recognised attention layers with their

Backwards-compatible way of querying the tensor model parallel world size. Similar to bwc_tensor_model_parallel_rank.

Calculates the learning rate at batch index. This function treats `self.last_batch_iteration` as the last batch index.

Autograd-aware Triton grouped GEMM (2D x 3D), drop-in for ``torch._grouped_mm``. Args: mat_a: ``[M, K]`` float16/bfloat16/float32.

Decorator that records an NVTX range for the duration of the function call. Skips NVTX instrumentation when torch.compile is active to avoid gr

Perform any necessary transformations of the parameters of this module. Parameters: param (torch.Tensor): Weight or bias

Validate that a DeepCompile pass schedule satisfies the registered pass contracts. ``schedule`` uses the ``[(step, passes), ...]`` format consume

Flatten dense tensors into a contiguous 1D buffer. Assume tensors are of same dense type. Since inputs are dense, the resulting tensor will b

Buffer associated with the batch metadata tensor that can be populated in preparation for passing a new input to the device.

Backwards-compatible way of querying the tensor model parallel rank from an ``mpu`` object. *Tensor* model parallelism means that tensors are

Traverse model and detect MoE layers. Returns list of MoELayerSpec.

Derive TP, dense-DP, EP, and EDP rank tables without process groups.

Return the number of processes along the given axis. For example: >>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3])

Helper to determine the number of local heads of a given shard. Args: shard_rank (int): The rank of the shard. num_shards (i

Return a list of all parameters that would have LoRA for the module.

Buffer associated with the metadata of each sequence in the ragged batch. If the device Tensor is requested, the Tensor is truncated

Query whether the current micro-batch is at the boundary of gradient accumulation, and thus will trigger gradient reductions and

True if this process is in the last stage in the pipeline.