Note: The FF layer is not auto-scaled when using a GLU type activation. Better do this manually and choose a sensible intermed_size that is nicely

Definitely can represent BERT, but also a lot of other things. To be used for MLM schemes.

Classification head and pooler.

Shampoo preconditioners for some generic layer. NOTE: Does not support sparse gradients at this time. Args: param (Tensor): Paramete

Adagrad/Adam/RMSProp preconditioner for a generic layer. Stores preconditioner using same format as parameter p. Operations are performed in-plac

FLASH as described in Transformer Quality in Linear Time. This is FLASH-QUAD, as we're not too interested in long-range sequences here.

Fourier mixing layer as described in the FNet paper. Layer takes input with size [Batch, Seq, Hidden] and returns output of the same size. Thi

Self-attention layer abstract class. This is a funnel crammed into the gpt neo-x implementation from: https://github.com/EleutherAI/gpt-neox/

:class:`LARS` [LARC in apex] is a pytorch implementation of both the scaling and clipping variants of LARS, in which the ratio between gradie

not again...

A funnel transformer variation. For now only implemented for fixed sequence lengths, but this is not a necessary limitation.

Definitely can represent BERT, but also a lot of other things. To be used for MLM schemes.

Definitely can represent BERT, but also a lot of other things. To be used for MLM schemes.

Definitely can represent BERT, but also a lot of other things. To be used for MLM schemes.

Self-attention layer abstract class. This is the gpt neo-x implementation from: https://github.com/EleutherAI/gpt-neox/blob/main/megatron/mod

Self-attention layer abstract class. Seq-first variant 1 THSA layer takes input with size [Seq, Batch, Hidden] and returns output of the

This is code based on https://github.com/taka-yamakoshi/TelephoneGame/blob/master/model/bert.py.

Modified pretraining for depth-recurrent models. Only works with models that expose the ScriptableRecurrentLM interface.

mini wrapper around BERT attention from huggingface for sanity checks.

Shampoo with blocking applied to the parameters. NOTE: Does not support sparse gradients at this time. Args: param (Tensor): Paramet

Cache a given dataset into RAM or SDRAM (GPU memory). This is only a good idea if you have enough RAM, especially if mapping into SDRAM.

Cross Entropy plus logit regularization via z_loss.

Wrapper for flash MHA.

Fourier mixing layer as described in the FNet paper. Layer takes input with size [Seq, batch, Hidden] and returns output of the same size. Thi

A funnel layer.

*-GLU activation functions. Implementation mostly following megatron

mini wrapper around BERT attention from huggingface for sanity checks.

CrossEntropyLoss, but only on incorrectly classified examples.

Lazy copy-paste from https://gist.github.com/MFreidank/821cc87b012c53fade03b0c7aba13958.

L1 Loss as a drop-in replacement for Cross Entropy Loss. Only for 2dim inputs and 1dim labels This implementation includes a mean reduction in ba

Implement LMDB caching and access. Originally based on https://github.com/pytorch/vision/blob/master/torchvision/datasets/lsun.py and htt

Shorthand for a learnable embedding.

MSE Loss as a drop-in replacement for Cross Entropy Loss. This implementation includes a mean reduction in batch dimension and a 1/num_classes/M

MSE Loss as a drop-in replacement for Cross Entropy Loss. Only for 2dim inputs and 1dim labels This implementation includes a mean reduction in b

The RMS variant of scaling norms.

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

Make big go fast.

Quick and simple scale norm implementation. Do we also need FixNorm (cosine in the last layer)? It's a maybe here: https://github.com/lucidra

Sinusoidal with scaling (see FLASH paper).

Definitely can represent BERT, but also a lot of other things. To be used for MLM schemes.

Classification head and pooler.

Minimal wrapper around pytorch self attention.

Modified sequential class.

Modified sequential class.

r"""Inject some information about the relative or absolute position of the tokens in the sequence. The positional encodings have the same dimensio

Regression directly back to input embedding. Remove the decoding layer if using this loss. As mentioned at https://twitter.com/ChrSzegedy/status/

This class mirrors deepspeed functionality.

A transformer-encoder structure based on the components from above.

A transformer-encoder structure based on the components from above.

Self-attention layer abstract class. This is a simplification of the tuformer implementationfrom https://github.com/xliu1231/fairseq_tuformer

Wrap arguments.

Implements Adahessian algorithm. It has been proposed in `ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning`. Arguments:

A FLASH-quad layer.

base class for a meta optimizer that wraps and modifies an existing pytorch optimizer.

Preconditioner class.

Classification head.

Self-attention layer abstract class. Seq-first variant 2 THSA layer takes input with size [Seq, Batch, Hidden] and returns output of the

Implements Shampoo algorithm. See details in: - https://arxiv.org/pdf/1802.09568.pdf - https://arxiv.org/pdf/2002.09018.pdf If root_

A transformer-encoder structure based on the components from above.