Mergoo: 다양한 LLM 전문가를 신뢰성 있게 통합하는 방법

최근 인공지능 분야에서 다양한 LLM(대형 언어 모델) 전문가들의 지식을 통합하는 작업이 중요해지고 있다. 이를 실현하기 위해 설계된 것이 바로 Mergoo다. Mergoo는 다양한 일반 및 도메인 특화 LLM 전문가들의 지식을 신뢰성 있고 투명하게 결합할 수 있는 파이프라인을 제공하며, 이를 통해 LLM 통합 작업을 보다 유연하고 효과적으로 수행할 수 있다.

 

Mergoo는 Mixture-of-Experts(MoE), Mixture-of-Adapters, 그리고 Layer-wise Merging과 같은 다양한 통합 기법을 지원하며, 이를 통해 사용자는 자신이 원하는 방식으로 LLM을 구축할 수 있다. 통합된 모델은 Hugging Face에서 제공하는 SFTrainer, PEFT, Trainer와 같은 도구를 사용하여 특정 용도에 맞게 추가로 미세 조정할 수 있다.

 

지원 모델 

• llama
• Mistral
• Bert
• Phi3

1. 완전 미세 조정된 LLM의 혼합

 

Branch-Train-Mix 방식을 사용하면, 도메인 특화된 LLM 전문가들을 Mixture-of-Experts(MoE) 레이어로 통합할 수 있다. 이때, 각 전문가의 feedforward 파라미터를 MoE 레이어의 전문가로 통합하고, 나머지 파라미터는 평균을 내서 통합하게 된다. MoE 레이어는 이후 토큰 수준의 라우팅을 학습하기 위해 추가로 미세 조정할 수 있다.

 

예를 들어, 다음과 같은 도메인 특화 LLM 전문가들을 하나로 통합할 수 있다:

• Base Model: meta-llama/Llama-2-7b-hf
• Code Expert: codellama/CodeLlama-7b-Python-hf
• WikiChat Expert: stanford-oval/Llama-2-7b-WikiChat-fused

이 전문가들을 통합하기 위한 구성(config)은 다음과 같다:

config = \
{
    "model_type": "llama",
    "num_experts_per_tok": 2,
    "experts":[
        {
            "expert_name" : "base_expert",
            "model_id" : "meta-llama/Llama-2-7b-hf"
        },
        {
            "expert_name" : "expert_1",
            "model_id" : "codellama/CodeLlama-7b-Python-hf"
        },
        {
            "expert_name" : "expert_2",
            "model_id" : "stanford-oval/Llama-2-7b-WikiChat-fused"
        }
    ],
    "router_layers":[
        "gate_proj",
        "up_proj",
        "down_proj"
    ],
}

 

이 구성에 따라 통합된 전문가의 체크포인트를 생성하고 저장하는 코드는 다음과 같다:

import torch
from mergoo.compose_experts import ComposeExperts
model_id = "mergoo_llama_code_wikichat"
expertmerger = ComposeExperts(config, torch_dtype=torch.float16)
expertmerger.compose()
expertmerger.save_checkpoint(model_id)

 

생성된 체크포인트는 인스트럭션 데이터셋을 사용하여 추가로 학습할 수 있다:

from mergoo.models.modeling_llama import LlamaForCausalLM
import torch
import datasets
import random
from trl import SFTTrainer
from transformers import TrainingArguments

# 통합된 체크포인트 로드
model = LlamaForCausalLM.from_pretrained(
    "mergoo_llama_code_wikichat", 
    device_map="auto", 
    torch_dtype=torch.bfloat16,
)# 'gate' 및 라우터 레이어는 학습되지 않았으므로 로드 시 경고가 발생할 수 있음

# 학습 데이터셋 로드
dataset = datasets.load_dataset("iamtarun/python_code_instructions_18k_alpaca")['train']
dataset = dataset['prompt']
random.shuffle(dataset)
train_dataset =  datasets.Dataset.from_dict(dict(prompt=dataset[:-1000]))
eval_dataset = datasets.Dataset.from_dict(dict(prompt=dataset[-1000:]))

# 학습 인자 설정
trainer_args = TrainingArguments(
    output_dir= "checkpoints/llama_moe",
    per_device_train_batch_size = 1,
    per_device_eval_batch_size = 1, 
    learning_rate= 1e-5,
    save_total_limit=1,
    num_train_epochs=1,
    eval_steps= 5000,
    logging_strategy="steps",
    logging_steps= 25,
    gradient_accumulation_steps=4,
    bf16=True
)

trainer = SFTTrainer(
    model,
    args= trainer_args,
    train_dataset= train_dataset,
    eval_dataset= eval_dataset,
    dataset_text_field="prompt",
)

# 학습 시작
trainer.train()

 

2. Mixture of Adapters

Mergoo는 여러 LoRA adapters를 MoE 스타일의 아키텍처로 병합하는 기능을 제공한다. 이 과정은 미세 조정된 LoRA 모델에 게이팅(gating)과 라우팅(routing) 레이어를 적용하여 이루어진다.

예를 들어, Customer Support 도메인의 다음과 같은 전문가들을 통합할 수 있다:

• predibase/customer_support
• predibase/customer_support_accounts
• predibase/customer_support_orders

이 LoRA 전문가들을 통합하기 위한 구성(config)은 다음과 같다:

import torch
from mergoo.compose_experts import ComposeExperts

model_id = "mergoo_customer_support_moe_lora"
config = {
    "model_type": "mistral",
    "num_experts_per_tok": 2,
    "base_model": "mistralai/Mistral-7B-v0.1",
    "experts": [
        {
            "expert_name": "adapter_1", # expert_name이 반드시 adapter_로 시작
            "model_id": "predibase/customer_support"
        },
        {
            "expert_name": "adapter_2", # expert_name이 반드시 adapter_로 시작
            "model_id": "predibase/customer_support_accounts"
        },
        {
            "expert_name": "adapter_3", # expert_name이 반드시 adapter_로 시작
            "model_id": "predibase/customer_support_orders"
        }
    ],
}

expertmerger = ComposeExperts(config, torch_dtype=torch.bfloat16)
expertmerger.compose()
expertmerger.save_checkpoint(model_id)

LoRA로 미세 조정된 전문가를 병합할 때는, expert_name이 반드시 adapter_로 시작해야 한다. 이후, 병합된 전문가 모델도 Mixture of Fully Fine-tuned LLMs에서 설명한 방식으로 추가 학습할 수 있다.

 

3) 결론

 

Mergoo는 다양한 LLM 전문가들의 지식을 통합하기 위한 신뢰성 있고 투명한 도구로, Mixture-of-Experts, Mixture-of-Adapters(MoE-LoRA), 그리고 Layer-wise Merging 같은 다양한 기법을 지원한다. 이를 통해 통합된 LLM은 downstream tasks에 맞춰 추가 학습할 수 있으며, 이를 통해 보다 신뢰성 있고 강력한 전문가 시스템을 구축할 수 있다.