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:

print message when get_dist_msg() deems it should be logged, see its docstring for details. Use this function instead of `log_dist` when the log

Returns a TensorRef to the operand

Assemble and return the fp32 parameter of a low-precision (e.g., fp16) parameter. Args: param (``torch.nn.Parameter``): A model p

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

Implements Top1Gating on logits.

Progress the pipeline to train the next batch of data. The engine will ingest ``self.train_batch_size()`` total samples collectively across al

param_dtype, if set, must match the enabled fp16/bf16 mode. The optimizer/master-weight/reduction paths derive the model dtype from t

Initialize AutoEP: detect MoE layers, create EP groups, replace with EP-enabled layers.

Extract model type from module config or class name.

Compute per-group m-start (exclusive prefix) and m-size on device (no D2H sync). A single Triton kernel reads ``offs`` and writes both outputs, r

Emulate the EP metadata all-to-alls seen by rank ``ep_rank``. Both the legacy per-rank exchange and the per-(rank, local-expert) exch

True when the module checkpoints activations with the reentrant function. Reentrant checkpointing needs the first stage's inputs to require g

Pure-PyTorch reference: out[rows_g] = a[rows_g] @ b[g].

Sample from logits with temperature and nucleus (top-p) filtering.

Build a routing assignment list whose histogram is exactly ``counts``.

Sets the random number generator state of the current GPU. Arguments: new_state (torch.ByteTensor): The desired state This function i