A class that keeps all conversation history.

The rotary position embeddings from RoFormer_ (Su et. al). A crucial insight from the method is that the query and keys are transformed b

GPT Transformer language model. Args: config (TransformerConfig): Transformer config transformer_layer_spec (ModuleSpec): Specifi

GPT Transformer language model. Args: config (TransformerConfig): Transformer config hybrid_transformer_layer_spec (ModuleSpec):

A class for processing a LLama text dataset

Mamba language model. Args: config (TransformerConfig): Transformer config mamba_stack_spec (ModuleSpec): Specifies the modules t

ChatGLM dataset class.

r""" Hedgehog feature map as introduced in `The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry <https://arxiv.org

Transformer class.

Starcoder dataset class

r""" Simple linear mapping feature map as in `Finetuning Pretrained Transformers into RNNs <https://arxiv.org/abs/2103.13076>`_

A class for processing a Bloom text dataset

r""" Deterministic Parameter-Free Projection (DPFP) feature map in `Linear Transformers Are Secretly Fast Weight Programmers <https://arxiv.or

MLP will take the input with h hidden state, project it to 4*h hidden dimension, perform nonlinear transformation, and project the state

MLP will take the input with h hidden state, project it to 4*h hidden dimension, perform nonlinear transformation, and project the state

Parallel self-attention layer abstract class. Self-attention layer takes input with size [s, b, h] and returns output of the same size.

MLP. MLP will take the input with h hidden state, project it to 4*h hidden dimension, perform nonlinear transformation, and project the s

Transformer class.

Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

Language model embeddings. Args: hidden_size: hidden size vocab_size: vocabulary size max_sequence_length: maximum size o

Implement the scaled dot product attention with softmax. Arguments --------- softmax_scale: The temperature to use for the softmax att

An efficient implementation of the Experts layer using CUTLASS GroupedGEMM. This class is designed to execute multiple experts in parallel, t

An efficient implementation of the Experts layer using CUTLASS GroupedGEMM. This class is designed to execute multiple experts in parallel, t

An efficient implementation of the Experts layer using CUTLASS GroupedGEMM. This class is designed to execute multiple experts in parallel, there

SentencePiece BPE tokenizer with Jieba integration

LLAMA dataset class for mmap format data

MLP will take the input with h hidden state, project it to 4*h hidden dimension, perform nonlinear transformation, and project the state

AllGather Based Token dispatcher.

AllGather Based Token dispatcher.

AlltoAll Based Token dispatcher.

AlltoAll Based Token dispatcher.

Token dispatcher without token dropping.

A single 'no-op' transformer layer. The sole purpose of this layer is for when a standalone embedding layer is used (i.e., args.standalone_em

A single transformer layer. Transformer layer takes input with size [s, b, h] and returns an output of the same size.

Pooler layer. Pool hidden states of a specific token (for example start of the sequence) and add a linear transformation followed by a tanh.

An implementation of the Experts layer using a sequence of MLP layers. This class executes each expert sequentially.

An implementation of the Experts layer using a sequence of MLP layers. This class executes each expert sequentially.

An implementation of the Experts layer using a sequence of MLP layers. This class executes each expert sequentially.

Routes input to one of N MLP "experts"

Route each token to the top-k experts.

Route each token to the top-k experts.

Route each token to the top-k experts.

Transformer class.

Transformer class.

Transformer language model. Args: transformer_hparams: transformer hyperparameters vocab_size: vocabulary size max_sequen

Hardcoded tokenizer.