Mixture of Agents（MoA） 実践！

何かと話題の「LLMをたくさん用意して組み合わせる系」の研究。
Mixture of Experts（MoE）に続き。
AlpacaEval2.0ベンチマークでGPT-4にオープンソースの組み合わせで勝利したとか何とか。https://github.com/togethercomputer/MoA

論文：https://arxiv.org/abs/2406.04692, Together AI

概要

論文中の図です。

 実装

異なるLLM（共通していても可）の出力をつなげながら次に渡していけば良さそう

問題点： 複数のLLMをloadすると, GPUがOutOfMemoryになる…

ちゃんとした原因は特定できていません。なぜか上手くいかなかった。

対処： 逐一モデルをload・解放することを繰り返す（当然時間は無駄にかかる）

推論にはvllmを利用しています。vllmのバージョンによってメモリ解放のコードが微妙に異なるので注意！

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoConfig
from peft import PeftModel
import glob
import json
import pandas as pd
import os
import gc
import torch
from vllm import LLM, SamplingParams
from vllm.distributed.parallel_state import destroy_model_parallel


def generate(model,text):
    outputs = model.generate([text], sampling_params)
    output = outputs[0]
    response = output.outputs[0].text
    return response

    
test_dataset = prepare_data(...) #データを用意
test_df = pd.DataFrame(test_dataset)
n = 3 #Layer内のLLMの数
L = 3 #Layerの数
responses = []
full_responses = []
model_list = ["Qwen/Qwen2-7B-Instruct",...] #model_idのリストを作成

for k in range(len(test_dataset)):
    test_data = test_dataset[k]
    text = "### 質問：\n" + test_data["problem_text"] + "\n" + ",".join(test_data["choices"]) + "\n### 回答：\n" # input prompt, 論文中のxi
    for layer in range(L):       
        print(f"Layer {layer}:")
        #論文中のプロンプトをGPTで翻訳したテンプレート↓
        prompt = "以下は、複数のオープンソースモデルから得られた応答を、単一の質の高い回答にまとめるための和訳です。これにより、提供された情報を批判的に評価し、バイアスや不正確な情報が含まれている可能性を認識することが重要です。回答は単に与えられた回答を複製するのではなく、洗練された、正確で包括的な返答を提供することを目指し、高い構造化、論理性、正確性、信頼性の基準を守ってください。\n\nモデルからの応答：\n"
        for i in range(n):
            sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=48)
            llm = LLM(model = model_list[i]) #異なるLLMを利用
            output = generate(llm,text)
            prompt += f"{i+1}. "
            prompt += output
            prompt += "\n"

            # メモリ解放
            destroy_model_parallel()
            del llm.llm_engine.model_executor.driver_worker
            del llm 
            gc.collect()
            torch.cuda.empty_cache()
           
        prompt += text
        text = prompt
        print(text)

    llm = LLM(model = model_list[0]) #最後の１匹
    final_output = generate(llm,text)
    # メモリ解放
    destroy_model_parallel()
    del llm.llm_engine.model_executor.driver_worker #これが必要そう
    del llm 
    gc.collect()
    torch.cuda.empty_cache()

    full_responses.append(final_output)
    responses.append(final_output.replace(prompt,""))
test_df["full_response"] = full_responses
test_df["response"] = responses

test_df.to_json('result.jsonl', orient='records', force_ascii=False, lines=True)
test_df.to_csv('result.csv',index=False)