StableAdamW optimizer. An AdamW-Adafactor hybrid with learning rate update clipping. This version is modified to only run foreach which has the o

r""" Args: t_warmup (str | Time): Warmup time. t_max (str | Time): The duration of this scheduler. Default = ``"1dur"``. a

Overall BERT model. Args: config: a BertConfig class instance with the configuration to build a new model Inputs: `input_ids

Overall BERT model. Args: config: a BertConfig class instance with the configuration to build a new model Inputs: `input_ids

Base class for masked language models outputs. Args: loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` i

The rotary position embeddings applied directly to unpadded sequences.

Base class for masked language models outputs. Args: loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` i

Construct the embeddings for words, ignoring position. There are no positional embeddings since we use ALiBi and token_type embeddings.

A stack of BERT layers providing the backbone of Mosaic BERT. This module is modeled after the Hugging Face BERT's :class:`~transformers.model.be

Composes the Mosaic BERT attention and FFN blocks into a single layer.

Chains attention, Dropout, and LayerNorm for Mosaic BERT.

Performs multi-headed self attention on a batch of unpadded sequences. If Flash Attention 2 is installed, this module uses Flash Attention to gre

Applies the FFN at the end of each Mosaic BERT layer. Compared to the default BERT architecture, this block replaces :class:`~transformers.model.

Computes the output of the attention layer. This module is modeled after the Hugging Face BERT's :class:`~transformers.model.bert.modeling_be

An IterableDataset that returns token samples for MDSWriter. Returns dicts of {'tokens': bytes} To use data created by this class and writte

Collator wrapper to add sequence_id to batch.

r"""Decays the learning rate according to a cosine decay followed by an inverse square root decay, with optional warmup and cooldown. Specificall

Measure how long it takes to return a batch from the dataloader.

Torchmetric that grabs the precomputed ce_loss value from the model outputs

Torchmetric that grabs the precomputed z_loss value from the model outputs

Torchmetric that computes cross entropy on language modeling outputs using flash_attn's Cross Entropy. Adds metric state variables: sum_l

Construct the embeddings from token embeddings without any positional embeddings.

Bert Model transformer with a sequence classification/regression head. This head is just a linear layer on top of the pooled output. Used for,

Logs the precomputed L1 and L2 gradient norms from StableAdamW

An IterableDataset that returns text samples for MDSWriter. Returns dicts of {'text': bytes}

A dataset class that can read data with raw mds-format (mosaic streaming-format without compression) from local. In comparison with `Streamin

Decays the learning rate according to a 1 - sqrt function, without an initial warmup phase. The learning rate decays from 1 to alpha_f accord

Records the packing efficiency for each batch.

Disable automatic garbage collection and collect garbage at interval. Args: batch_interval (int): Number of batches between calls to gc.c

Generic text dataset using MosaicML's StreamingDataset. Args: tokenizer (Tokenizer): HuggingFace tokenizer to tokenize sample

Bert Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a softmax) e.g. for RocStories/S

Bert Model with a span classification head. This is used for extractive question-answering tasks like SQuAD (a linear layers on top of the hi

Bert Model transformer with a sequence classification/regression head. This head is just a linear layer on top of the pooled output. Used for,

BoolQ.

CB.

COLA.

COPA.

Eurlex multi-label classification.

Encapsulates a fine-tuning job. Tasks should subclass FineTuneJob and implement the get_trainer() method. Args: name (str, optio

Construct the embeddings with absolute positional embeddings.

A FlexBERT attention base class for type hints.

Composes the FlexBERT attention and MLP blocks into a single layer using pre-normalization.

A FlexBERT embeddings base class for type hints.

A FlexBERT base class for type hints.

Bert Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a softmax) e.g. for RocStories/S

Applies the GLU at the end of each FlexBERT layer. Compared to the default BERT architecture, this block replaces :class:`~transformers.model.ber

A FlexBERT Layer base class for type hints.

Applies the MLP at the end of each FlexBERT layer. Compared to the default BERT architecture, this block replaces :class:`~transformers.model.ber

A FlexBERT MLP base class for type hints.

Performs multi-headed self attention on a batch of padded sequences. This module supports two attention implementations: 1. Flash Attention 2

A stack of BERT layers providing the backbone of FlexBERT. This module is modeled after the Hugging Face BERT's :class:`~transformers.model.bert.

Performs multi-headed self attention on a batch of padded sequences. This module supports two attention implementations: 1. Flash Attention 2

Composes the FlexBERT attention and MLP blocks into a single layer using pre-normalization.

Composes the FlexBERT attention and MLP blocks into a single layer using post-normalization.

Composes the FlexBERT attention and MLP blocks into a single layer using pre-normalization.

Performs multi-headed self attention on a batch of padded sequences. This module supports two attention implementations: 1. Flash Attention 2

Performs multi-headed self attention on a batch of padded sequences. This module supports two attention implementations: 1. Flash Attention 2

Applies the GLU at the end of each FlexBERT layer using intermediate_ff computed in parallel of the attention. Compared to the default BERT archi

An abstract class to handle custom weights initialization of modules

Construct the embeddings from token embeddings without any positional embeddings.

Performs multi-headed self attention on a batch of unpadded sequences. If Flash Attention 2 is installed, this module uses Flash Attention to imp

A stack of BERT layers providing the backbone of FlexBERT. This module is modeled after the Hugging Face BERT's :class:`~transformers.model.bert.

Computes the output of the multi-headed self parallel attention on a batch of unpadded sequences If Flash Attention 2 is installed, this module u