Check if the value is a loss tensor. Conditions: - The value must be a tensor. - The tensor must have exactly one element. - The tenso

Move tensor on to specified device by changing the storage. Works on individual tensors, and tensors contained/nested in lists, tuples, and d

Split [0, total) into one chunk per thread. Chunk boundaries are rounded up to a multiple of `align` so each slice's AVX/scalar split lines up with th

Record a grad-bearing forward whose fallback gathers must survive until backward.

Recursively set the attribute of a module. Args: model (`torch.nn.Module`) The model to set the attribute in. mod

Instruct DeepSpeed to coordinate ``parameter``'s collection and partitioning in the forward and backward passes of ``module``. This is used w

Run the graph with profiling enabled. args: inputs to the graph. Tensors in the inpusts must be real tensors, not fake tensors. args can cont

Shared driver for the frozen-param + checkpoint regression tests. Each iteration checks: (1) backward runs without CheckpointError, (2) grads mat

Sets a log level in the passed `logger` and its handlers from string. e.g. "info" => `logging.INFO` Args: log_level_str: one of 'debug',

Utility method for sharding an attention output parameter.

Utility method for sharding an MLP 1 parameter. Both biases and weights are supported, as well as for fused weights for MoE. Args:

Utility method for sharding an MLP 2 parameter. Args: param (torch.Tensor): The parameter to shard. shard_rank (int): Which

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

Move the parameter to the given device. Then, if the device is a cuda device, quantize it.

Implements Top2Gating on logits.

Passes a QKV parameter to the underlying implementation for any necessary transformations. Args: param (torch.Te

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

Validate that the leading dim of the shard is of the expected size and aligns with the sharding logic for the attention computation itself.

Set bf16/fp16 in config_dict based on dtype; skip if not supported.

Fully parameterized testing entry point.

Partition config that matches q_proj and o_proj via regex.

Build the restore-facing parameter UC metadata. Restore metadata stays on the parameter object and may include details that are intentionally

Run _replace_module and capture full_name values that match a spec.

Check if the checkpoint process is still alive. Note: Only call this when self.ckpt_process is not None. Some ranks don't have a chec

Divide a tensor among the tensor parallel ranks