triton transformer

Transformador em Tritão (wip)

Implementação de um Transformer, mas totalmente em Triton. Sou completamente novo no código de rede neural de nível inferior, então este repositório será principalmente uma experiência de aprendizado, com o objetivo final sendo um transformador básico que seja mais rápido e eficiente de treinar.

Resultados

Layernorm para frente

Layernorm para frente e para trás

Softmax para frente e para trás

Instalar

$ pip install triton-transformer

Uso

 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)

Para treinar, basta passar os rótulos com os labels de palavras-chave adiante, e a perda de entropia cruzada será retornada para backprop.

ex. BERTO

 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 ()

Teste - treinamento GPT

$ python train.py

Pendência

• softmax
• entropia cruzada (usando operações triton)
• norma de camada para frente
• norma de camada para trás
• multiplicação de matriz de lote + ação fundida para frente
• otimizar o padrão de camada para trás (descobrir quanto armazenar versus recalcular)
• use o abandono eficiente de memória dos tutoriais do Triton
• multiplicação de matriz de lote + ação fundida para trás
• atenção fundida (expandir em softmax)
• use triton matmul para outras projeções
• comparar e otimizar
• kernels condicionais à inferência vs treinamento
• kernel matmul triangular eficiente para atenção causal

Citações

 @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 }
}