triton transformer

Transformateur en Triton (wip)

Implémentation d'un Transformer, mais entièrement en Triton. Je suis complètement nouveau dans le code de réseau neuronal de niveau inférieur, donc ce référentiel sera principalement une expérience d'apprentissage, l'objectif final étant un transformateur vanille plus rapide et plus efficace à former.

Résultats

Layernorm en avant

Layernorm en avant et en arrière

Softmax en avant et en arrière

Installer

$ pip install triton-transformer

Usage

 import torch
from triton_transformer import Transformer

model = Transformer (
    num_tokens = 256 ,       # vocab size
    max_seq_len = 1024 ,     # maximum sequence length
    dim = 512 ,              # dimension
    depth = 6 ,              # depth
    heads = 8 ,              # number of heads
    dim_head = 64 ,          # dimension per head
    causal = True ,          # autoregressive or not
    attn_dropout = 0.1 ,     # attention dropout
    ff_dropout = 0.1 ,       # feedforward dropout
    use_triton = True       # use this to turn on / off triton
). cuda ()

x = torch . randint ( 0 , 256 , ( 1 , 1024 )). cuda ()
logits = model ( x ) # (1, 1024, 256)

Pour vous entraîner, transmettez simplement les étiquettes avec les labels de mots-clés en avant, et la perte d'entropie croisée sera renvoyée pour le backprop.

ex. BERTE

 import torch
from triton_transformer import Transformer

model = Transformer (
    num_tokens = 20000 ,
    max_seq_len = 512 ,
    dim = 512 ,
    depth = 12 ,
    heads = 8 ,
    dim_head = 64 ,
    use_triton = True
). cuda ()

x = torch . randint ( 0 , 20000 , ( 1 , 512 )). cuda ()
labels = torch . randint ( 0 , 20000 , ( 1 , 512 )). cuda ()
mask = torch . ones ( 1 , 512 ). bool (). cuda ()

loss = model ( x , mask = mask , labels = labels )
loss . backward ()

Test - Formation GPT

$ python train.py

Faire

• softmax
• entropie croisée (en utilisant les opérations Triton)
• couchenorme en avant
• layernorm à l'envers
• multiplication de la matrice par lots + transfert d'actes fusionnés
• optimiser layernorm à l'envers (déterminer combien stocker ou recalculer)
• utiliser l'abandon efficace de la mémoire des didacticiels Triton
• multiplication de la matrice par lots + acte fusionné à l'envers
• attention fusionnée (développer sur softmax)
• utilisez triton matmul pour d'autres projections
• comparer et optimiser
• noyaux conditionnels à l'inférence vs à la formation
• noyau matmul triangulaire efficace pour l'attention causale

Citations

 @article { Tillet2019TritonAI ,
    title   = { Triton: an intermediate language and compiler for tiled neural network computations } ,
    author  = { Philippe Tillet and H. Kung and D. Cox } ,
    journal = { Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages } ,
    year    = { 2019 }
}

 @misc { vaswani2017attention ,
    title   = { Attention Is All You Need } , 
    author  = { Ashish Vaswani and Noam Shazeer and Niki Parmar and Jakob Uszkoreit and Llion Jones and Aidan N. Gomez and Lukasz Kaiser and Illia Polosukhin } ,
    year    = { 2017 } ,
    eprint  = { 1706.03762 } ,
    archivePrefix = { arXiv } ,
    primaryClass = { cs.CL }
}

 @misc { so2021primer ,
    title   = { Primer: Searching for Efficient Transformers for Language Modeling } ,
    author  = { David R. So and Wojciech Mańke and Hanxiao Liu and Zihang Dai and Noam Shazeer and Quoc V. Le } ,
    year    = { 2021 } ,
    eprint  = { 2109.08668 } ,
    archivePrefix = { arXiv } ,
    primaryClass = { cs.LG }
}

 @article { chowdhery2022PaLM ,
  title   = { PaLM: Scaling Language Modeling with Pathways } ,
  author  = { Chowdhery, Aakanksha et al } ,
  year    = { 2022 }
}