레이어별

레이어별 (재)로딩이란? (What is Layerwise (Re)loading?)

포스트 트레이닝(사후 훈련) 워크플로에서는 트레이너가 만든 새 가중치를 이미 떠 있는 추론 인스턴스에 넣어야 할 때가 있어요. 문제는 그때 CUDA 그래프나 다른 런타임 아티팩트를 다시 컴파일하고 싶지 않다는 거죠. 레이어별 재로딩(layerwise reloading)은 바로 그 상황을 처리하는 시스템이에요. 옆에서 설명하자면, 기존 가중치 데이터 목적지에 새 가중치 데이터를 "레이어 단위로" 끼워 넣어 CUDA 그래프 재컴파일을 피하는 방식이에요.

출처: vLLM 공식 문서 — layerwise

레이어별 재로딩은 무엇을 하는가

레이어별 재로딩은 새 가중치 데이터를 기존 가중치 데이터 목적지에 로드하면서, CUDA 그래프나 다른 런타임 아티팩트의 재컴파일을 유발하지 않는 시스템이에요. 이 시스템은 QeRL 스타일의 포스트 트레이닝 플로우를 가능하게 하기 위해 사용돼요. 전체 정밀도(full-precision)의 트레이너 가중치가 양자화되어 빠른 고탐험 롤아웃을 위한 대상 vLLM 인스턴스에 로드되는 방식이죠. 핵심 구현은 layerwise.py에서 찾을 수 있어요.

QeRL을 위한 레이어별 재로딩

새 가중치를 기존 가중치 데이터 목적지에 로드하기 위해선 가중치가 다음 작업들을 거쳐야 해요.

• Transfer (전송): 가중치를 트레이너 모델에서 대상 노드/디바이스로 전송
• Fuse (융합): 가중치 파티션을 융합, 예: qkv/gate_up
• Process (처리): 일반적으로 온라인 양자화와 커널별 패딩 또는 스트라이딩을 의미
• Shard (분할): 선택한 병렬 전략에 따라 가중치를 분할
• Copy (복사): 가중치를 기존 가중치 데이터 목적지에 복사

레이어별 재로딩은 다음 단계로 이를 달성해요.

1. 가중치를 트레이너에서 대상으로 전송 (자세한 내용은 weight_transfer 참고)
2. model.load_weights로 가중치를 로드하며, 이때 분할 및 융합됨
3. 레이어의 모든 가중치가 로드되는 즉시 온라인 방식으로 가중치를 처리
4. 가중치를 기존 가중치 데이터 목적지에 복사

구현에 대한 자세한 내용은 Low Level layerwise API를 참고하세요.

온라인 양자화와 함께하는 레이어별 로딩

온라인 양자화(online quantization)는 사용자가 전체 정밀도 가중치를 제공하고 그 가중치가 모델에 로드될 때 즉석에서 양자화되는 경우를 말해요. 레이어별 재로딩 시스템은 온라인 양자화를 처리(processing) 단계로 취급해서 이를 다룹니다. 그리고 이 처리는 첫 로드 때와 재로드 때 모두 온라인 방식으로 처리되죠. 일반적인 온라인 양자화 메서드 구현은 다음과 같아요.

class Fp8OnlineLinearMethod(Fp8LinearMethod):
    """Online version of Fp8LinearMethod which loads a full precision
    checkpoint and quantizes weights during loading."""
    uses_meta_device: bool = True

    def create_weights(self, layer: torch.nn.Module, ...):
        # weight is materialized and processed during loading
        layer.weight = ModelWeightParameter(
            data=torch.empty(..., device="meta"),
            weight_loader=weight_loader,
        )
        # set up online processing
        initialize_online_processing(layer)

    def process_weights_after_loading(self, layer: Module) -> None:
        if getattr(layer, "_already_called_process_weights_after_loading", False):
            return
        layer.weight, layer.weight_scale = ops.scaled_fp8_quant(layer.weight)
        # Prevent duplicate processing (e.g., during weight reload)
        layer._already_called_process_weights_after_loading = True

사용 예시 (Example Usages)

고수준 Weight Transfer API

레이어별 재로딩 시스템은 포스트 트레이닝 가중치 전송 시스템과 통합되어 있어요. 가중치 전송 시스템과 함께 레이어별 재로딩을 쓰려면 여기에 있는 예시를 따라가세요. 레이어별 재로딩은 WeightTransferUpdateInfo.is_checkpoint_format 플래그로 제어되며 기본값은 True예요.

중간 수준 reload_weights API

레이어별 재로딩은 reload_weights API로도 노출돼요. 이 인터페이스는 다음 코드로 호출할 수 있습니다.

from vllm import LLM

llm = LLM("Qwen/Qwen3-0.6B")
llm.collective_rpc("reload_weights")

이 인터페이스는 로드할 체크포인트 경로를 선택할 수 있는 weights_path 지정도 허용해요.

from vllm import LLM

# fine tuned model checkpoints for testing
mul_path = "inference-optimization/Qwen3-0.6B-debug-multiply"
add_path = "inference-optimization/Qwen3-0.6B-debug-add"

llm = LLM("Qwen/Qwen3-0.6B")
llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_path})
llm.generate("3 4 = ")  # 12

llm.collective_rpc("reload_weights", kwargs={"weights_path": add_path})
llm.generate("3 4 = ")  # 7

마지막으로 weights_iterator를 직접 제공할 수도 있어요. 이 이터레이터는 lazy 또는 eager하게 정의할 수 있습니다.

from vllm import LLM

weights_iterator = [("q_proj", ...), ("k_proj", ...), ...]
llm = LLM("Qwen/Qwen3-0.6B")
llm.collective_rpc("reload_weights", kwargs={"weights_iterator": weights_iterator})

저수준 layerwise API

layerwise.py는 그 생애주기를 실행하기 위해 다음 함수들을 구현해요.

이 라이프사이클에 직접 플러그인할 수도 있어요. initialize_layerwise_reload를 호출하고 가중치를 로드한 뒤 finalize_layerwise_processing을 호출하면 됩니다.

from vllm import LLM
from vllm.model_executor.model_loader.reload import (
    initialize_layerwise_reload, finalize_layerwise_processing)

llm = LLM("Qwen/Qwen3-0.6B")
# this model path requires `VLLM_ENABLE_V1_MULTIPROCESSING=0` and is not stable
model = llm.llm_engine.engine_core.engine_core.model_executor.driver_worker.worker.get_model()
# layerwise reload
initialize_layerwise_reload(model)
model.load_weights(...)
finalize_layerwise_processing(model, llm.model_config)

과도한 메모리 사용 문제 해결 (Troubleshooting Excessive Memory Usage)

레이어별 재로딩은 사용자가 가중치를 모델에 로드하면서 점진적으로 로드·처리할 수 있게 해줘요. 이 시스템은 레이어의 모든 가중치가 로드될 때까지 디바이스에 레이어 가중치를 버퍼링하는 데 의존합니다. 하지만 오프로딩 없이는, 가중치가 순서에 맞지 않게(out of order) 로드되면 과도한 버퍼링이 생길 수밖에 없어요.

그래서 사용자는 모델에 재로딩할 때 가중치의 순서에 주의해야 해요. 가중치는 "순서대로(in order)" 로드되어야 하고, 즉 각 레이어의 가중치가 다음 레이어의 가중치를 로드하기 전에 완전히 채워져야 합니다. "순서에 맞지 않는" 로딩은 다른 레이어 가중치가 로드되는 동안 레이어 가중치가 버퍼에 남아 있어 과도한 메모리를 사용하게 만들 수 있어요. 아래 예시에서 q_proj, k_proj, v_proj, up_proj가 모두 동시에 버퍼링되어, up_proj가 q_proj, k_proj, v_proj 이후에 로드된 경우보다 더 많은 메모리를 쓰게 됩니다.

가중치가 순서에 맞지 않게 로드되면 아래와 같은 경고가 나타나요.

WARNING [layerwise.py:198] Allocating 28.5 MB of device memory to buffers to load ["QKVParallelLinear", "MergedColumnParallelLinear"] layers. This extra memory usage can be avoided by ordering weights by their parent layer when reloading.

더 알아보기 (Learn more)

• 가중치 전송 (Weight Transfer)
• NCCL 가중치 전송
• 비동기 강화학습 (Async RL)