ColossalAI Documentation

File: README.md Copy Chapter

Colossal-AI

Instantly Run Colossal-AI on Enterprise-Grade GPUs Copy Chapter

Skip the setup. Access a powerful, pre-configured Colossal-AI environment on HPC-AI Cloud.

Train your models and scale your AI workload in one click!

• NVIDIA Blackwell B200s: Experience the next generation of AI performance (See Benchmarks). Now available on cloud from $2.47/hr.
• Cost-Effective H200 Cluster: Get premier performance with on-demand rental from just $1.99/hr.

Get Started Now & Claim Your Free Credits →

Instant Access Top Open Models at Half the Cost Copy Chapter

Skip the hassle. Access powerful, long-context LLMs seamlessly through HPC-AI Model APIs.

Build your AI agents, chatbots, and RAG applications with HPC-AI Model APIs!

• Latest & Greatest Models: Experience state-of-the-art performance with Kimi 2.5, MiniMax 2.5, and GLM 5.1. Perfect for massive 2M+ context windows and complex coding tasks.

• Unbeatable Pricing: Stop overpaying for API endpoints. Get premier inference speed at up to 50% cheaper than OpenRouter.

Get Started Now & Claim Your $4 Free Credits →

Colossal-AI Benchmark Copy Section

To see how these performance gains translate to real-world applications, we conducted a large language model training benchmark using Colossal-AI on Llama-like models. The tests were run on both 8-card and 16-card configurations for 7B and 70B models, respectively.

The results from the Colossal-AI benchmark provide the most practical insight. For the 7B model on 8 cards, the B200 achieved a 50% higher throughput and a significant increase in TFLOPS per GPU. For the 70B model on 16 cards, the B200 again demonstrated a clear advantage, with over 70% higher throughput and TFLOPS per GPU. These numbers show that the B200's performance gains translate directly to faster training times for large-scale models.

Latest News Copy Chapter

• [2025/02] DeepSeek 671B Fine-Tuning Guide Revealed—Unlock the Upgraded DeepSeek Suite with One Click, AI Players Ecstatic!
• [2024/12] The development cost of video generation models has saved by 50%! Open-source solutions are now available with H200 GPU vouchers [code] [vouchers]
• [2024/10] How to build a low-cost Sora-like app? Solutions for you
• [2024/09] Singapore Startup HPC-AI Tech Secures 50 Million USD in Series A Funding to Build the Video Generation AI Model and GPU Platform
• [2024/09] Reducing AI Large Model Training Costs by 30% Requires Just a Single Line of Code From FP8 Mixed Precision Training Upgrades
• [2024/06] Open-Sora Continues Open Source: Generate Any 16-Second 720p HD Video with One Click, Model Weights Ready to Use
• [2024/05] Large AI Models Inference Speed Doubled, Colossal-Inference Open Source Release
• [2024/04] Open-Sora Unveils Major Upgrade: Embracing Open Source with Single-Shot 16-Second Video Generation and 720p Resolution
• [2024/04] Most cost-effective solutions for inference, fine-tuning and pretraining, tailored to LLaMA3 series

Table of Contents Copy Chapter

<li>[Open-Sora: Revealing Complete Model Parameters, Training Details, and Everything for Sora-like Video Generation Models](#Open-Sora)</li>
 <li>[Colossal-LLaMA-2: One Half-Day of Training Using a Few Hundred Dollars Yields Similar Results to Mainstream Large Models, Open-Source and Commercial-Free Domain-Specific Llm Solution](#Colossal-LLaMA-2)</li>
 <li>[ColossalChat: An Open-Source Solution for Cloning ChatGPT With a Complete RLHF Pipeline](#ColossalChat)</li>
 <li>[AIGC: Acceleration of Stable Diffusion](#AIGC)</li>
 <li>[Biomedicine: Acceleration of AlphaFold Protein Structure](#Biomedicine)</li>

<li>[LLaMA 1/2/3 ](#LLaMA3)</li>
 <li>[MoE](#MoE)</li>
 <li>[GPT-3](#GPT-3)</li>
 <li>[GPT-2](#GPT-2)</li>
 <li>[BERT](#BERT)</li>
 <li>[PaLM](#PaLM)</li>
 <li>[OPT](#OPT)</li>
 <li>[ViT](#ViT)</li>
 <li>[Recommendation System Models](#Recommendation-System-Models)</li>

<li>[GPT-2](#GPT-2-Single)</li>
 <li>[PaLM](#PaLM-Single)</li>

<li>[Colossal-Inference: Large AI  Models Inference Speed Doubled](#Colossal-Inference)</li>
 <li>[Grok-1: 314B model of PyTorch + HuggingFace Inference](#Grok-1)</li>
 <li>[SwiftInfer:Breaks the Length Limit of LLM for Multi-Round Conversations with 46% Acceleration](#SwiftInfer)</li>

<li>[PyPI](#PyPI)</li>
 <li>[Install From Source](#Install-From-Source)</li>

Why Colossal-AI Copy Chapter

(back to top)

Features Copy Chapter

Colossal-AI provides a collection of parallel components for you. We aim to support you to write your
distributed deep learning models just like how you write your model on your laptop. We provide user-friendly tools to kickstart
distributed training and inference in a few lines.

• Parallelism strategies

  • Data Parallelism
  • Pipeline Parallelism
  • 1D, 2D, 2.5D, 3D Tensor Parallelism
  • Sequence Parallelism
  • Zero Redundancy Optimizer (ZeRO)
  • Auto-Parallelism

• Heterogeneous Memory Management

  • PatrickStar

• Friendly Usage

  • Parallelism based on the configuration file

(back to top)

Colossal-AI in the Real World Copy Chapter

Open-Sora Copy Section

Open-Sora：Revealing Complete Model Parameters, Training Details, and Everything for Sora-like Video Generation Models
[code]
[blog]
[Model weights]
[Demo]
[GPU Cloud Playground]
[OpenSora Image]

(back to top)

Colossal-LLaMA-2 Copy Section

[GPU Cloud Playground]
[LLaMA3 Image]

• 7B: One half-day of training using a few hundred dollars yields similar results to mainstream large models, open-source and commercial-free domain-specific LLM solution.
  [code]
  [blog]
  [HuggingFace model weights]
  [Modelscope model weights]

• 13B: Construct refined 13B private model with just $5000 USD.
  [code]
  [blog]
  [HuggingFace model weights]
  [Modelscope model weights]

ColossalChat Copy Section

ColossalChat: An open-source solution for cloning ChatGPT with a complete RLHF pipeline.
[code]
[blog]
[demo]
[tutorial]

• Up to 10 times faster for RLHF PPO Stage3 Training

• Up to 7.73 times faster for single server training and 1.42 times faster for single-GPU inference

• Up to 10.3x growth in model capacity on one GPU

• A mini demo training process requires only 1.62GB of GPU memory (any consumer-grade GPU)

• Increase the capacity of the fine-tuning model by up to 3.7 times on a single GPU

• Keep at a sufficiently high running speed

(back to top)

AIGC Copy Section

Acceleration of AIGC (AI-Generated Content) models such as Stable Diffusion v1 and Stable Diffusion v2.

• Training: Reduce Stable Diffusion memory consumption by up to 5.6x and hardware cost by up to 46x (from A100 to RTX3060).

• DreamBooth Fine-tuning: Personalize your model using just 3-5 images of the desired subject.

• Inference: Reduce inference GPU memory consumption by 2.5x.

(back to top)

Biomedicine Copy Section

Acceleration of AlphaFold Protein Structure

• FastFold: Accelerating training and inference on GPU Clusters, faster data processing, inference sequence containing more than 10000 residues.

• FastFold with Intel: 3x inference acceleration and 39% cost reduce.

• xTrimoMultimer: accelerating structure prediction of protein monomers and multimer by 11x.

(back to top)

Parallel Training Demo Copy Chapter

LLaMA3 Copy Section

• 70 billion parameter LLaMA3 model training accelerated by 18%
  [code]
  [GPU Cloud Playground]
  [LLaMA3 Image]

LLaMA2 Copy Section

• 70 billion parameter LLaMA2 model training accelerated by 195%
  [code]
  [blog]

LLaMA1 Copy Section

• 65-billion-parameter large model pretraining accelerated by 38%
  [code]
  [blog]

MoE Copy Section

• Enhanced MoE parallelism, Open-source MoE model training can be 9 times more efficient
  [code]
  [blog]

GPT-3 Copy Section

• Save 50% GPU resources and 10.7% acceleration

GPT-2 Copy Section

• 11x lower GPU memory consumption, and superlinear scaling efficiency with Tensor Parallelism

• 24x larger model size on the same hardware

• over 3x acceleration

BERT Copy Section

• 2x faster training, or 50% longer sequence length

PaLM Copy Section

• PaLM-colossalai: Scalable implementation of Google's Pathways Language Model (PaLM).

OPT Copy Section

• Open Pretrained Transformer (OPT), a 175-Billion parameter AI language model released by Meta, which stimulates AI programmers to perform various downstream tasks and application deployments because of public pre-trained model weights.
• 45% speedup fine-tuning OPT at low cost in lines. [Example] [Online Serving]

Please visit our documentation and examples for more details.

ViT Copy Section

• 14x larger batch size, and 5x faster training for Tensor Parallelism = 64

Recommendation System Models Copy Section

• Cached Embedding, utilize software cache to train larger embedding tables with a smaller GPU memory budget.

(back to top)

Single GPU Training Demo Copy Chapter

GPT-2 Copy Section

• 20x larger model size on the same hardware

• 120x larger model size on the same hardware (RTX 3080)

PaLM Copy Section

• 34x larger model size on the same hardware

(back to top)

Inference Copy Chapter

Colossal-Inference Copy Section

• Large AI models inference speed doubled, compared to the offline inference performance of vLLM in some cases.
  [code]
  [blog]
  [GPU Cloud Playground]
  [LLaMA3 Image]

Grok-1 Copy Section

• 314 Billion Parameter Grok-1 Inference Accelerated by 3.8x, an easy-to-use Python + PyTorch + HuggingFace version for Inference.

[code]
[blog]
[HuggingFace Grok-1 PyTorch model weights]
[ModelScope Grok-1 PyTorch model weights]

SwiftInfer Copy Section

• SwiftInfer: Inference performance improved by 46%, open source solution breaks the length limit of LLM for multi-round conversations

(back to top)

Installation Copy Chapter

Requirements:

• PyTorch >= 2.2
• Python >= 3.10 and < 3.13
• CUDA >= 11.0
• NVIDIA GPU Compute Capability >= 7.0 (V100/RTX20 and higher)
• Linux OS

If you encounter any problem with installation, you may want to raise an issue in this repository.

Install from PyPI Copy Section

You can easily install Colossal-AI with the following command. By default, we do not build PyTorch extensions during installation.

pip install colossalai

Note: only Linux is supported for now.

However, if you want to build the PyTorch extensions during installation, you can set BUILD_EXT=1.

BUILD_EXT=1 pip install colossalai

Otherwise, CUDA kernels will be built during runtime when you actually need them.

We also keep releasing the nightly version to PyPI every week. This allows you to access the unreleased features and bug fixes in the main branch.
Installation can be made via

pip install colossalai-nightly

Download From Source Copy Section

The version of Colossal-AI will be in line with the main branch of the repository. Feel free to raise an issue if you encounter any problems. :)

git clone https://github.com/hpcaitech/ColossalAI.git
cd ColossalAI

# install colossalai
pip install .

By default, we do not compile CUDA/C++ kernels. ColossalAI will build them during runtime.
If you want to install and enable CUDA kernel fusion (compulsory installation when using fused optimizer):

BUILD_EXT=1 pip install .

For Users with CUDA 10.2, you can still build ColossalAI from source. However, you need to manually download the cub library and copy it to the corresponding directory.

# clone the repository
git clone https://github.com/hpcaitech/ColossalAI.git
cd ColossalAI

# download the cub library
wget https://github.com/NVIDIA/cub/archive/refs/tags/1.8.0.zip
unzip 1.8.0.zip
cp -r cub-1.8.0/cub/ colossalai/kernel/cuda_native/csrc/kernels/include/

# install
BUILD_EXT=1 pip install .

(back to top)

Use Docker Copy Chapter

Pull from DockerHub Copy Section

You can directly pull the docker image from our DockerHub page. The image is automatically uploaded upon release.

Build On Your Own Copy Section

Run the following command to build a docker image from Dockerfile provided.

Building Colossal-AI from scratch requires GPU support, you need to use Nvidia Docker Runtime as the default when doing docker build. More details can be found here.
We recommend you install Colossal-AI from our project page directly.

cd ColossalAI
docker build -t colossalai ./docker

Run the following command to start the docker container in interactive mode.

docker run -ti --gpus all --rm --ipc=host colossalai bash

(back to top)

Community Copy Chapter

Join the Colossal-AI community on Forum,
Slack,
and WeChat(微信) to share your suggestions, feedback, and questions with our engineering team.

Contributing Copy Chapter

Referring to the successful attempts of BLOOM and Stable Diffusion, any and all developers and partners with computing powers, datasets, models are welcome to join and build the Colossal-AI community, making efforts towards the era of big AI models!

You may contact us or participate in the following ways:

1. Leaving a Star ⭐ to show your like and support. Thanks!
2. Posting an issue, or submitting a PR on GitHub follow the guideline in Contributing
3. Send your official proposal to email [email protected]

Thanks so much to all of our amazing contributors!

(back to top)

CI/CD Copy Chapter

We leverage the power of GitHub Actions to automate our development, release and deployment workflows. Please check out this documentation on how the automated workflows are operated.

Cite Us Copy Chapter

This project is inspired by some related projects (some by our team and some by other organizations). We would like to credit these amazing projects as listed in the Reference List.

To cite this project, you can use the following BibTeX citation.

@inproceedings{10.1145/3605573.3605613,
author = {Li, Shenggui and Liu, Hongxin and Bian, Zhengda and Fang, Jiarui and Huang, Haichen and Liu, Yuliang and Wang, Boxiang and You, Yang},
title = {Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training},
year = {2023},
isbn = {9798400708435},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
url = {https://doi.org/10.1145/3605573.3605613},
doi = {10.1145/3605573.3605613},
abstract = {The success of Transformer models has pushed the deep learning model scale to billions of parameters, but the memory limitation of a single GPU has led to an urgent need for training on multi-GPU clusters. However, the best practice for choosing the optimal parallel strategy is still lacking, as it requires domain expertise in both deep learning and parallel computing. The Colossal-AI system addressed the above challenge by introducing a unified interface to scale your sequential code of model training to distributed environments. It supports parallel training methods such as data, pipeline, tensor, and sequence parallelism and is integrated with heterogeneous training and zero redundancy optimizer. Compared to the baseline system, Colossal-AI can achieve up to 2.76 times training speedup on large-scale models.},
booktitle = {Proceedings of the 52nd International Conference on Parallel Processing},
pages = {766–775},
numpages = {10},
keywords = {datasets, gaze detection, text tagging, neural networks},
location = {Salt Lake City, UT, USA},
series = {ICPP '23}
}

Colossal-AI has been accepted as official tutorial by top conferences NeurIPS, SC, AAAI,
PPoPP, CVPR, ISC, NVIDIA GTC ,etc.

(back to top)

File: applications/ColossalQA/colossalqa/prompt/README.md Copy Chapter

Prompt Design Guide

For the retriever conversation system, users can customize three prompts.

The Retrieval QA Prompt Copy Chapter

This is the prompt for retrieval QA, the input is user's inputs, the retrieved documents, the historical conversation.

Chinese Copy Section

你是一个善于解答用户问题的AI助手。在保证安全的前提下，回答问题要尽可能有帮助。你的答案不应该包含任何有害的、不道德的、种族主义的、性别歧视的、危险的或非法的内容。请确保你的回答是公正和积极的。
如果不能根据给定的上下文推断出答案，请不要分享虚假、不确定的信息。
使用提供的背景信息和聊天记录对用户的输入作出回应或继续对话。您应该只生成一个回复。不需要跟进回答。请使用中文作答。

背景信息:
[retrieved documents]

聊天记录:
[historical conversation, overlength chat history will be summarized]

用户: [question]
Assistant:

English Copy Section

[INST] <<SYS>>Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
If the answer cannot be inferred based on the given context, please don't share false information.<</SYS>>
Use the context and chat history to respond to the human's input at the end or carry on the conversation. You should generate one response only. No following up is needed.

context:
[retrieved documents]

chat history
[historical conversation, overlength chat history will be summarized]

Human: {question}
Assistant:

Summarization Prompt Copy Chapter

This prompt is used by the memory module to recursively summarize overlength conversation to shrink the length of the prompt.

Disambiguity Prompt Copy Chapter

This prompt is used to perform zero-shot reference resolution to disambiguate entity references within user's questions.

Final Prompt Examples Copy Chapter

Assume k=3 for the retriever.

English Copy Section

Note that the "[INST] <>...<>" template is the specific prompt format used in LLaMA2.

Normal Length

[INST] <<SYS>>Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
If the answer cannot be inferred based on the given context, please don't share false information.<</SYS>>
Use the context and chat history to respond to the human's input at the end or carry on the conversation. You should generate one response only. No following up is needed.

context:
[document 1]

[document 2]

[document 3]

chat history
Human: XXX
Assistant: XXX
...

Human: {question}
Assistant:

Overlength

[INST] <<SYS>>Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
If the answer cannot be inferred based on the given context, please don't share false information.<</SYS>>
Use the context and chat history to respond to the human's input at the end or carry on the conversation. You should generate one response only. No following up is needed.

context:
[document 1]

[document 2]

[document 3]

chat history
A summarization of historical conversation:
[one line summary of historical conversation]
Most recent conversation:
Human: XXX
Assistant: XXX
...

Human: {question}
Assistant:

Chinese Copy Section

Normal Length

你是一个善于解答用户问题的AI助手。在保证安全的前提下，回答问题要尽可能有帮助。你的答案不应该包含任何有害的、不道德的、种族主义的、性别歧视的、危险的或非法的内容。请确保你的回答是公正和积极的。
如果不能根据给定的上下文推断出答案，请不要分享虚假、不确定的信息。
使用提供的背景信息和聊天记录对用户的输入作出回应或继续对话。您应该只生成一个回复。不需要跟进回答。请使用中文作答。

背景信息:
[document 1]

[document 2]

[document 3]

聊天记录:
用户: XXX
Assistant: XXX
...

用户: [question]
Assistant:

Overlength

你是一个善于解答用户问题的AI助手。在保证安全的前提下，回答问题要尽可能有帮助。你的答案不应该包含任何有害的、不道德的、种族主义的、性别歧视的、危险的或非法的内容。请确保你的回答是公正和积极的。
如果不能根据给定的上下文推断出答案，请不要分享虚假、不确定的信息。
使用提供的背景信息和聊天记录对用户的输入作出回应或继续对话。您应该只生成一个回复。不需要跟进回答。请使用中文作答。

背景信息:
[document 1]

[document 2]

[document 3]

聊天记录:
历史对话概要:
[one line summary of historical conversation]
最近的对话:
用户: XXX
Assistant: XXX
...

用户: [question]
Assistant:

File: applications/ColossalQA/examples/webui_demo/README.md Copy Chapter

ColossalQA WebUI Demo

This demo provides a simple WebUI for ColossalQA, enabling you to upload your files as a knowledge base and interact with them through a chat interface in your browser.

The server.py initializes the backend RAG chain that can be backed by various language models (e.g., ChatGPT, Huawei Pangu, ChatGLM2). Meanwhile, webui.py launches a Gradio-supported chatbot interface.

Usage

Installation Copy Chapter

First, install the necessary dependencies for ColossalQA:

git clone https://github.com/hpcaitech/ColossalAI.git
cd ColossalAI/applications/ColossalQA/
pip install -e .

Install the dependencies for ColossalQA webui demo:

pip install -r requirements.txt

Configure the RAG Chain Copy Chapter

Customize the RAG Chain settings, such as the embedding model (default: moka-ai/m3e), the language model, and the prompts, in the config.py. Please refer to Prepare configuration file for the details of config.py.

For API-based language models (like ChatGPT or Huawei Pangu), provide your API key for authentication. For locally-run models, indicate the path to the model's checkpoint file.

Prepare configuration file Copy Section

All configs are defined in ColossalQA/examples/webui_demo/config.py. You can primarily modify the bolded sections in the config to switch the embedding model and the large model loaded by the backend. Other parameters can be left as default or adjusted based on your specific requirements.

• embed:
  • embed_name: the embedding model name
  • embed_model_name_or_path: path to embedding model, could be a local path or a huggingface path
  • embed_model_device: device to load the embedding model
• model:
  • mode: "local" for loading models, "api" for using model api
  • model_name: "chatgpt_api", "pangu_api", or your local model name
  • model_path: path to the model, could be a local path or a huggingface path. don't need if mode="api"
  • device: device to load the LLM
• splitter:
  • name: text splitter class name, the class should be imported at the beginning of config.py
• retrieval:
  • retri_top_k: number of retrieval text which will be provided to the model
  • retri_kb_file_path: path to store database files
  • verbose: Boolean type, to control the level of detail in program output
• chain:
  • mem_summary_prompt: summary prompt template
  • mem_human_prefix: human prefix for prompt
  • mem_ai_prefix: AI assistant prefix for prompt
  • mem_max_tokens: max tokens for history information
  • mem_llm_kwargs: model's generation kwargs for summarizing history
    • max_new_tokens: int
    • temperature: int
    • do_sample: bool
  • disambig_prompt: disambiguate prompt template
  • disambig_llm_kwargs: model's generation kwargs for disambiguating user's input
    • max_new_tokens: int
    • temperature: int
    • do_sample: bool
  • gen_llm_kwargs: model's generation kwargs
    • max_new_tokens: int
    • temperature: int
    • do_sample: bool
  • gen_qa_prompt: generation prompt template
  • verbose: Boolean type, to control the level of detail in program output

Run WebUI Demo Copy Chapter

Execute the following command to start the demo:

1. If you want to use a local model as the backend model, you need to specify the model name and model path in config.py and run the following commands.

export TMP="path/to/store/tmp/files"
# start the backend server
python server.py --http_host "host" --http_port "port"

# in an another terminal, start the ui
python webui.py --http_host "your-backend-api-host" --http_port "your-backend-api-port"

2. If you want to use chatgpt api as the backend model, you need to change the model mode to "api", change the model name to "chatgpt_api" in config.py, and run the following commands.

export TMP="path/to/store/tmp/files"

# Auth info for OpenAI API
export OPENAI_API_KEY="YOUR_OPENAI_API_KEY"

# start the backend server
python server.py --http_host "host" --http_port "port"

# in an another terminal, start the ui
python webui.py --http_host "your-backend-api-host" --http_port "your-backend-api-port"

3. If you want to use pangu api as the backend model, you need to change the model mode to "api", change the model name to "pangu_api" in config.py, and run the following commands.

export TMP="path/to/store/tmp/files"

# Auth info for Pangu API
export URL=""
export USERNAME=""
export PASSWORD=""
export DOMAIN_NAME=""

# start the backend server
python server.py --http_host "host" --http_port "port"

# in an another terminal, start the ui
python webui.py --http_host "your-backend-api-host" --http_port "your-backend-api-port"

After launching the script, you can upload files and engage with the chatbot through your web browser.

File: applications/ColossalQA/README.md Copy Chapter

ColossalQA - Langchain-based Document Retrieval Conversation System

Table of Contents Copy Chapter

• Table of Contents
• Overall Implementation
• Install
• How to Use
• Examples
  • A Simple Web UI Demo
  • Local Chinese Retrieval QA + Chat
  • Local English Retrieval QA + Chat
  • Local Bi-lingual Retrieval QA + Chat
  • Experimental AI Agent Based on Chatgpt + Chat
• Use cases
  • English customer service chatbot
  • Chinese customer service intent classification

As Colossal-AI is undergoing some major updates, this project will be actively maintained to stay in line with the Colossal-AI project.

Overall Implementation Copy Chapter

Highlevel Design Copy Section

Retrieval-based Question Answering (QA) is a crucial application of natural language processing that aims to find the most relevant answers based on the information from a corpus of text documents in response to user queries. Vector stores, which represent documents and queries as vectors in a high-dimensional space, have gained popularity for their effectiveness in retrieval QA tasks.

Step 1: Collect Data

A successful retrieval QA system starts with high-quality data. You need a collection of text documents that's related to your application. You may also need to manually design how your data will be presented to the language model.

Step 2: Split Data

Document data is usually too long to fit into the prompt due to the context length limitation of LLMs. Supporting documents need to be split into short chunks before constructing vector stores. In this demo, we use neural text splitter for better performance.

Step 3: Construct Vector Stores

Choose a embedding function and embed your text chunk into high dimensional vectors. Once you have vectors for your documents, you need to create a vector store. The vector store should efficiently index and retrieve documents based on vector similarity. In this demo, we use Chroma and incrementally update indexes of vector stores. Through incremental update, one can update and maintain a vector store without recalculating every embedding.
You are free to choose any vector store from a variety of vector stores supported by Langchain. However, the incremental update only works with LangChain vector stores that support:

• Document addition by id (add_documents method with ids argument)
• Delete by id (delete method with)

Step 4: Retrieve Relative Text

Upon querying, we will run a reference resolution on user's input, the goal of this step is to remove ambiguous reference in user's query such as "this company", "him". We then embed the query with the same embedding function and query the vector store to retrieve the top-k most similar documents.

Step 5: Format Prompt

The prompt carries essential information including task description, conversation history, retrieved documents, and user's query for the LLM to generate a response. Please refer to this README for more details.

Step 6: Inference

Pass the prompt to the LLM with additional generation arguments to get agent response. You can control the generation with additional arguments such as temperature, top_k, top_p, max_new_tokens. You can also define when to stop by passing the stop substring to the retrieval QA chain.

Step 7: Update Memory

We designed a memory module that automatically summarize overlength conversation to fit the max context length of LLM. In this step, we update the memory with the newly generated response. To fix into the context length of a given LLM, we summarize the overlength part of historical conversation and present the rest in round-based conversation format. Fig.2. shows how the memory is updated. Please refer to this README for dialogue format.

Supported Language Models (LLMs) and Embedding Models Copy Section

Our platform accommodates two kinds of LLMs: API-accessible and locally run models. For the API-style LLMs, we support ChatGPT, Pangu, and models deployed through the vLLM API Server. For locally operated LLMs, we are compatible with any language model that can be initiated using transformers.AutoModel.from_pretrained. However, due to the dependence of retrieval-based QA on the language model's abilities in zero-shot learning, instruction following, and logical reasoning, smaller models are typically not advised. In our local demo, we utilize ChatGLM2 for Chinese and LLaMa2 for English. Modifying the base LLM requires corresponding adjustments to the prompts.

Here are some sample codes to load different types of LLM.

# For locally-run LLM
from colossalqa.local.llm import ColossalAPI, ColossalLLM
api = ColossalAPI('chatglm2', 'path_to_chatglm2_checkpoint')
llm = ColossalLLM(n=1, api=api)

# For LLMs running on the vLLM API Server
from colossalqa.local.llm import VllmAPI, VllmLLM
vllm_api = VllmAPI("Your_vLLM_Host", "Your_vLLM_Port")
llm = VllmLLM(n=1, api=vllm_api)

# For ChatGPT LLM
from langchain.llms import OpenAI
llm = OpenAI(openai_api_key="YOUR_OPENAI_API_KEY")

# For Pangu LLM
# set up your authentication info
from colossalqa.local.pangu_llm import Pangu
os.environ["URL"] = ""
os.environ["URLNAME"] = ""
os.environ["PASSWORD"] = ""
os.environ["DOMAIN_NAME"] = ""

llm = Pangu(id=1)
llm.set_auth_config()

Regarding embedding models, we support all models that can be loaded via "langchain.embeddings.HuggingFaceEmbeddings". The default embedding model used in this demo is "moka-ai/m3e-base", which enables consistent text similarity computations in both Chinese and English.

In the future, supported LLM will also include models running on colossal inference and serving framework.

Install Copy Chapter

Install colossalqa

# python>=3.10,<3.13
cd ColossalAI/applications/ColossalQA
pip install -e .

To use the vLLM for providing LLM services via an API, please consult the official guide here to start the API server. It's important to set up a new virtual environment for installing vLLM, as there are currently some dependency conflicts between vLLM and ColossalQA when installed on the same machine.

How to Use Copy Chapter

Collect Your Data Copy Section

For ChatGPT based Agent we support document retrieval and simple sql search.
If you want to run the demo locally, we provided document retrieval based conversation system built upon langchain. It accept a wide range of documents. After collecting your data, put your data under a folder.

Read comments under ./colossalqa/data_loader for more detail regarding supported data formats.

Run The Script Copy Section

We provide a simple Web UI demo of ColossalQA, enabling you to upload your files as a knowledge base and interact with them through a chat interface in your browser. More details can be found here

We also provided some scripts for Chinese document retrieval based conversation system, English document retrieval based conversation system, Bi-lingual document retrieval based conversation system and an experimental AI agent with document retrieval and SQL query functionality. The Bi-lingual one is a high-level wrapper for the other two classes. We write different scripts for different languages because retrieval QA requires different embedding models, LLMs, prompts for different language setting. For now, we use LLaMa2 for English retrieval QA and ChatGLM2 for Chinese retrieval QA for better performance.

To run the bi-lingual scripts.

python retrieval_conversation_universal.py \
    --en_model_path /path/to/Llama-2-7b-hf \
    --zh_model_path /path/to/chatglm2-6b \
    --zh_model_name chatglm2 \
    --en_model_name llama \
    --sql_file_path /path/to/any/folder

To run retrieval_conversation_en.py.

python retrieval_conversation_en.py \
    --model_path /path/to/Llama-2-7b-hf \
    --model_name llama \
    --sql_file_path /path/to/any/folder

To run retrieval_conversation_zh.py.

python retrieval_conversation_zh.py \
    --model_path /path/to/chatglm2-6b \
    --model_name chatglm2 \
    --sql_file_path /path/to/any/folder

To run retrieval_conversation_chatgpt.py.

python retrieval_conversation_chatgpt.py \
    --open_ai_key_path /path/to/plain/text/openai/key/file \
    --sql_file_path /path/to/any/folder

To run conversation_agent_chatgpt.py.

python conversation_agent_chatgpt.py \
    --open_ai_key_path /path/to/plain/text/openai/key/file

After running the script, it will ask you to provide the path to your data during the execution of the script. You can also pass a glob path to load multiple files at once. Please read this guide on how to define glob path. Follow the instruction and provide all files for your retrieval conversation system then type "ESC" to finish loading documents. If csv files are provided, please use "," as delimiter and """ as quotation mark. For json and jsonl files. The default format is

{
  "data":[
    {"content":"XXX"},
    {"content":"XXX"}
    ...
  ]
}

For other formats, please refer to this document on how to define schema for data loading. There are no other formatting constraints for loading documents type files. For loading table type files, we use pandas, please refer to Pandas-Input/Output for file format details.

We also support another kay-value mode that utilizes a user-defined key to calculate the embeddings of the vector store. If a query matches a specific key, the value corresponding to that key will be used to generate the prompt. For instance, in the document below, "My coupon isn't working." will be employed during indexing, whereas "Question: My coupon isn't working.\nAnswer: We apologize for ... apply it to?" will appear in the final prompt. This format is typically useful when the task involves carrying on a conversation with readily accessible conversation data, such as customer service, question answering.

Document(page_content="My coupon isn't working.", metadata={'is_key_value_mapping': True, 'seq_num': 36, 'source': 'XXX.json', 'value': "Question: My coupon isn't working.\nAnswer:We apologize for the inconvenience. Can you please provide the coupon code and the product name or SKU you're trying to apply it to?"})

For now, we only support the key-value mode for json data files. You can run the script retrieval_conversation_en_customer_service.py by the following command.

python retrieval_conversation_en_customer_service.py \
    --model_path /path/to/Llama-2-7b-hf \
    --model_name llama \
    --sql_file_path /path/to/any/folder

The Plan Copy Chapter

• build document retrieval QA tool
• Add memory
• Add demo for AI agent with SQL query
• Add customer retriever for fast construction and retrieving (with incremental update)

Reference Copy Chapter

@software{Chase_LangChain_2022,
author = {Chase, Harrison},
month = oct,
title = {{LangChain}},
url = {https://github.com/hwchase17/langchain},
year = {2022}
}

@inproceedings{DBLP:conf/asru/ZhangCLLW21,
  author    = {Qinglin Zhang and
               Qian Chen and
               Yali Li and
               Jiaqing Liu and
               Wen Wang},
  title     = {Sequence Model with Self-Adaptive Sliding Window for Efficient Spoken
               Document Segmentation},
  booktitle = {{IEEE} Automatic Speech Recognition and Understanding Workshop, {ASRU}
               2021, Cartagena, Colombia, December 13-17, 2021},
  pages     = {411--418},
  publisher = {{IEEE}},
  year      = {2021},
  url       = {https://doi.org/10.1109/ASRU51503.2021.9688078},
  doi       = {10.1109/ASRU51503.2021.9688078},
  timestamp = {Wed, 09 Feb 2022 09:03:04 +0100},
  biburl    = {https://dblp.org/rec/conf/asru/ZhangCLLW21.bib},
  bibsource = {dblp computer science bibliography, https://dblp.org}
}

@misc{touvron2023llama,
      title={Llama 2: Open Foundation and Fine-Tuned Chat Models},
      author={Hugo Touvron and Louis Martin and Kevin Stone and Peter Albert and Amjad Almahairi and Yasmine Babaei and Nikolay Bashlykov and Soumya Batra and Prajjwal Bhargava and Shruti Bhosale and Dan Bikel and Lukas Blecher and Cristian Canton Ferrer and Moya Chen and Guillem Cucurull and David Esiobu and Jude Fernandes and Jeremy Fu and Wenyin Fu and Brian Fuller and Cynthia Gao and Vedanuj Goswami and Naman Goyal and Anthony Hartshorn and Saghar Hosseini and Rui Hou and Hakan Inan and Marcin Kardas and Viktor Kerkez and Madian Khabsa and Isabel Kloumann and Artem Korenev and Punit Singh Koura and Marie-Anne Lachaux and Thibaut Lavril and Jenya Lee and Diana Liskovich and Yinghai Lu and Yuning Mao and Xavier Martinet and Todor Mihaylov and Pushkar Mishra and Igor Molybog and Yixin Nie and Andrew Poulton and Jeremy Reizenstein and Rashi Rungta and Kalyan Saladi and Alan Schelten and Ruan Silva and Eric Michael Smith and Ranjan Subramanian and Xiaoqing Ellen Tan and Binh Tang and Ross Taylor and Adina Williams and Jian Xiang Kuan and Puxin Xu and Zheng Yan and Iliyan Zarov and Yuchen Zhang and Angela Fan and Melanie Kambadur and Sharan Narang and Aurelien Rodriguez and Robert Stojnic and Sergey Edunov and Thomas Scialom},
      year={2023},
      eprint={2307.09288},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}

@article{zeng2022glm,
  title={Glm-130b: An open bilingual pre-trained model},
  author={Zeng, Aohan and Liu, Xiao and Du, Zhengxiao and Wang, Zihan and Lai, Hanyu and Ding, Ming and Yang, Zhuoyi and Xu, Yifan and Zheng, Wendi and Xia, Xiao and others},
  journal={arXiv preprint arXiv:2210.02414},
  year={2022}
}

@inproceedings{du2022glm,
  title={GLM: General Language Model Pretraining with Autoregressive Blank Infilling},
  author={Du, Zhengxiao and Qian, Yujie and Liu, Xiao and Ding, Ming and Qiu, Jiezhong and Yang, Zhilin and Tang, Jie},
  booktitle={Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
  pages={320--335},
  year={2022}
}

File: applications/ColossalMoE/README.md Copy Chapter

Mixtral

Usage Copy Chapter

1. Installation Copy Section

Please install the latest ColossalAI from source.

CUDA_EXT=1 pip install -U git+https://github.com/hpcaitech/ColossalAI

Then install dependencies.

cd ColossalAI/applications/ColossalMoE
pip install -e .

Additionally, we recommend you to use torch 1.13.1. We've tested our code on torch 1.13.1 and found it's compatible with our code.

2. Inference Copy Section

Yon can use colossalai run to launch inference:

bash infer.sh

If you already have downloaded model weights, you can change name to your weights position in infer.sh.

3. Train Copy Section

You first need to create ./hostfile, listing the ip address of all your devices, such as:

111.111.111.110
111.111.111.111

Then yon can use colossalai run to launch train:

bash train.sh

It requires 16 H100 (80G) to run the training. The number of GPUs should be divided by 8. If you already have downloaded model weights, you can change name to your weights position in train.sh.

File: applications/ColossalEval/README.md Copy Chapter

Table of Contents Copy Chapter

• Table of Contents
• Overview
• Leaderboard
  • Model with ~13 Billion Parameters
  • Model with ~7 Billion Parameters
• Install
• Evaluation Process
  • Inference
    • Dataset Preparation
    • Configuration
    • How to Use
  • Evaluation
    • Dataset Evaluation
      • Configuration
      • How to Use
    • GPT Evaluation
      • Configuration
      • How to Use
• More Details
  • Inference
  • Evaluation
    • Metrics
  • Examples
    • Dataset Evaluation Example
    • GPT Evaluation Example
• FAQ
  • How to Add a New Metric?
  • How to Add a New Dataset?
  • How to Add a New Model?
• To do
• Citations

Overview Copy Chapter

ColossalEval is a project which provides a uniform pipeline to help evaluate language models on different public dataset or your own dataset using both classic metrics and the help from GPTs. Currently we support AGIEval, CEval, CMMLU, CValues, GAOKAO-Bench, GSM8K, LongBench, MMLU, MtBench and SafetyBench. More details can be found in the following sections.

Leaderboard Copy Chapter

Model with ~13 Billion Parameters Copy Section

We conducted comprehensive evaluation on 5 datasets and compare our Colossal-Llama-2-13b-base model with various models.

• We use 5-shot for MMLU and calculate scores based on the logits of first predicted token.
• We use 5-shot for CMMLU and calculate scores based on the logits of first predicted token.
• We use 8-shot for GSM and calculate scores based on the logits of first predicted token.
• We use 5-shot for AGIEval and only calculate scores for 4-choice questions using a combination metric of exact match and the logits of first predicted token. If any of the exact match or logits of first predicted token is correct, the model will get the score.
• We use 0-shot for GAOKAO-Bench and only calculate scores for 4-choice questions based on the logits of first predicted token.
• The generation config for all dataset is greedy search.
• We also provided CEval scores from its latest leaderboard or the official repository of the model.

The score in parentheses corresponds to the scores in the official repository of the model.

More details about metrics can be found in Metrics.

Model with ~7 Billion Parameters Copy Section

We conducted comprehensive evaluation on 4 datasets and compare our Colossal-Llama-2-7b-base model with various models.

• We use 5-shot for MMLU and calculate scores based on the logits of first predicted token.
• We use 5-shot for CMMLU and calculate scores based on the logits of first predicted token.
• We use 5-shot for AGIEval and only calculate scores for 4-choice questions using a combination metric of exact match and the logits of first predicted token. If any of the exact match or logits of first predicted token is correct, the model will get the score.
• We use 0-shot for GAOKAO-Bench and only calculate scores for 4-choice questions based on the logits of first predicted token.
• The generation config for all dataset is greedy search.
• We also provided CEval scores from its latest leaderboard or the official repository of the model.

More details about metrics can be found in Metrics.

The score in parentheses corresponds to the scores in the official repository of the model.

We use zero-shot for ChatGLM models.

To evaluate Qwen-7B on dataset MMLU, the prompt would be "xxx Answer:"(remove the space after ":") and we calculate the logits over " A", " B", " C" and " D" for Qwen-7B. Both the original and updated versions of Qwen-7B tend to be much more deterministic than other models. For example, the logits over " A" can be -inf and softmax would be exact 0.

For other models and other dataset, we calculate logits over "A", "B", "C" and "D".

Our model achieves a much better score over all other Llama-1 or Llama-2 based models and also stands out among popular open source LLMs.

Install Copy Chapter

You should install ColossalEval in order to use it and colossal_eval is the package installed.

git clone https://github.com/hpcaitech/ColossalAI.git
cd ColossalAI/applications/ColossalEval
pip install .

If you want to add customized dataset or models, use pip install -e . in stead to ensure that any changes you make to the source code will immediately affect the package you install.

Evaluation Process Copy Chapter

The evaluation process involves 2 steps which are inference and evaluation. You need to set the config for each step.

Inference Copy Section

The inference process consists of two parts. We now support tensor parallel inference for large models using ShardFormer in the example script.

1. Preprocess and convert the original dataset.
2. Config your tokenizer and model arguments to perform zero-shot or few-shot prompting.

Dataset Preparation

In this step, the original dataset(either in csv or jsonl format) will be loaded and converted into a dict. In the conversion process, we carefully parse each subcategory and assign specific inference arguments for this subcategory.

Inference arguments are stored in a dict. The following is an example.

inference_kwargs = {
    "calculate_loss": True,
    "all_classes": ["A", "B", "C", "D"],
    "language": "Chinese",
    "calculate_overall_loss": False,
    "max_new_tokens": 32
}

The inference_kwargs currently contains 5 fields:

• calculate_loss (bool, compulsory): Whether the loss on target tokens will be calculated
• all_classes (Optional[list], compulsory): Whether the subcategory is a single-choice question. Specify all available options in a list or otherwise None.
• language (str, compulsory): The language for the subcategory.
• calculate_overall_loss (bool, compulsory): Whether to calculate the overall loss of sentences or not if the dataset is a pretrain dataset. It is usually used for calculate perplexity when you want to evaluate a model with extended context length.
• max_new_tokens (int, compulsory): The number of new tokens to generate during inference.

For example, for dataset MMLU, each subcategory consists of single-choice questions with options A, B, C and D by default and we can assign value ["A", "B", "C", "D"] to keyall_classes. For dataset C-Eval, target answers aren't provided in the test split so calculate_loss should be set as False. However, other dataset such as GAOKAO-bench contains different formats of questions and lacks some keys or metadata which can reveal what type (single-choice or multi-choice) of questions it is. Before assigning inference arguments, we first parse the dataset to decide which type of questions the subcategory belongs to and set the inference arguments accordingly.

Other than inference_kwargs, data is a list containing questions of a same subcategory. The following is a converted dataset.

{
    "dev": {
        "category 1": {"data": [], "inference_kwargs": {}},
        "category 2": {"data": [], "inference_kwargs": {}}
    },
    "test": {
        "category 1": {"data": [], "inference_kwargs": {}},
        "category 2": {"data": [], "inference_kwargs": {}}
    }
}

A data sample basically follow the format of Alpaca. It should contain the following keys:

• dataset (str, compulsory): The name of the dataset.
• split (str, compulsory): The split of the instruction.
• category (str, compulsory): The category of the instruction.
• instruction (str, compulsory): The instruction for the LLM.
• input (str, optional): The additional context of the instruction.
• output (str, optional): The model output of the instruction.
• target (str, optional): The target answer for the instruction.

Example:

{
    "dev": {
        "Abstract Algebra": [
            {
                "dataset": "mmlu",
                "split": "dev",
                "category": "Abstract Algebra",
                "instruction": "The following is a single-choice question on Abstract Algebra. Answer the question by replying A, B, C or D.",
                "input": "Question: Find all c in Z_3 such that Z_3[x]/(x^2 + c) is a field.\nA. 0\nB. 1\nC. 2\nD. 3\nAnswer: ",
                "output": "",
                "target": "B"
            },
        ]
    },
    "test": {
        "Abstract Algebra": [
            {
                "dataset": "mmlu",
                "split": "test",
                "category": "Abstract Algebra",
                "instruction": "The following is a single-choice question on Abstract Algebra. Answer the question by replying A, B, C or D.",
                "input": "Question: Find the degree for the given field extension Q(sqrt(2), sqrt(3), sqrt(18)) over Q.\nA. 0\nB. 4\nC. 2\nD. 6\nAnswer: ",
                "output": "",
                "target": "B"
            },
        ]
    }
}

Configuration

In this step, you will configure your tokenizer and model arguments to infer on the given datasets.

A config file consists of two parts.

1. Model config. In model config, you need to specify model name, model path, model class, tokenizer arguments and model arguments. For model class, currently we support HuggingFaceModel, HuggingFaceCausalLM, ChatGLMModel, ChatGLMModel2 and vLLMModel. HuggingFaceModel is for models that can be loaded with AutoModel and HuggingFaceCausalLM is for models that can be loaded with AutoModelForCausalLM. ChatGLMModel and ChatGLMModel2 are for ChatGLM and ChatGLM2 models respectively. vLLMModel is for models that can be loaded with vllm offline inference LLM class. You can check all model classes in colossal_eval/models/__init__.py. If your model should set trust_remote_code as true, specify it in the tokenizer_kwargs and model_kwargs fields.
2. Dataset config. In dataset config, you need to specify dataset name, path and dataset class. Currently, we support zero-shot on dataset MMLU, CMMLU, AGIEval, GAOKAO-Bench, GSM8K and LongBench and few-shot on dataset MMLU, CMMLU AGIEval and GSM8K. If you want to enable few shot, set few_shot as true. You can check all model classes in colossal_eval/dataset/__init__.py.

Once you have all config ready, the program will run inference on all the given datasets on all the given models.

An example config using model class HuggingFaceCausalLM and dataset class CMMLUDataset can be:

{
    "model": [
        {
            "name": "model name",
            "model_class": "HuggingFaceCausalLM",
            "parameters": {
                "path": "path to model",
                "model_max_length": 2048,
                "tokenizer_path": "path to tokenizer",
                "tokenizer_kwargs": {
                    "use_fast": false,
                    "trust_remote_code": true
                },
                "peft_path": null,
                "model_kwargs": {
                    "trust_remote_code": true
                },
                "prompt_template": "plain",
                "batch_size": 4
            }
        }
    ],
    "dataset": [
        {
            "name": "dataset name",
            "dataset_class": "CMMLUDataset",
            "debug": false,
            "few_shot": true,
            "path": "path to original dataset",
            "save_path": "path to save converted dataset"
        }
    ]
}

An example config using model class vLLMModel and dataset class CMMLUDataset can be:

{
    "model": [
        {
            "name": "model name",
            "model_class": "vLLMModel",
            "parameters": {
                "path": "path to model",
                "model_max_length": 2048,
                "tokenizer_path": "",
                "tokenizer_kwargs": {
                    "trust_remote_code": true
                },
                "model_kwargs": {
                    "trust_remote_code": true
                },
                "prompt_template": "plain",
                "batch_size": 4
            }
        }
    ],
    "dataset": [
        {
            "name": "dataset name",
            "dataset_class": "CMMLUDataset",
            "debug": false,
            "few_shot": true,
            "path": "path to original dataset",
            "save_path": "path to save converted dataset"
        }
    ]
}

Currently, we support Hugging Face models as well as vLLM models. For Hugging Face models, the tokenizer_kwargs is the arguments used in AutoTokenizer.from_pretrained(). The model_kwargs is the arguments used in AutoModel.from_pretrained or AutoModelForCausalLM.from_pretrained(). For vLLM model, the tokenizer_kwargs and model_kwargs are loaded together in LLM class.few_shot will be set true if you want to enable few-shot prompting for the dataset. debug will be set true if you want to verify whether your prompt is right or wrong.

For GSM8K dataset, you can set additional flags load_train or load_reference for dataset configuration as true and during the inference process, the program will calculate loss summation over all tokens for each data sample. During the evaluation process, you can use metric loss_over_all_tokens to calculate the overall loss and use it for data leakage evaluation.

How to Use

An example script can be the following. The configs/dataset_evaluation/inference.py is the same in all examples provided.

torchrun --nproc_per_node=4 inference.py \
    --config "path to config file" \
    --load_dataset \
    --tp_size 2 \
    --inference_save_path "path to save inference results"

You should specify the path to config file in config. You can run the script without specifying load_dataset if you already save the converted dataset or otherwise set it to first load the original dataset and save the converted dataset. You should specify the path to save inference results in inference_save_path. If you want to use tensor parallel inference, specify the tensor parallel size in --tp_size and the process will automatically calculate data parallel size (currently not support for vLLMModel).

Evaluation Copy Section

In the evaluation process, you only need to configure your evaluation parameters. You can use either public dataset or help from GPTs to do evaluation. We will introduce configuration for dataset evaluation and GPT evaluation.

Dataset Evaluation

In dataset evaluation, we calculate different metrics on the given inference results and public dataset.

Configuration

A config file for dataset evaluation consists of two parts.

1. Model config. In model config, you need to specify model name. If you want to evaluate perplexity over a pretrain dataset and calculate per-byte-perplexity, you have to add your tokenizer config and model max length.
2. Dataset config. In dataset config, you need to specify the evaluation metrics for the dataset.

Once you have all config ready, the program will run evaluation on inference results for all given models and dataset.

An example config can be:

{
    "model": [
        {
            "name": "model name"
        }
    ],
    "dataset": [
        {
            "name": "dataset name",
            "metrics": ["first_token_accuracy"]
        }
    ]
}

The above config specifies that the program will evaluate the inference results using first_token_accuracy metric.

How to Use

An example script can be the following.

python eval_dataset.py \
    --config "path to config file" \
    --inference_results_path "path to inference results" \
    --evaluation_results_save_path "path to save evaluation results"

You should specify the path to config file in config, the path to inference results in inference_results_path and the path to save evaluation results in evaluation_save_path.

GPT Evaluation

In GPT evaluation, we provide a prompt template which can fit in different pre-defined metrics with Chain-of-Thoughts. In the following sections, we will only introduce how you can evaluate model answers using GPTs. More details can be found in colossal_eval/evaluate/GPT Evaluation.md.

Configuration

The following is an example of a English config file. The configuration file can control how the pipeline evaluates the model. You need to specify GPT evaluation metrics. You can find an example English config file in configs/gpt_evaluation.

{
    "language": "en",
    "category": {
        "brainstorming": {
            "GPT": [
                "language organization",
                "relevance",
                "creativity",
                "practicality",
                "reasonableness"
            ]
        },
    }
}

How to Use

After setting the config file, you can evaluate the model using examples/gpt_evaluation/eval.py. If you want to make comparisons between answers of two different models, you should specify two answer files in the argument answer_file_list and two model names in the argument model_name_list(details can be found in colossal_eval/evaluate/GPT Evaluation.md). If you want to evaluate one answer file, the length of both answer_file_list and model_name_list should be 1 and the program will perform evaluation using GPTs. The prompt files for battle and gpt evaluation can be found in configs/gpt_evaluation/prompt. target file is the path to the converted dataset you save during inference time.

An example script is provided as follows:

python eval.py \
    --config_file "path to the config file" \
    --battle_prompt_file "path to the prompt file for battle" \
    --gpt_evaluation_prompt_file "path to the prompt file for gpt evaluation" \
    --target_file "path to the target answer file" \
    --answer_file_list "path to the answer file" \
    --model_name_list "the names of the model" \
    --gpt_model "which GPT model to use for evaluation" \
    --save_path "path to save results" \
    --openai_key "your openai key" \

More Details Copy Chapter

Inference Copy Section

In the inference process, we will do generation, calculate loss over target tokens, calculate number of target tokens, softmax over given options (for example, "A", "B", "C", and "D") according to the inference arguments.

For tokenization, we adopt tokenization strategy in LongBench to preserve crucial instructions on the left and right side and keep all target tokens.

For labeling target tokens, we adopt method from FastChat, but it doesn't always hold true due to tokenizers' different behavior. We plan to insert special tokens to correctly label the target tokens.

For calculating loss, we return per-sample-loss instead of per-batch-loss if we directly use model(batch).loss provided in HuggingFace.

Evaluation Copy Section

To make it more easier to set the config, you only need to specify all metrics you want to use in key metrics. However, the program will only use a subset of metrics you give for different subcategories. Applying all metrics to all subcategories is obviously unsuitable. The suggested metrics for specific categories should be defined in colossal_eval/evaluate/dataset_evaluator/metrics.py.

Metrics

• combined_single_choice_accuracy: A combination of first_token_logit and single_choice_accuracy. If one of these is correct, the model will get the score. It can be used in all dataset that contains single-choice questions.
• first_token_logit: Calculate score based on softmax score over the given choices. If the argmax of the softmax is equal to the reference, the model will get the score. If there is NaN in softmax score, it will calculate the score using exact match. It can be used in all dataset that contains single-choice questions.
• single_choice_accuracy: Calculate score using exact match. It will only get the first uppercase letter such as A, B, C or D that is not surrounded by lowercase letters. If the uppercase letter is equal to the reference, the model will get the score. It can be used in all dataset that contains single-choice questions.
• multi_choice_accuracy: Calculate score on multi-choice questions. It will get a set of all uppercase letters such as A, B, C or D that is not surrounded by lowercase letters. If the prediction contains uppercase letters that are not in reference. The model will get 0 score. If the prediction contains a uppercase letter that is in reference, the model will get a score of 1/len(reference). It is used in AGIEval and GAOKAO-Bench.
• math_equivalence: Code from hendrycks. Compute scores over the prediction math formula and reference math formula. It is used in AGIEval and GAOKAO-Bench.
• f1_score: Calculate English f1 score between prediction and reference. It is used in Longbench.
• f1_zh_score: Calculate Chinese f1 score between prediction and reference. It is used in Longbench.
• rouge_score: Calculate English f1 score between prediction and reference. It is used in GAOKAO-Bench and LongBench.
• rouge_zh_score: Calculate Chinese rouge score between prediction and reference. It is used in GAOKAO-Bench and LongBench.
• retrieval_score: Calculate English retrieval score between prediction and reference. It determines whether the output(which paragraph) corresponds to the given abstract. It is used in Longbench.
• retrieval_zh_score: Calculate Chinese retrieval score between prediction and reference. It determines whether the output(which paragraph) corresponds to the given abstract. It is used in Longbench.
• classification_score: Calculate classification score between prediction and reference. It determines whether the output(a class) is equal to the reference. It is used in Longbench.
• code_sim_score: Calculate similarity score between prediction and reference. It is used in Longbench.
• count_score: Calculate count score between prediction and reference. It determines whether the output(number of given passages) is equal to the reference. It is used in Longbench.
• gsm_accuracy: Calculate scores between prediction and reference.. It is used in GSM8K.
• perplexity: Calculate perplexity. The formula is $ perplexity = \frac{1}{n} \sum_i e^{loss_i} $ where $n$ is the number of samples and $ loss_i $ is the average loss for sample $ i $. It can be used in all dataset.
• ppl_score: Calculate perplexity score. The formula is $ ppl_score = \frac{1}{n} \sum_i e^{-loss_i} $ where $n$ is the number of samples and $ loss_i $ is the average loss for sample $ i $. It can be used in all dataset.
• ppl_score_over_choices: Calculate perplexity score over choices. The formula is $ ppl_score_over_choices= \frac{1}{n} \sum_i e^{-loss_over_choices_i} $ where $n$ is the number of samples and $ loss_over_choices_i $ is the loss on the first predicted token for sample $ i $. It can be used in all dataset that contains single-choice questions.
• per_byte_perplexity: Calculate per byte perplexity. The formula is $ \frac{1}{n} \sum_i e^{\frac{loss_i}{byte_i}} $ where $n$ is the number of samples, $ loss_i $ is the total loss for sample $ i $ and $ byte_i $ is the number of bytes sample $ i $ occupies. It can be used in all dataset.
• per_byte_ppl_score: Calculate per byte perplexity score. The formula is $ \frac{1}{n} \sum_i e^{-\frac{loss_i}{byte_i}} $ where $n$ is the number of samples, $ loss_i $ is the total loss for sample $ i $ and $ byte_i $ is the number of bytes sample $ i $ occupies. It can be used in all dataset.
• loss_over_all_tokens: Calculate loss over all tokens. The formula is $ loss_over_all_tokens = \frac{1}{n} \sum_i loss_i $ where $n$ is the total number of tokens of the dataset and $ loss_i $ is the loss summation for sample $ i $ over all tokens and $ \sum_i loss_i $ is the loss summation for all samples. It can be used in all dataset.

We use combined_single_choice_accuracy and first_token_logit in the leaderboard.

Examples Copy Section

We provide 2 examples for you to explore our colossal_eval package.

Dataset Evaluation Example

This example is in folder examples/dataset_evaluation.

1. cd examples/dataset_evaluation
2. Fill in your inference config file in config/inference/config.json. Set the model and dataset parameters.
3. Run inference.sh to get inference results.
4. Fill in your evaluation config file in config/evaluation/config.json. Set the model and dataset parameters.
5. Run eval_dataset.sh to get evaluation results.

GPT Evaluation Example

The examples is in folder examples/gpt_evaluation.

1. cd examples/gpt_evaluation
2. Fill in your inference config file in config/inference/config.json. Set the model and dataset parameters. If you want to use the example dataset we provide, the dataset is ColossalDataset.
3. Run inference.sh to get inference results.
4. Fill in your evaluation config file in config/evaluation/config.json.
5. Run eval.sh to get evaluation results.

FAQ Copy Chapter

How to Add a New Metric? Copy Section

If you want to add a customized metric, we recommend using pip install -e . to ensure that any changes you make to the source code will immediately affect the package you install.

To add a new metric, you can follow the example of multi_choice_accuracy in line 339 in colossal_eval/evaluate/dataset_evaluator/metric.py. The method take one data sample's prediction and reference as input and return a score ranging from 0 to 1.

A skeleton of code is the following.

def CustomizedMetric(prediction: str, reference: str):
	score = xxx
	return score

Once you have successfully added your own metric, you should specify your metric both in colossal_eval/evaluate/dataset_evaluator/metric.py (suggest which subcategories should the metric be applied to) and your evaluation config.

How to Add a New Dataset? Copy Section

If you want to add customized dataset, we recommend using pip install -e . to ensure that any changes you make to the source code will immediately affect the package you install.

To add a new dataset, you can follow the example of colossal_eval/dataset/mmlu.py. You need to make sure that the format of questions in one subcategory should be the same. For example, all questions should have target answers or all questions should be single-choice questions.

A skeleton of code is the following.

class CustomizedDataset(BaseDataset):
    @staticmethod
    def load():
        # 1. Load and convert the original dataset format.
    	# 2. Assign inference arguments for each subcategory.
    	# 3. Return the converted dataset.
    	pass

Once you have successfully added your own dataset, you can specify your dataset class in your inference config.

How to Add a New Model? Copy Section

If you want to add customized models, we recommend using pip install -e . to ensure that any changes you make to the source code will immediately affect the package you install.

To add a new model, you can follow the example of colossal_eval/models/huggingface.py. You need to provide a way to load the model and tokenizer, calculate loss and generate.

A skeleton of code is the following.

class CustomizedModel(BaseModel):
    def __init__(self):
        super().__init__()
		self._load_tokenizer()
		self._load_model()

	def _load_tokenizer():
		pass

	def _load_model():
		pass

	def _calculate_loss():
		pass

	def get_loss():
		self._calculate_loss()

	def inference(samples):
		# 1. Load samples from the same subcategory.
		# 2. Infer in a batch way according to inference arguments.
		# 3. Return results.
		batch_samples = xxx
		self.get_loss(batch_samples)
		self.generate(batch_samples)

		return inference_results

	def generate():
		pass

Once you have successfully added your own model, you can specify your model class in your inference config.

Citations Copy Chapter

/* Detailed source-code truncated for AI context efficiency. */

File: applications/ColossalChat/benchmarks/README.md Copy Chapter

Benchmarks

Benchmark OPT with LoRA on dummy prompt data Copy Chapter

We provide various OPT models (string in parentheses is the corresponding model name used in this script):

• OPT-125M (125m)
• OPT-350M (350m)
• OPT-700M (700m)
• OPT-1.3B (1.3b)
• OPT-2.7B (2.7b)
• OPT-3.5B (3.5b)
• OPT-5.5B (5.5b)
• OPT-6.7B (6.7b)
• OPT-10B (10b)
• OPT-13B (13b)

We also provide various training strategies:

• gemini: ColossalAI GeminiPlugin with placement_policy="cuda", like zero3
• gemini_auto: ColossalAI GeminiPlugin with placement_policy="cpu", like zero3-offload
• zero2: ColossalAI zero2
• zero2_cpu: ColossalAI zero2-offload
• 3d: ColossalAI HybridParallelPlugin with TP, DP support

How to Run Copy Chapter

cd ../tests
# Prepare data for benchmark
SFT_DATASET=/path/to/sft/data/ \
PROMPT_DATASET=/path/to/prompt/data/ \
PRETRAIN_DATASET=/path/to/ptx/data/ \
PREFERENCE_DATASET=/path/to/preference/data \
./test_data_preparation.sh
# Start benchmark
./benchmark_ppo.sh

File: applications/ColossalChat/coati/distributed/zero_bubble/README.md Copy Chapter

Zero Bubble Distributed RL Framework for Language Model Fine-Tuning

This folder contains code for the Zero Bubble distributed RL framework. It currently supports GRPO and DAPO. See the main README for general installation instructions and usage.

Note: This project is under active development — expect changes.

🛠 Installation Copy Chapter

1. Follow the general installation guide in the main README.
2. Install pygloo. Build pygloo for Ray from source following the instructions in its repository README.

Design idea Copy Chapter

We aim to reduce the “bubble” — the idle time that occurs between rollouts and training steps (illustrated in Fig. 1).

Fig. 1 - In an all-sync online RL framework, rollout workers wait for the trainer to finish training and synchronize weights, and the trainer waits for rollouts. This causes large GPU idle time.

Fig. 2 - Our Zero Bubble pipeline follows a producer–consumer pattern:

• A global data buffer temporarily stores rollouts produced by inference workers.
• A weights distributor buffers updated model weights and distributes them to inference workers.
• When the data buffer has enough data, the trainer continuously consumes from it and pushes updated weights to the weights distributor.
• After finishing a mini-batch, each inference worker checks the weights distributor and synchronizes to a newer weight version if available.

Under ideal conditions (inference workers produce data at the same rate the trainer consumes it), the pipeline eliminates idle time. We call it zero bubble because, with an unlimited data buffer, inference and training can run indefinitely without waiting. In practice, to avoid wasted compute and stale/off-policy data, we set a bounded buffer size so inference workers will briefly wait when the buffer is full.

Usage Copy Chapter

In addition to the general parameters (see the main README), the Zero Bubble pipeline introduces one additional parameter:

• data_actor_buffer_size_limit - Maximum number of rollout batches the data buffer may hold. Defaults to twice the trainer’s mini-batch size. Avoid setting this too large — a very large buffer increases off-policy training. For DAPO, since only effective prompts count, you may need to raise data_actor_buffer_size_limit depending on sample utility.

Example: RL training on 8 GPUs with Zero Bubble (zero2)

python rl_example_zero_bubble.py \
  --dataset /path/to/your/dataset.jsonl \
  --model /path/to/your/model \
  -t 4 -i 4 -b vllm -a DAPO \
  -imbs 8 -ibs 8 -tbs 8 -e 2 -rt boxed \
  -si 25 -s "Please reason step by step, and put your final answer within \\boxed{}." \
  -tMbs 2 -tmbs 2 -p Rebase_Experiments -zero 2 -mpt 512 -mnt 3584

Performance Copy Chapter

Fig. 3 - Performance of the Zero Bubble pipeline tested with an unlimited buffer size.

File: applications/ColossalChat/coati/distributed/README.md Copy Chapter

Distributed RL Framework for Language Model Fine-Tuning

This repository implements a distributed Reinforcement Learning (RL) training framework designed to fine-tune large language models using algorithms such as GRPO and DAPO. It supports multi-node and multi-GPU setups, scalable rollout generation, and policy optimization using libraries like VLLM. Currently, we support two Reinforcement Learning with Verifiable Reward (RLVR) tasks: solving math problems and code generation.

Please note that we are still under intensive development, stay tuned.

🚀 Features Copy Chapter

• Distributed Training with Ray: Scalable to multiple machines and GPUs.
• Support for GRPO and DAPO: Choose your preferred policy optimization algorithm.
• Model Backends: Support vllm as inference backends.
• Rollout and Policy Decoupling: Efficient generation and consumption of data through parallel inferencer-trainer architecture.
• Evaluation Integration: Easily plug in task-specific eval datasets.
• Checkpoints and Logging: Configurable intervals and directories.
• [New]: Zero Bubble training framework that supports GRPO and DAPO. (read more)

🛠 Installation Copy Chapter

Prepare Develop Environment Copy Section

Install Colossalai & ColossalChat

git clone https://github.com/hpcaitech/ColossalAI.git
git checkout grpo-latest
BUILD_EXT=1 pip install -e .

cd ./applications/ColossalChat
pip install -e .

Install vllm

pip install vllm==0.7.3

Install Ray.

pip install ray

Install Other Dependencies

pip install cupy-cuda12x
python -m cupyx.tools.install_library --cuda 12.x --library nccl

To support long input/output sequence length (e.g., 32K), you may need to manually change the default setting (180 seconds) for the timeout_s variable in your ray installation to a larger value as shown in the screenshot below.

Prepare Model & dataset

huggingface-cli download --local-dir-use-symlinks False Qwen/Qwen2.5-7B --local-dir /models/Qwen/Qwen2.5-7B

Architecture Design Copy Chapter

Producer-Consumer Pattern: a classic software design pattern used for managing resources, data, or tasks between two different processes or threads.

• Producer: inference engine which rollouts out examples and saves them into a shared buffer.
• Consumer: training framework which takes training examples from the shared buffer and train the policy model.

Key features for Producer-Consumer Pattern:

• Buffer: Acts as a shared queue where the producer adds data and the consumer removes data.
• Concurrency: Rollout and training can work concurrently.

🧠 Data Format Copy Chapter

Samples in the training or evaluation .jsonl file should follow the format specific to the type of task. We currently support two RLVR tasks: solving math problems and code generation.

Math Data Format Copy Section

{
  "messages": {
    "role": "user",
    "content": "Simplify $\\sqrt[3]{1+8} \\cdot \\sqrt[3]{1+\\sqrt[3]{8}}$."
  },
  "gt_answer": "3"
}

Code Data Format Copy Section

We support Prime code dataset format. Inputs and outputs in test cases should be two lists containing only strings and matching in the number of elements. Your prompt must properly instruct the LLM to generate code to read test cases from stdin and output results to stdout.

{
    "messages": {
        "role": "user",
        "content": "Solve the following coding problem using the programming language python:\n\nMikhail walks on a Cartesian plane. He starts at the point $(0, 0)$, and in one move he can go to any of eight adjacent points. For example, ..."
    },
    "test_cases": {
        "inputs": [
            "3\n2 2 3\n4 3 7\n10 1 9\n"
        ],
        "outputs": [
            "1\n6\n-1\n"
        ]
    }
}

⚙️ Hyperparameters & Arguments Copy Chapter

Distributed Training Copy Section

Sampling Copy Section

GRPO Specific Copy Section

Logging and Checkpointing Copy Section

Miscellaneous Copy Section

⚙️ GRPO Settings Copy Chapter

In addition to the two default training settings provided—GRPO and DAPO—users can customize their training by changing the following hyperparameters in grpo_config in rl_example.py.

🔄 Constraints and Notes Copy Chapter

• num_inferencer + num_trainer == NUM_GPUs

• num_inferencer % num_trainer == 0

• (num_inferencer * inference_batch_size) % (num_trainer * train_batch_size) == 0

• train_batch_size >= train_minibatch_size >= train_microbatch_size

• inference_batch_size >= inference_microbatch_size

• Set microbatch sizes based on VRAM capacity

• To use tensor parallelism on inferencer

  • set backend to vllm
  • change tensor_parallel_size in inference_model_config in rl_example.py
  • set num_inferencer = NUM_INFERENCE_GPUs / tensor_parallel_size

• To set tensor parallelism / pipeline parallelism / zero stage

  • change corresponding settings in plugin_config in rl_example.py

• Ensure rollout generation rate matches trainer consumption:

  text Copy

  num_inferencer * inference_batch_size % (
    num_trainer * train_batch_size /
    train_pipeline_parallelism_size /
    train_tensor_parallelism_size
  ) == 0

• Model weights sync every:

  text Copy

  (num_inferencer * inference_batch_size) /
  (num_trainer * train_batch_size /
    train_pipeline_parallelism_size /
    train_tensor_parallelism_size)

• Reward Type

  We currently support three reward types--- think_answer_tags, boxed, code, each varies in details such as how answer is extracted and the reward calculation process. Please select one from think_answer_tags, boxed for math problem solving and use code for code generation. The default system prompt for each reward type is as follows. Please make sure your system prompt provides information for the answer to be correctly extracted from model responses.

  • think_answer_tags

    Answer extraction: extract the content between the last <answer>, </answer> tags.

    text Copy

    You are a helpful assistant. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within   and<answer> </answer> tags, respectively, i.e.,  reasoning process here <answer> answer here </answer>. Now the user asks you to solve a math problem that involves reasoning. After thinking, when you finally reach a conclusion, clearly output the final answer without explanation within the <answer> </answer> tags, i.e., <answer> 123 </answer>.\n\n

  • boxed

    Answer extraction: extract the last content marked by \\boxed{}

    text Copy

    Please reason step by step, and put your final answer within \\boxed{}.

  • code

    Answer extraction: extract code inside ```python\n...```

    text Copy

    You are a helpful assistant.

🧪 Example: single machine 8-GPU Zero2 Strategy Copy Chapter

python rl_example.py \
  --dataset /path/to/train_data.jsonl \
  --model /path/to/Qwen2.5-3B/ \
  -t 4 -i 4 \
  -b vllm \
  -ibs 2 -tbs 4 -tMbs 1 -tmbs 4 -imbs 1 \
  -rt boxed \
  -g 4 \
  -ibs 1 \
  -tbs 2 \
  -tMbs 1 \
  -tmbs 2 \
  -imbs 1 \
  -s "Please reason step by step, and put your final answer within \\boxed{}." \
  -tMbs 8 \
  -p GRPO-Train-Align-Debug \

🧪 Example: multi-machine TP+PP Strategy Copy Chapter

Create ray cluster on multi-machine Copy Section

For example, now we have 4 nodes and their IPs are 10.0.0.3, 10.0.0.4, 10.0.0.5, 10.0.0.6.
We use 10.0.0.3 as master node. First we start a ray cluster on 10.0.0.3:

ray start --head --node-ip-address=10.0.0.3

Then, for each slave node (10.0.0.4/10.0.0.5/10.0.0.6), we add to the ray cluster by following code:

ray start --address='10.0.0.3:6379'

Modify plugin_config in ./applications/ColossalChat/rl_example.py

plugin_config={
  "tp_size": 4,
  "pp_size": 2,
  "microbatch_size": max(
    1, args.train_microbatch_size // 2
  ),  # microbatch size should be set to train_microbatch_size // pp_size
  "zero_stage": 1,
  "max_norm": 1.0,
  },  # for pp, tp

# Hint1: replace /models/Qwen/Qwen2.5-7B to your model path
#        replace /datasets/train-alignment.jsonl to your dataset path
python rl_example.py
  -m /path/to/Qwen2.5-Math-7B/ \
  -d /path/to/train_data.jsonl \
  --master_address '10.0.0.3'
  -t 16 \
  -i 16 \
  -p GRPO-Train-Align-Debug \
  -g 2 \
  -ibs 1 \
  -tbs 2 \
  -tMbs 1 \
  -tmbs 2 \
  -imbs 1 \
  -b vllm \
  -e 2 \
  -rt boxed \
  -s "Please reason step by step, and put your final answer within \\boxed{}."

Acknowledgement Copy Chapter

Colossal-RL is a distributed version of ColossalChat and inspired by a few awesome open-source projects. We would like to express our gratitude to the following awesome open-source projects and algorithms: GRPO, DAPO, TRL, Verl, OpenRLHF, StreamRL, Qwen, Logic-RL.

File: applications/ColossalChat/coati/ray/README.md Copy Chapter

:warning: This content may be outdated since the major update of Colossal Chat. We will update this content soon.

Distributed PPO Training on Stage 3

Detach Experience Makers and Trainers Copy Chapter

We can completely separate the trainers and makers.

• The experience maker performs inference, produces experience, and remotely delivers it to the trainer (1).
• The trainer consumes experience to train models, and periodically transmits new model parameters to the maker (2.1, 2.2).
• Using an experience buffer to overlap transmission and computing.

In this manner, each node will work continuously without model idle time, and different optimization strategies can be applied for inference and training to meet the needs of speed or storage. It is also helpful for scalability.

DetachedPPOTrainer and ExperienceMakerHolder are Ray Actors (distinguished from Actor Model), representing Trainer and Experience Maker on the graph above, respectively.

More about Ray Core

Usage Copy Chapter

See examples at ColossalAI/application/Chat/examples/ray

Setup Makers Copy Section

• define makers' environment variables :

  python Copy

  env_info_makers = [{
      'local_rank': '0',
      'rank': str(rank),
      'world_size': str(num_makers),
      'master_port': maker_port,
      'master_addr': master_addr
  } for rank in range(num_makers)]

• define maker models :

  python Copy

  def model_fn():
      actor = get_actor_from_args(...)
      critic = get_critic_from_args(...)
      reward_model = get_reward_model_from_args(...)
      initial_model = get_actor_from_args(...)
      return actor, critic, reward_model, initial_model

• set experience_holder_refs :

  python Copy

  experience_holder_refs = [
      ExperienceMakerHolder.options(
          name=f"maker_{i}",
          num_gpus=1,
          max_concurrency=2
      ).remote(
          detached_trainer_name_list=[f"trainer_{x}" for x in target_trainers(...)],
          model_fn=model_fn,
          ...)
      for i, env_info_maker in enumerate(env_info_makers)
  ]

  The names in the detached_trainer_name_list refer to the target trainers that the maker should send experience to.
  We set a trainer's name the same as a maker, by .options(name="str"). See below.

Setup Trainers Copy Section

• define trainers' environment variables :

  python Copy

  env_info_trainers = [{
      'local_rank': '0',
      'rank': str(rank),
      'world_size': str(num_trainers),
      'master_port': trainer_port,
      'master_addr': master_addr
  } for rank in range(num_trainers)]

• define trainer models :

  python Copy

  def trainer_model_fn():
      actor = get_actor_from_args(...)
      critic = get_critic_from_args(...)
      return actor, critic

• set trainer_refs :

  python Copy

  trainer_refs = [
      DetachedPPOTrainer.options(
          name=f"trainer{i}",
          num_gpus=1,
          max_concurrency=2
      ).remote(
          experience_maker_holder_name_list=[f"maker{x}" for x in target_makers(...)],
          model_fn = trainer_model_fn(),
          ...)
      for i, env_info_trainer in enumerate(env_info_trainers)
  ]

  The names in experience_maker_holder_name_list refer to the target makers that the trainer should send updated models to.
  By setting detached_trainer_name_list and experience_maker_holder_name_list, we can customize the transmission graph.

Launch Jobs Copy Section

• define data_loader :

  python Copy

  def data_loader_fn():
      return = torch.utils.data.DataLoader(dataset=dataset)

• launch makers :

  python Copy

  wait_tasks = []
  for experience_holder_ref in experience_holder_refs:
      wait_tasks.append(
          experience_holder_ref.workingloop.remote(data_loader_fn(),
                                                   num_steps=experience_steps))

• launch trainers :

  python Copy

  for trainer_ref in trainer_refs:
      wait_tasks.append(trainer_ref.fit.remote(total_steps, update_steps, train_epochs))

• wait for done :

  python Copy

  ray.get(wait_tasks)

Flexible Structure Copy Chapter

We can deploy different strategies to makers and trainers. Here are some notions.

2 Makers 1 Trainer Copy Section

2 Makers 2 Trainer Copy Section

Maker Inference Quantization Copy Section

Tensor Parallel Copy Section

TODO Copy Chapter

• Support LoRA
• Support TP & PP