Fork the cuda rng state, perform operations, and exit with the original state.

Get an attribute from a wrapped model

Construct lists suitable for a communicator group along axis ``axis``. Example: >>> topo = Topo(axes=['pipe', 'data', 'model'],

Return a list of the axis names in the ordering of the topology.

Get cuda rng tracker.

Return the communication group with all data-parallel ranks

Gets directory path where lazy files are stored.

Build the model.

get the associated module name from the model based on the key_word provided by users

Get the pipeline model parallel group the caller rank belongs to.

Return world size for the sequence parallel group.

Get rng states. Copy the dictionary so we have direct pointers to the states, not just a pointer to the dictionary.

Get norm of an iterable of parameters. This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and added functionality to handle model

Returns list of include paths, relative to root of deepspeed package (i.e., DeepSpeed/deepspeed)

Set global variables, initialize distributed, and set autoresume and random seeds. `allow_no_cuda` should not be set unless using megatron for

Is an activation to be checkpointed

Check if all non-python dependencies are satisfied to build this op

Load a model checkpoint and return the iteration.

Replaces values in dict d with values in dict u. Args: d (dict): the target dict to overwrite u (dict): the dict containing the v

Set to the initial state (no tracker).

Reset the buffers when the shape of activation and gradient change. For example, for curriculum learning that changes the seqlen of each

Returns list of source files for your op, relative to root of deepspeed package (i.e., DeepSpeed/deepspeed)

If no closure is supplied, :attr:`step` should be called after ``fp16_optimizer_obj.backward(loss)``. :attr:`step` updates th

Compute the trimmed mean of a list of numbers. Args: data (list): List of numbers. trim_percent (float): Percentage of data to tr

A context manager to partition the model parameters during the model construction with MiCS partition strategy. Model states are partitioned

Join the splits of the output into a single result. Args: outputs (List[Any]): The reduced outputs for each output split.

Divide a tensor among the tensor parallel ranks

Make sure the input variable is not None.

Gather tensors and concatinate along the last dimension.

Gather tensors and concatinate along the last dimension.

Gather tensors and concatenate them along a dimension

Get the data parallel group the caller rank belongs to.

Get the expert data parallel group the caller rank belongs to.