Resolve (param_dtype, buffer_dtype) for the module cast. param_dtype follows the enabled fp16/bf16 mode (None if neither). buffer_dty

Helper for embedding test to limit the number of tests to run. Params: embed_dim (int): Model dimension vocab_size (int): Le

Helper for `release_[component]` methods. Accepts a list of tuples where the first element is the module param that needs to be delet

r"""Functional API that performs AdamW algorithm computation. See :class:`~torch.optim.AdamW` for details.

Adds a pointer offset in units of Element

Backwards-compatible way of querying the pipeline parallel world size.

If all heads require same sparsity layout, it propagate first head layout to all heads Arguments: layout: required: a tensor of

< Efficiently disables all accesses guarded by mask

Count routed tokens per expert. Fast path uses ``torch.bincount`` on the current device. If ``deterministic_safe=True`` and deterministic alg

The number of pipelines.

Get the sequence descriptor for the given sequence id. If the sequence does not exist, then None is returned.

Initialize the DeepSpeed Engine. Arguments: args: an object containing local_rank and deepspeed_config fields. This is option

True if this process is in the first stage in the pipeline.

Pointer to where the list of KV ids associated with a sequence are. If the device Tensor is requested, the Tensor is truncated to the

Loads a fragment from memory

Loads a fragment from memory

Check if this optimizer requires loss scaling for correct backward pass. Returns True if any of the following conditions are met:

Returns a TensorRef to the operand

Set dependency can be used for managing dependencies when a mapping is provided in the class definition for the layer. The dep_name h

Update the model parameters. .. note:: This method will be called internally by ZeRO-Offload. DeepSpeed users should

Not supporting closure.

Store a fragment to memory