r"""Coalesce ZeRO 1/2/3 gradient reduction across multiple engine.backward() calls. One with-block == one optimizer step: every backward insid

Convert HF tp_plan to DeepSpeed layer specs. A style outside SUPPORTED_STYLES raises ValueError and invalidates the whole plan, so a

Extract tp_plan from HuggingFace model. Merge plans from the model config, model class, and runtime instance. HuggingFace may replace the ins

< Efficiently disables all accesses guarded by mask

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

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

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

fp32 config toggling managed_gradient_accumulation; micro-batch 1 so total_samples == micro-batch count.

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

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

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

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.

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.

Return True if the given activation function is gated.

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

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

True if the engine is executing a gradient reduction or optimizer step instruction. This is overridden from :class:`DeepSpeedEngine` to force

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