1_How does ChatML affect the performance of LLMs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How does ChatML affect the quality of LLMs?\n",
    "\n",
    "ChatML is short for Chat Makeup Language. More details can be [find here](https://github.com/MicrosoftDocs/azure-docs/blob/main/articles/ai-services/openai/includes/chat-markup-language.md).\n",
    "\n",
    "ChatML make explicit to the model the source of each piece of text, and particularly show the boundary between human and AI. For ChatGPT, manage conversations in a format like following:\n",
    "\n",
    "```python\n",
    "[\n",
    "    {\"role\": \"system\", \"content\": \"Provide some context and/or instructions to the model.\"},\n",
    "    [\"role\": \"user\", \"content\": \"The user’s message goes here\"],\n",
    "]\n",
    "```\n",
    "\n",
    "Then convert the conversation messages into following form as input to the model, To adhere the instruction format, append the assistant token at the end:\n",
    "\n",
    "```text\n",
    "\n",
    "<|im_start|>system \n",
    "Provide some context and/or instructions to the model.\n",
    "<|im_end|> \n",
    "<|im_start|>user \n",
    "The user’s message goes here\n",
    "<|im_end|> \n",
    "<|im_start|>assistant \n",
    "\n",
    "```\n",
    "\n",
    "Different models follow different ChatML formats, Make sure you generate the right prompt to achieve the expected model performance.\n",
    "\n",
    "\n",
    "\n",
    "## Compare Huggingface(w/ ChatML), vLLM (w/o ChatML)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/kky/miniconda3/envs/gomars/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Beijing, as the capital of China, is a modern city with a long history. It has many advantages and disadvantages compared to New York.\n",
      "\n",
      "Advantages:\n",
      "1. Political importance: As the capital of China, Beijing is an important political, economic, and cultural center in the country.\n",
      "2. History and culture: With a long history, Beijing is home to many historical sites and cultural relics, such as the Forbidden City, the Great Wall, and the Temple of Heaven.\n",
      "3. Transportation: Beijing has a well-developed transportation system, including high-speed trains, subways, and buses.\n",
      "4. Education: Beijing is a major educational center in China, with many top universities and research institutions.\n",
      "5. International influence: As the political and cultural center of China, Beijing has a growing international influence.\n",
      "\n",
      "Disadvantages:\n",
      "1. Air pollution: Due to the large number of vehicles and industrial emissions, Beijing's air quality can be poor at times.\n",
      "2. Traffic congestion: The dense traffic in Beijing can cause long delays and congestion during peak hours.\n",
      "3. High cost of living: The cost of living in Beijing is generally higher than in other cities in China, due to the high price of real estate and other living expenses.\n",
      "4. Language barrier: For non-Chinese speakers, learning Chinese and adapting to the local culture can be challenging.\n",
      "5. Food and drink: While there are numerous Chinese cuisine options in Beijing, Western food may not be as readily available or affordable as in New York.\n",
      "\n",
      "New York, located in the United States, is a global city known for its diversity, economy, and culture. It also has its own set of advantages and disadvantages compared to Beijing.\n",
      "\n",
      "Advantages:\n",
      "1. Diversity: New York is famous for its diverse population and multicultural environment, which offers a wide range of food, entertainment, and social opportunities.\n",
      "2. Economy: New York is one of the world's leading financial centers, with a strong stock market, banking, and investment sector.\n",
      "3. Culture: New York is a major center of art, fashion, and entertainment, with numerous museums, theaters, and music venues.\n",
      "4. Transportation: New York has an efficient and extensive public transportation system, including subways, buses, and taxis.\n",
      "5. International influence: As the largest city in the United States, New York has a significant global influence in various fields, such as politics, business, and media.\n",
      "\n",
      "Disadvantages:\n",
      "1. Crime rate: New York has a higher crime rate than many other cities, particularly in lower-income areas.\n",
      "2. Cost of living: The cost of living in New York is generally higher than in other cities, due to the high price of real estate, rent, and other living expenses.\n",
      "3. Weather: New York's weather can be variable, with cold winters and humid summers.\n",
      "4. Language barrier: While English is the official language of the city, learning American English and adapting to the local culture can still be challenging for non-native speakers.\n",
      "5. Urban sprawl: The vast size of New York can make it difficult to navigate for those not familiar with the city, especially outside of Manhattan.\n"
     ]
    }
   ],
   "source": [
    "# huggingface version with right ChatML.\n",
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "from transformers.generation.utils import GenerationConfig\n",
    "model_id = \"/data/hf/Baichuan2-7B-Chat/\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id,\n",
    "    revision=\"v2.0\",\n",
    "    use_fast=False,\n",
    "    trust_remote_code=True)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_id,\n",
    "    revision=\"v2.0\",\n",
    "    device_map=\"auto\",\n",
    "    torch_dtype=torch.bfloat16,\n",
    "    trust_remote_code=True)\n",
    "model.generation_config = GenerationConfig.from_pretrained(model_id, revision=\"v2.0\")\n",
    "messages = []\n",
    "messages.append({\"role\": \"user\", \"content\": \"Compare the advantages and disadvantages of Beijing and New York\"})\n",
    "response = model.chat(tokenizer, messages)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/kky/miniconda3/envs/gomars/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "2024-01-04 10:31:17,015\tINFO util.py:159 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.\n",
      "2024-01-04 10:31:17,362\tINFO util.py:159 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.\n",
      "2024-01-04 10:31:19,175\tINFO worker.py:1673 -- Started a local Ray instance.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 01-04 10:31:20 llm_engine.py:73] Initializing an LLM engine with config: model='/data/hf/Baichuan2-7B-Chat/', tokenizer='/data/hf/Baichuan2-7B-Chat/', tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=True, dtype=torch.bfloat16, max_seq_len=4096, download_dir=None, load_format=auto, tensor_parallel_size=2, quantization=None, enforce_eager=False, seed=0)\n",
      "WARNING 01-04 10:31:20 tokenizer.py:62] Using a slow tokenizer. This might cause a significant slowdown. Consider using a fast tokenizer instead.\n",
      "INFO 01-04 10:31:35 llm_engine.py:223] # GPU blocks: 3047, # CPU blocks: 1024\n",
      "\u001b[36m(RayWorkerVllm pid=933127)\u001b[0m INFO 01-04 10:31:37 model_runner.py:394] Capturing the model for CUDA graphs. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m(RayWorkerVllm pid=933126)\u001b[0m [W CUDAGraph.cpp:145] Warning: Waiting for pending NCCL work to finish before starting graph capture. (function operator())\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(RayWorkerVllm pid=933127)\u001b[0m INFO 01-04 10:32:11 model_runner.py:437] Graph capturing finished in 35 secs.\n",
      "\u001b[36m(RayWorkerVllm pid=933126)\u001b[0m INFO 01-04 10:31:37 model_runner.py:394] Capturing the model for CUDA graphs. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processed prompts: 100%|██████████| 1/1 [00:00<00:00,  3.81it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " in terms of culture and lifestyle\n",
      "Beijing and New York are two of the world\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# vLLM version without ChatML\n",
    "from vllm import LLM, SamplingParams\n",
    "model_id = \"/data/hf/Baichuan2-7B-Chat/\"\n",
    "\n",
    "sampling_params = SamplingParams(repetition_penalty=1.05, top_k=5, top_p=0.85, max_tokens=2048, temperature=0.3)\n",
    "llm = LLM(model=model_id, trust_remote_code=True, tensor_parallel_size=2)\n",
    "\n",
    "response = llm.generate(\"Compare the advantages and disadvantages of Beijing and New York\")\n",
    "print(response[0].outputs[0].text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The HF version of the model provides high-quality responses, while the vLLM version give very poor response quality. \n",
    "\n",
    "Model response of HF version:\n",
    "```text\n",
    "\n",
    "Beijing, as the capital of China, is a modern city with a long history. It has many advantages and disadvantages compared to New York.\n",
    "\n",
    "Advantages:\n",
    "1. Political importance: As the capital of China, Beijing is an important political, economic, and cultural center in the country.\n",
    "2. History and culture: With a long history, Beijing is home to many historical sites and cultural relics, such as the Forbidden City, the Great Wall, and the Temple of Heaven.\n",
    "3. Transportation: Beijing has a well-developed transportation system, including high-speed trains, subways, and buses.\n",
    "4. Education: Beijing is a major educational center in China, with many top universities and research institutions.\n",
    "5. International influence: As the political and cultural center of China, Beijing has a growing international influence.\n",
    "\n",
    "Disadvantages:\n",
    "1. Air pollution: Due to the large number of vehicles and industrial emissions, Beijing's air quality can be poor at times.\n",
    "2. Traffic congestion: The dense traffic in Beijing can cause long delays and congestion during peak hours.\n",
    "3. High cost of living: The cost of living in Beijing is generally higher than in other cities in China, due to the high price of real estate and other living expenses.\n",
    "4. Language barrier: For non-Chinese speakers, learning Chinese and adapting to the local culture can be challenging.\n",
    "5. Food and drink: While there are numerous Chinese cuisine options in Beijing, Western food may not be as readily available or affordable as in New York.\n",
    "\n",
    "New York, located in the United States, is a global city known for its diversity, economy, and culture. It also has its own set of advantages and disadvantages compared to Beijing.\n",
    "\n",
    "Advantages:\n",
    "1. Diversity: New York is famous for its diverse population and multicultural environment, which offers a wide range of food, entertainment, and social opportunities.\n",
    "2. Economy: New York is one of the world's leading financial centers, with a strong stock market, banking, and investment sector.\n",
    "3. Culture: New York is a major center of art, fashion, and entertainment, with numerous museums, theaters, and music venues.\n",
    "4. Transportation: New York has an efficient and extensive public transportation system, including subways, buses, and taxis.\n",
    "5. International influence: As the largest city in the United States, New York has a significant global influence in various fields, such as politics, business, and media.\n",
    "\n",
    "Disadvantages:\n",
    "1. Crime rate: New York has a higher crime rate than many other cities, particularly in lower-income areas.\n",
    "2. Cost of living: The cost of living in New York is generally higher than in other cities, due to the high price of real estate, rent, and other living expenses.\n",
    "3. Weather: New York's weather can be variable, with cold winters and humid summers.\n",
    "4. Language barrier: While English is the official language of the city, learning American English and adapting to the local culture can still be challenging for non-native speakers.\n",
    "5. Urban sprawl: The vast size of New York can make it difficult to navigate for those not familiar with the city, especially outside of Manhattan.\n",
    "\n",
    "\n",
    "```\n",
    "\n",
    "Model response of vLLM version:\n",
    "\n",
    "```text\n",
    " in terms of culture and lifestyle\n",
    "Beijing and New York are two of the world\n",
    "```\n",
    "\n",
    "This is due to the fact that Baichuan chat model used a unique ChatML format during supervised fine-tuning, Following the fine-tuning ChatML format is essential to achive the expected model performance. Always we can found the ChatML format definition on huggingface model page or code.\n",
    "\n",
    "Let's try to incorporate [Baichuan ChatML](https://huggingface.co/baichuan-inc/Baichuan2-7B-Chat/blob/main/generation_utils.py) into vLLM and observe the outcome."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Incorporate ChatML into vLLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from queue import Queue\n",
    "\n",
    "\n",
    "chatml_cfg = {\"user_token_id\": 195, \"assistant_token_id\": 196, \"max_new_tokens\": 2048, \"model_max_length\": 4096}\n",
    "\n",
    "def build_chat_input(config, tokenizer, messages: List[dict]):\n",
    "    def _parse_messages(messages, split_role=\"user\"):\n",
    "        system, rounds = \"\", []\n",
    "        round = []\n",
    "        for i, message in enumerate(messages):\n",
    "            if message[\"role\"] == \"system\":\n",
    "                assert i == 0\n",
    "                system = message[\"content\"]\n",
    "                continue\n",
    "            if message[\"role\"] == split_role and round:\n",
    "                rounds.append(round)\n",
    "                round = []\n",
    "            round.append(message)\n",
    "        if round:\n",
    "            rounds.append(round)\n",
    "        return system, rounds\n",
    "\n",
    "    max_new_tokens = config.get(\"max_new_tokens\")\n",
    "    max_input_tokens = config.get(\"model_max_length\") - max_new_tokens\n",
    "    system, rounds = _parse_messages(messages, split_role=\"user\")\n",
    "    system_tokens = tokenizer.encode(system)\n",
    "    max_history_tokens = max_input_tokens - len(system_tokens)\n",
    "\n",
    "    history_tokens = []\n",
    "    for round in rounds[::-1]:\n",
    "        round_tokens = []\n",
    "        for message in round:\n",
    "            if message[\"role\"] == \"user\":\n",
    "                round_tokens.append(config.get(\"user_token_id\"))\n",
    "            else:\n",
    "                round_tokens.append(config.get(\"assistant_token_id\"))\n",
    "            round_tokens.extend(tokenizer.encode(message[\"content\"]))\n",
    "        if len(history_tokens) == 0 or len(history_tokens) + len(round_tokens) <= max_history_tokens:\n",
    "            history_tokens = round_tokens + history_tokens  # concat left\n",
    "            if len(history_tokens) < max_history_tokens:\n",
    "                continue\n",
    "        break\n",
    "\n",
    "    input_tokens = system_tokens + history_tokens\n",
    "    if messages[-1][\"role\"] != \"assistant\":\n",
    "        input_tokens.append(chatml_cfg[\"assistant_token_id\"])\n",
    "    input_tokens = input_tokens[-max_input_tokens:]  # truncate left\n",
    "    return input_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processed prompts: 100%|██████████| 1/1 [00:07<00:00,  7.51s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "As two famous cities in the world, Beijing and New York have their own advantages and disadvantages. The following is a comparison between them:\n",
      "\n",
      "Advantages of Beijing:\n",
      "1. History and culture: Beijing is the capital of China and has a long history. It is full of historical sites and cultural relics, such as the Forbidden City, the Great Wall, and the Temple of Heaven.\n",
      "2. Politics and economy: Beijing is an important political and economic center in China. Many government departments and multinational companies have headquarters here.\n",
      "3. Education: Beijing has many well-known universities and research institutions, providing high-quality education for students.\n",
      "4. Transportation: Beijing has a complete transportation system, including high-speed railways, subways, buses, and taxis.\n",
      "5. Natural environment: Although Beijing is a large city, it still has some green spaces, such as the Summer Palace and the Forest Park.\n",
      "\n",
      "Disadvantages of Beijing:\n",
      "1. Air pollution: Due to the large number of vehicles and industrial emissions, the air quality in Beijing is not good.\n",
      "2. Traffic congestion: With the rapid development of Beijing's economy, traffic congestion has become more serious.\n",
      "3. High living costs: The cost of living in Beijing is relatively high, especially housing prices.\n",
      "\n",
      "Advantages of New York:\n",
      "1. Culture and entertainment: New York is a world-renowned cultural and entertainment center. It has many famous museums, theaters, and entertainment venues, such as the Metropolitan Museum of Art, Broadway, and Times Square.\n",
      "2. Finance: New York is one of the world's most important financial centers. Many stock exchanges and banks are located here.\n",
      "3. Fashion: New York is known as the fashion capital of the world. Many famous designers and brands come from here.\n",
      "4. Diversity: New York is a city with a diverse population. People come from all over the world here, which brings rich cultural experiences.\n",
      "\n",
      "Disadvantages of New York:\n",
      "1. Crime rate: New York has a relatively high crime rate, especially in terms of violence and theft.\n",
      "2. High living costs: Like Beijing, the cost of living in New York is also relatively high, especially housing prices.\n",
      "3. Bad weather: New York has bad weather in winter, especially snow and cold weather.\n",
      "\n",
      "In general, Beijing and New York each have their own advantages and disadvantages. Choosing a suitable city depends on personal preferences and needs.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# vLLM chatml version\n",
    "from vllm import LLM, SamplingParams\n",
    "\n",
    "\n",
    "model_id = \"/data/hf/Baichuan2-7B-Chat/\"\n",
    "\n",
    "sampling_params = SamplingParams(repetition_penalty=1.05, top_k=5, top_p=0.85, max_tokens=2048, temperature=0.3)\n",
    "llm = LLM(model=model_id, trust_remote_code=True, tensor_parallel_size=2)\n",
    "\n",
    "prompt = \"Compare the advantages and disadvantages of Beijing and New York\"\n",
    "prompt_token_ids = build_chat_input(chatml_cfg, llm.get_tokenizer(), [{\"role\": \"user\", \"content\": prompt}])\n",
    "\n",
    "response = llm.generate(prompt_token_ids=[prompt_token_ids], sampling_params=sampling_params)\n",
    "print(response[0].outputs[0].text)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model's response quality has significantly improved!!\n",
    "\n",
    "Be carefull when you use any LLM inference libraries like vLLM, fastGPT, etc."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gomars",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}