Return a list of blocks to the free pool. If a single invalid block is provided (i.e., one that is out of range of the allocator or i

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

Get cuda rng tracker.

Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with ``load_state_dict()`` and used for training

Compute total from a list of norms

The group of ranks within the same pipeline.

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

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

The input ids tensor for the ragged batch. If the device Tensor is requested, the Tensor is truncated to the number of tokens in the

Is an activation to be checkpointed

Tag a router/gate parameter as the *replicated* folded family (AVERAGE). This is the ONLY family marker applied on the live forward path today:

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

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

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

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.

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

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