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.

Create layout tensor for the given sequence length Arguments: seq_len: required: an integer determining number of attention head

Utility for sharding a parameter. This will return the slice of the parameter that should exist on the given shard_rank given the sharding co

Split parameters into different MoE groups for optimizer Args: param_groups (Union[Dict[str, Any], Tuple[Dict[str, Any], ...], List[Dict[

Stores a fragment to memory

Mapping of token to which sequence it belongs to in the ragged batch. If the device Tensor is requested, the Tensor is truncated to t

Cast params to param_dtype; cast buffers only when buffer_dtype is set.