EasyLM

EasyLM

Modelos de linguagem grande (LLMs) facilitados, EasyLM é uma solução completa para pré-treinamento, ajuste fino, avaliação e veiculação de LLMs em JAX/Flax. EasyLM pode ampliar o treinamento LLM para centenas de aceleradores TPU/GPU aproveitando a funcionalidade pjit do JAX.

Baseado nos transformadores e conjuntos de dados do Hugginface, este repositório fornece uma base de código fácil de usar e personalizar para treinar grandes modelos de linguagem sem a complexidade de muitas outras estruturas.

EasyLM é construído com JAX/Flax. Ao aproveitar o utilitário pjit do JAX, o EasyLM é capaz de treinar modelos grandes que não cabem em um único acelerador, fragmentando os pesos do modelo e treinando dados em vários aceleradores. Atualmente, o EasyLM oferece suporte a vários treinamentos de TPU/GPU em um único host, bem como treinamento de vários hosts em Google Cloud TPU Pods.

Atualmente, os seguintes modelos são suportados:

• LLaMA
• Lhama 2
• Lhama 3

Servidor de discórdia

Estamos administrando uma comunidade não oficial do Discord (não afiliada ao Google) para discussões relacionadas ao treinamento de LLMs em JAX. Siga este link para ingressar no servidor Discord. Temos canais dedicados para vários frameworks LLM baseados em JAX, incluindo EasyLM, JaxSeq, Alpa e Levanter.

Modelos treinados com EasyLM

OpenLLaMA

OpenLLaMA é nossa reprodução licenciada permissivamente do LLaMA, que pode ser usada para fins comerciais. Confira a página principal do projeto aqui. O OpenLLaMA pode servir como substituto dos pesos LLaMA no EasyLM. Consulte a documentação do LLaMA para obter mais detalhes.

Coala

Koala é nosso novo chatbot aprimorado com base no LLaMA. Se você estiver interessado em nosso chatbot Koala, você pode conferir a postagem do blog e a documentação para executá-lo localmente.

Instalação

O método de instalação difere entre hosts GPU e hosts Cloud TPU. A primeira etapa é extrair do GitHub.

git clone https://github.com/young-geng/EasyLM.git
cd EasyLM
export PYTHONPATH= " ${PWD} : $PYTHONPATH " 

Instalando no host GPU

O ambiente GPU pode ser instalado via Anaconda.

conda env create -f scripts/gpu_environment.yml
conda activate EasyLM

Instalação no host do Cloud TPU

A VM host da TPU vem com Python e PIP pré-instalados. Basta executar o script a seguir para configurar o host TPU.

./scripts/tpu_vm_setup.sh

Documentações

As documentações do EasyLM podem ser encontradas no diretório docs.

Referência

Se você achou o EasyLM útil em suas pesquisas ou aplicações, cite o seguinte BibTeX:

 @software{geng2023easylm,
  author = {Geng, Xinyang},
  title = {EasyLM: A Simple And Scalable Training Framework for Large Language Models},
  month = March,
  year = 2023,
  url = {https://github.com/young-geng/EasyLM}
}

Créditos

• A implementação LLaMA é de JAX_llama
• A implementação JAX/Flax GPT-J e RoBERTa são de transformadores
• A maioria dos utilitários JAX são do mlxu
• A base de código é fortemente inspirada em JAXSeq