레이어별

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

레이어별 재로딩(layerwise reloading)은 cuda graph와 기타 런타임 아티팩트의 재컴파일을 유발하지 않고, 새 가중치 데이터를 기존 가중치 데이터 저장소에 로드하는 것을 처리하는 시스템입니다. 이 시스템은 QeRL 스타일의 post-training 흐름을 가능하게 하며, 전체 정밀도 트레이너 가중치를 양자화해 타겟 vLLM 인스턴스에 로드함으로써 빠르고 높은 탐험(high-exploration) rollout을 만듭니다.

출처: 문서

본문

레이어별 재로딩은 cuda graph와 기타 런타임 아티팩트의 재컴파일을 유발하지 않고 새 가중치 데이터를 기존 가중치 데이터 저장소에 로드하는 시스템입니다. 이 시스템은 QeRL 스타일의 post-training 흐름을 가능하게 하는데, 여기서 전체 정밀도 트레이너 가중치가 양자화되어 빠르고 높은 탐험 rollout을 위해 타겟 vLLM 인스턴스에 로드됩니다. 핵심 구현은 layerwise.py에 있습니다.

QeRL을 위한 레이어별 재로딩 (Layerwise Reloading for QeRL)

새 가중치를 기존 가중치 데이터 저장소에 로드하려면 가중치는 다음 연산을 거쳐야 합니다:

Transfer: 가중치를 트레이너 모델에서 타겟 노드/디바이스로 전송
Fuse: 가중치 파티션을 융합, 예: qkv/gate_up
Process: 보통 온라인 양자화와 커널별 패딩·스트라이딩을 의미
Shard: 선택한 병렬화 전략에 따라 가중치를 샤딩
Copy: 가중치를 기존 가중치 데이터 저장소에 복사

레이어별 재로딩은 다음 단계로 이를 달성합니다:

가중치를 트레이너에서 타겟으로 전송(transfer)(weight_transfer 참고)
model.load_weights로 가중치를 로드하는데, 이때 샤딩되고 융합됩니다
레이어의 모든 가중치가 로드되는 즉시 가중치를 온라인 방식으로 처리(process)
가중치를 기존 가중치 데이터 저장소로 복사

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

온라인 양자화와 레이어별 로딩 (Layerwise Loading with Online Quantization)

온라인 양자화는 사용자가 전체 정밀도 가중치를 제공하고, 그 가중치가 모델에 로드되면서 즉석(on-the-fly)으로 양자화되는 것을 말합니다. 레이어별 재로딩 시스템은 온라인 양자화를 처리(processing) 단계로 취급해, 최초 로드와 재로드 중 모두 온라인 방식으로 처리합니다. 일반적인 온라인 양자화 메서드 구현은 다음과 같습니다:

class Fp8PerTensorOnlineLinearMethod(LinearMethodBase):
    """Online version of FP8 per-tensor quantization 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)

하이레벨 가중치 전송 API (High Level Weight Transfer API)

레이어별 재로딩 시스템은 post-training 가중치 전송 시스템과 통합됩니다. 가중치 전송 시스템과 함께 레이어별 재로딩을 사용하려면 여기의 예시를 따르세요. Checkpoint 형식 가중치 전송 엔진(예: NCCL, IPC 백엔드)은 start_weight_update/finish_weight_update 수명주기 안에서 레이어별 재로딩을 자동으로 실행합니다.

중간 레벨 reload_weights API (Mid Level reload_weights API)

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

from vllm import LLM

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

이 인터페이스는 로드할 checkpoint 경로를 선택하는 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 (Low Level layerwise API)

layerwise.py는 수명주기를 실행하기 위해 다음 함수들을 구현합니다:

함수	용도	양자화된 재로드	온라인 양자화
`record_metadata_for_reloading`	레이어를 meta 디바이스에 복원할 수 있도록 텐서 메타데이터 기록	`BaseModelLoader`가 호출	`BaseModelLoader`가 호출
`restore_layer_on_meta`	재로드 시작 시 레이어를 모델 형식으로 복원	`initialize_layerwise_reload`가 호출	호출 안 함. 온라인 양자화 가중치는 `...OnlineLinearMethod.create_weights`를 통해 이미 meta 디바이스에서 시작
`initialize_online_processing`	모든 레이어 가중치가 로드될 때까지 가중치를 버퍼링하는 `online_process_loader` 래퍼로 가중치 로더를 감쌈	`initialize_layerwise_reload`가 호출	`...OnlineLinearMethod.create_weights`가 호출
`_layerwise_process`	모든 가중치가 로드되면 레이어 처리	로딩 중 `online_process_loader`가 호출	로딩 중 `online_process_loader`가 호출
`_copy_and_restore_kernel_tensors`	컴파일된 cuda graph에 영향을 주도록 처리된 가중치를 원래 텐서 위치에 복사	`process_weights_after_loading` 후 `_layerwise_process`가 호출	호출 안 함. 아직 컴파일된 cuda graph 없음
`finalize_layerwise_processing`	모든 가중치를 로드하지 않은 레이어(예: 어텐션 가중치나 패딩이 있는 가중치)를 잡아냄	`BaseModelLoader`가 호출	`BaseModelLoader`가 호출

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"로 로드해야 합니다. 즉 다음 레이어의 가중치를 시작하기 전에 각 레이어의 가중치가 완전히 로드되어야 합니다. "out of 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)

가중치 전송 — 학습-추론 가중치 동기화
비동기 RL — 가중치 업데이트 pause/resume API
가중치 전송: IPC — IPC 엔진