엔진 초기화에서 가중치 로딩 건너뛰기

엔진 초기화에서 가중치 로딩 건너뛰기 (Skip Loading Weights In Engine Init)

엔진 초기화 시 가중치 로딩을 건너뛰는 방법을 보여주는 예제입니다. load_format="dummy"로 빈 가중치로 LLM을 띄운 뒤, collective_rpc로 load format을 auto로 바꾸고 실제 가중치를 제자리(in-place)로 다시 로드해 출력이 정상화되는 것을 확인합니다. RL/튜닝 워크플로에서 가중치를 나중에 주입하는 상황에 유용합니다.

출처: 문서

원본: https://github.com/vllm-project/vllm/blob/main/examples/rl/skip_loading_weights_in_engine_init.py

본문

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from vllm import LLM, RequestOutput, SamplingParams

# Sample prompts.
prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]
# Create a sampling params object.
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

def print_prompts_and_outputs(outputs: list[RequestOutput]) -> None:
    print("-" * 60)
    for output in outputs:
        prompt = output.prompt
        generated_text = output.outputs[0].text
        print(f"Prompt:    {prompt!r}")
        print(f"Output:    {generated_text!r}")
        print("-" * 60)

def main():
    # Create an LLM without loading real weights
    llm = LLM(
        model="Qwen/Qwen3-0.6B",
        load_format="dummy",
        enforce_eager=True,
        tensor_parallel_size=4,
    )
    outputs = llm.generate(prompts, sampling_params)
    print("\nOutputs do not make sense:")
    print_prompts_and_outputs(outputs)

    # Update load format from `dummy` to `auto`
    llm.collective_rpc(
        "update_config", args=({"load_config": {"load_format": "auto"}},)
    )
    # Now reload real weights inplace
    llm.collective_rpc("reload_weights")

    # Check outputs make sense
    outputs = llm.generate(prompts, sampling_params)
    print("\nOutputs make sense after loading real weights:")
    print_prompts_and_outputs(outputs)

if __name__ == "__main__":
    main()

흐름을 정리하면 (1) load_format="dummy"로 실제 가중치 없이 LLM을 생성해 무의미한 출력을 확인하고, (2) collective_rpc("update_config", ...)로 load format을 auto로 바꾼 뒤 (3) collective_rpc("reload_weights")로 실제 가중치를 제자리 로드해, 이후 출력이 의미를 가지는지 확인합니다.

더 알아보기 (Learn more)