Utility method for getting the size of the intermediate dimension of a sharded MLP. Args: intermediate_size (int): The size of the i

Compare two tensors or non-tensor numbers for their closeness. Add msg=blah to add an additional comment to when assert fails. For defa

Helper to create a subclass of ParametrizedList with the desired `count_attr`. In this manner, we can annotate the type of a Parameter depen

Initialize the Sparsity Pattern Config. For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial Arguments:

Helper function for testing blocked flash attention. Used to enable parametrize to only set up a subset of parameters before being passed to

Detect the silent single-rank fallback described in #8084. When a multi-process launcher (``deepspeed``, ``torchrun``, accelerate, ...) sets ``WO

Ensure that numerator is divisible by the denominator.

Returns a dict of name to shape mapping, only for the flattened fp32 weights saved by the optimizer. the names are exactly as in state_dict. T

Helper function. If the op provided is not a torch.dist.ReduceOp, convert it and return

View a flat buffer using the sizes of tensors. Assume that tensors are of same dense type, and that flat is given by _flatten_dense_tensors.

Compare two tensors across all ranks. We want to make sure that all ranks succeed or fail together.

Clears the predicate set efficiently

Recursively scan *model* and return an :class:`SPModelInfo`. The function identifies: * ViT encoder attention layers → wrapped with :class:`U

Clears the predicate set efficiently

Wrap ``module._parameters`` in :class:`ZeROOrderedDict` if not already. PyTorch 2.5+ defaults ``nn.Module._parameters`` to a plain ``dict`` (

Completes construction of the ragged batch by flushing the host buffers to the device.

Get effective partition dimension.

Test for existence of a function within a tuple of libraries. This is used as a smoke test to check whether a certain library is ava

Incrementally insert a sequence into the ragged batch. This will update the metadata for the ragged batch and the sequence.

Parse the 'expert_parallel' section from DS config JSON.

simultaneously reduce-scatter a list of tensors - this can be done more efficiently than individual reduce scatter calls TODO. see if PyTorch

Register "ulysses" attn_implementation with HF transformers and return mpu (Megatron-LM-style parallel state groups object). If seque

Gather TP-partitioned assignment outputs and combine back by token index. The all-gather rebuilds an identical full output on every TP peer, so a

Rollback the optimizer state for a specific subgroup.