User-facing configuration parsed from DS config JSON.

Sets parameters for ZeRO optimizations.

Offload only marked checkpoint inputs. The marker is installed on ``GradientCheckpointingLayer.__call__`` while this context is active and re

A config bundle is a collection of configs that are used to instantiate a model implementation.

Configuration for Automatic Tensor Parallelism. Example usage: config = AutoTPConfig( tp_size=4, layer_specs

Initialize the DeepSpeed Transformer Config. Arguments: hidden_size: The hidden size of the transformer layer intermed

Rollout engine using DeepSpeed hybrid engine. Args: engine: DeepSpeed engine wrapping the model. tokenizer: HuggingFace tokenizer

Lightweight metadata describing what an optimization pass expects and produces. Contracts let DeepCompile validate a pass schedule before it runs

Modules to be parallelized with pipeline parallelism. The key constraint that enables pipeline parallelism is the representation of the forwa

CUDA implementation of gated activation kernel. This kernel assumes that the input tensor has gate and activation values in adjacent channels

Column-parallel linear layer with sub-parameter support. Handles cases where weights contain multiple logical sub-parameters that need t

Unified specification for tensor parallel layer partitioning. This is inspired by Universal Checkpointing's SubparamShape but extended f

A throwaway git repo seeded with BASELINE, committed on ``master``.

Automatic Expert Parallelism: detect and replace MoE layers.

Track eager-only ZeRO-3 gathers and restore partitioned state around compiled forwards.

For float16 configuration

Allocator class for managing which blocks are free/used in the blocked KV-cache. This is a simple allocator that uses a linked list to ke

Config class for DSLinearBase.

Configuration settings for quantization for LoRAOptimizedLinear, QuantizedLinear, and QuantizedParameter Attributes: q_bits (int

Config class for both DSPreLN and DSPostLN.

Initialize the DeepSpeed GPT Transformer Layer.

Configuration for HybridEngineRollout.

Initialize an MoE layer. Arguments: hidden_size (int): the hidden dimension of the model, importantly this is also the input and output d

The DummySequenceDescriptor is an empty object that allows us to perform schedulability checks before formally tracking a sequence.

Quantized parameter class that implements weight quantization. Weights are stored in quantized form on GPUs, and can be dequantized on-the-fl

Container for all the auxiliary Tensors used in the management of a ragged batch. For each Tensor, we maintain a shadow Tensor on the host.

CUDA implementation of top-1 gating. This will perform a softmax on the logits, and return the scale as well as its idx within that expert's

Group of timers. Borrowed from Nvidia Megatron code

CUDA implementation of bias activation kernel. This kernel should be deprecated once we are fusing the bias activation into the linear kernel

Initialize the DeepSpeed Transformer Config. Arguments: batch_size: The maximum batch size used for running the kernel on each GP

InternVL-style splice: replace IMG_CONTEXT token runs with visual tokens. InternVL encodes each image as ``<IMG_START> <IMG_CONTEXT>×N <IMG_END>`

CUDA Kernel implementation that will perform rotary position embeddings on the queries and keys before copying into a blocked KV cache.

Preset configuration for a known MoE model family.

Qwen2-VL-style splice: visual tokens enclosed by vision_start/end tokens. Qwen2-VL wraps each image's visual tokens with a pair of special bounda

Sequence-parallel wrapper for an opaque ViT attention module. Parameters ---------- attn: The original ViT attention layer (any `

Wrapper around the BLAS matmul kernel for FP16/BF16/FP32 for CUDA/RoCM. Performs z = x @ y

CUDA Kernel implementation that will perform rotary position embeddings on the queries and keys before copying into a blocked KV cache.

CUDA Kernel implementation that will perform rotary position embeddings on the queries and keys before copying into a blocked KV cache.

Floating point layer norm kernel for CUDA/RoCM. Performs: z = ln(x)

Dispatches generation forwards to a per-position CUDA graph. Args: forward: The original ``forward``, used both for capture and as the

Sets parameters for DeepSpeed Inference Engine.

Sets parameters for various monitoring methods.

Configuration settings for LoRAOptimizedLinear. Attributes: lora_r (int): LoRA attention dimension, also known as the rank. Defaults

Sets the learning rate of each parameter group according to 1Cycle learning rate policy (1CLR). 1CLR is a variation of the Cyclical Learning R

Optimized version of nn.Linear that adds features such as: * LoRA w. base weight sharding * FP [6,8,12] quantization Arguments:

A model where some parameters use Muon and some use AdamW. This simulates the scenario where: - Hidden layers use Muon (ndim >= 2) -