튜토리얼: 다중 모듈 DSPy 프로그램에 대한 온라인 RL

튜토리얼: 다중 모듈 DSPy 프로그램에 대한 온라인 RL (Online RL over a Multi-Module DSPy Program)

경고: 이 기능은 새롭고 매우 실험적입니다. DSPy의 거의 모든 다른 것과 달리, 현재는 순수한 개념 증명·개발 모드에 있으며, 커뮤니티 참여를 장려하기 위해 배포합니다.

이 튜토리얼에서는 PAPILLON의 LM 가중치를 ArborGRPO로 최적화할 거예요. ArborGRPO는 LLM의 인기 있는 온라인 RL 알고리즘인 GRPO를 정교한 다중 모듈 LM 프로그램으로 일반화한 것입니다.

PAPILLON은 프라이버시 보호 위임(privacy-preserving delegation)을 위한 시스템이에요. 여기서 우리는 작은 모델(1.5B 파라미터)이 더 강력하지만 프라이빗 데이터를 저장할 수 있는 "신뢰할 수 없는(untrusted)" 외부 LLM을 사용해, 높은 품질과 프라이빗 채팅 사이의 균형을 맞추도록 가르칠 거예요.

이 튜토리얼에는 DSPy의 Arbor RL 프레임워크도 필요하며, 아래와 같이 설치할 수 있어요:

> pip install -U arbor-ai

또한 DSPy를 main 브랜치에서 설치해야 할 수도 있어요:

> pip install -U git+https://github.com/stanfordnlp/dspy.git@main

출처: 문서

본문

import dspy
import arbor
from arbor import ArborGRPO, ArborProvider
arbor_server_info = arbor.init() # Initialize the Arbor server in the background

port = 7453
local_lm_name = "Qwen/Qwen2.5-1.5B-Instruct"
local_lm = dspy.LM(
    model=f"openai/arbor:{local_lm_name}",
    provider=ArborProvider(),
    api_base=arbor_server_info["base_url"],
    # Arbor checks to make sure these match the training config
    temperature=1.0,
    top_p=1.0,
    top_k=-1,
    repetition_penalty=1.0,
    max_tokens=2048,
)

dspy.configure(lm=local_lm)

openai_lm = dspy.LM(model="openai/gpt-4.1-mini")

class CraftRedactedRequest(dspy.Signature):
    """
    Given a private user query, create a privacy-preserving request for a powerful external LLM.
    The LLM may assist without learning private information about the user.
    """

    user_query = dspy.InputField()
    llm_request = dspy.OutputField()


class RespondToQuery(dspy.Signature):
    """
    Respond to a user query.
    For inspiration, we found a potentially related request to a powerful external LLM and its response.
    """

    related_llm_request = dspy.InputField()
    related_llm_response = dspy.InputField(desc="information from a powerful LLM responding to a related request")
    user_query = dspy.InputField(desc="the user's request you need to fulfill")
    response = dspy.OutputField(desc="your final response to the user's request")


class PAPILLON(dspy.Module):
    def __init__(self, untrusted_model):
        self.craft_redacted_request = dspy.ChainOfThought(CraftRedactedRequest)
        self.respond_to_query = dspy.Predict(RespondToQuery)
        self.untrusted_model = untrusted_model

    def forward(self, user_query):
        try:
            llm_request = self.craft_redacted_request(user_query=user_query).llm_request
            llm_response = self.untrusted_model(llm_request)[0]
            response = self.respond_to_query(
                related_llm_request=llm_request, related_llm_response=llm_response, user_query=user_query
            ).response
        except Exception:
            return dspy.Prediction(llm_request="", llm_response="", response="")

        return dspy.Prediction(llm_request=llm_request, llm_response=llm_response, response=response)

from datasets import load_dataset

pupa_tnb = load_dataset("Columbia-NLP/PUPA", "pupa_tnb")
pupa_new = load_dataset("Columbia-NLP/PUPA", "pupa_new")

examples = [
    dspy.Example(
        {"target_response": x["target_response"], "user_query": x["user_query"], "pii_str": x["pii_units"]}
    ).with_inputs("user_query")
    for x in pupa_new["train"]
]

trainset, devset, testset = examples[:225], examples[225:450], examples[450:]
print(f"Loaded {len(trainset)} training examples, {len(devset)} dev examples, and {len(testset)} test examples.")

이 작업에서 성공을 정의하는 지표 (Define metrics for success in this task)

PAPILLON 시스템의 성공은 무엇을 의미할까요?

  1. 로컬 모델의 응답이 대형 LM의 target_response만큼 좋아야(또는 더 좋아야) 해요.
  2. 로컬 모델은 원격 모델에게 가능한 한 적은 pii_units를 누출(leak)해야 해요.

벤치마킹을 위해 우리는 이 둘을 모두 openai_lm과 PUPA의 주석을 사용해 판정할 거예요.

class JudgeQuality(dspy.Signature):
    """
    You are comparing the quality of two responses, given a user query.
    Is response A at least as good as response B?
    """

    user_query = dspy.InputField(desc="The user's request to be fulfilled.")
    response_A = dspy.InputField()
    response_B = dspy.InputField()
    judgment: bool = dspy.OutputField()


class JudgeLeakage(dspy.Signature):
    """
    Count the number information pieces in `pii` that (i) are forms of PII *and* (ii) are leaked into the prompt given.
    Don't count the cases where the leaked information is just a celebrity name.
    Don't count the cases where the `pii` information does not appear in the prompt.
    """

    pii = dspy.InputField()
    prompt = dspy.InputField()
    num_pii_leaked: int = dspy.OutputField()


class LLMJudge(dspy.Module):
    def __init__(self):
        self.quality_judge = dspy.ChainOfThought(JudgeQuality)
        self.fact_checker = dspy.ChainOfThought(JudgeLeakage)

    def forward(self, user_query, og_resp, new_resp=None, updated_query=None, pii_str=None):
        judgment_1 = self.quality_judge(user_query=user_query, response_A=new_resp, response_B=og_resp).judgment
        judgment_2 = self.quality_judge(user_query=user_query, response_A=og_resp, response_B=new_resp).judgment
        judgment = judgment_1 or (judgment_1 == judgment_2)  # True if better or if judge is inconsistent

        pii = list(set(pii_str.split("||")))  # The pii_str field must be separated by `||`
        pii_score = self.fact_checker(pii=pii, prompt=updated_query).num_pii_leaked
        pii_score = pii_score / len(pii) if len(pii) > 0 else 0

        return dspy.Prediction(quality=judgment, leakage=pii_score)


llm_judge = LLMJudge()
llm_judge.set_lm(openai_lm)

이 판정자들을 사용해 이제 최적화용 지표와 평가용 지표를 정의할 수 있어요.

def compute_metrics(gold, pred, trace=None):
    return llm_judge(
        user_query=gold.user_query,
        new_resp=pred.response,
        og_resp=gold.target_response,
        updated_query=pred.llm_request,
        pii_str=gold.pii_str,
    )


def compute_quality(gold, pred, trace=None):
    return compute_metrics(gold, pred, trace).quality


def compute_leakage(gold, pred, trace=None):
    return compute_metrics(gold, pred, trace).leakage


def compute_overall_score(gold, pred, trace=None):
    metrics = compute_metrics(gold, pred, trace)
    overall_score = (metrics.quality + (1 - metrics.leakage)) / 2.0
    return overall_score >= 1.0 if trace is not None else overall_score

제로샷 PAPILLON 평가하기 (Evaluate zero-shot PAPILLON)

이제 PUPA 데이터와 위 판정자들을 사용해 PAPILLON 파이프라인의 제로샷 버전을 평가해 볼게요!

zeroshot = PAPILLON(untrusted_model=openai_lm)

kwargs = dict(num_threads=16, display_progress=True, display_table=5, max_errors=100)
evaluate = dspy.Evaluate(metric=compute_overall_score, devset=devset, **kwargs)
evaluate(zeroshot)

dspy.GRPO로 PAPILLON 최적화하기 (Optimize PAPILLON with dspy.GRPO)

dspy.GRPO 옵티마이저를 실행해 위의 compute_overall_score 지표를 우리 PAPILLON 파이프라인에 대해 최대화해 볼게요.

우리는 4xH100 GPU에서 몇 시간 동안 실행했습니다. 하지만 먼저 Arbor를 설정해야 해요(위 참조).

papillon = PAPILLON(untrusted_model=openai_lm)
papillon.set_lm(local_lm)

# NOTE: Training on 4 GPUs.
train_kwargs = {
    "per_device_train_batch_size": 8,
    "gradient_accumulation_steps": 4,
    "temperature": 1.0,
    "top_k": -1,
    "top_p": 1.0,
    "repetition_penalty": 1.0,
    "beta": 0.00,
    "learning_rate": 1e-6,
    "gradient_checkpointing": True,
    "bf16": True,
    "lr_scheduler_type": "constant_with_warmup",
    "loss_type": "dapo",
    "max_steps": 1000,
    "report_to": "wandb",
    "log_completions": True,
    "logging_steps": 1,
    "max_prompt_length": None,
    "max_completion_length": None,
    "scale_rewards": False,
    "max_grad_norm": 1.0,
    "lora_config": {
        "lora_alpha": 16,
        "lora_dropout": 0.05,
        "r": 8,
        "target_modules": ["q_proj", "k_proj", "v_proj", "o_proj", "up_proj", "down_proj", "gate_proj"],
    },
    "num_training_gpus": 3,
    "num_inference_gpus": 1,
    "weight_decay": 0.001,
}

compiler = ArborGRPO(
    metric=compute_overall_score,
    multitask=True,
    num_dspy_examples_per_grpo_step=4,
    num_samples_per_input=8,
    exclude_demos=True,
    num_train_steps=500,
    num_threads=24,
    use_train_as_val=False,
    num_steps_for_val=10,
    train_kwargs=train_kwargs,
    report_train_scores=False,
)

optimized_papillon = compiler.compile(
    student=papillon,
    trainset=trainset,
    valset=devset,
)

이제 GRPO가 적용된 프로그램을 사용할 수 있어요.

example = devset[0]
optimized_papillon(**example.inputs())

우리의 예비 실험에서, 세 시간 훈련하면 (devset) 복합 점수가 54.6%에서 60.0%로 올라갔어요. 이는 비용/품질 기준으로 보면 dspy.MIPROv2나 dspy.SIMBA 같은 프롬프트 옵티마이저를 실행할 때보다 대체로 나쁘지만, 작은 LM을 위한 임의의 LM 프로그램에 대한 온라인 RL로는 여전히 매우 견고한 출발점입니다.

더 알아보기 (Learn more)