Apply `function` to every Tensor in `value`. Args: functional: The function class to apply. value (Any): Target object to ap

For ZeRO-3, assert every non-persistent param is released after backward. The recompute bug left frozen params gathered after backward, so the re

Ensure cached ``ep_size_N`` rank lists match the requested folding spec.

For ZeRO-3, assert every persistent param stays gathered (AVAILABLE) after backward. Persistent frozen params are added to a recompute owner duri

Queue parameter for optimization in the worker process.

Classify one folded TP/SP gradient as ``sum``, ``average``, or ``skip``. TP means Tensor Parallel and SP means Sequence Parallel. The parallel mo

Return built-in AutoEP preset names in registry order.

Factory: construct the rollout engine specified by ``rollout_cfg.engine``. Args: rollout_cfg: :class:`RolloutConfig` (or any object with

Generates a name from the acronyms of the tuning keys in the config dict. TRAIN_MICRO_BATCH_SIZE_PER_GPU is always included in the tuning keys. A

Return ceil(a / b).

Ensure the Torch internal APIs needed by `count_used_parameters_in_backward` exist.

< Efficiently disables all accesses guarded by mask

Clips gradient norm of an iterable of parameters. This has been adapted from Nvidia megatron. We add norm averaging to consider MoE params wh

Mark a component as completed. This should be called by the relevant setter of a direct property or a ParametrizedList. This method w

Consolidate per-expert checkpoint files into full-expert universal format. For each AutoEP layer, loads all per-expert files, stacks into [E_

Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be loaded with ``torch.load(file)`` + ``load_stat

Return a copy of tensor on specified device. Works on individual tensors, and tensors contained/nested in lists, tuples, and dicts. Param

create a logger Args: name (str): name of the logger level: level of logger Raises: ValueError i

Deletes a key from a dictionary if it exists. Args: t (string): target key to delete d (dict): dictionary to delete from

< CUTLASS_HOST_DEVICE enables all accesses guarded by mask

Enter backward context. Call at the start of backward pass.

Exit backward context. Call at the end of backward pass.

Extend the KV-cache for the sequence. Arguments: new_ids (Union[List[torch.IntTensor], torch.IntTensor]): For each alloc

This method does the following (in order): 1. kick off fetch for parameters in immediately required sub module 2. kick off fetch for n

Pure-PyTorch CPU reference for filling permutation indices.

Return the list of ranks whose coordinates match the provided criteria. Example: >>> X = ProcessTopology(axes=['pipe', 'data', 'm