WEB GPU를 통한 로컬 LLM 서비스 구현하기 (2)

deepseek API 쓰다가 보안 취약점을 고민하게 된 이후, 로컬에서 LLM을 돌리면 보안 문제는 물론 네트워크 딜레이도 확 줄일 수 있겠다는 생각이 들었음. 그래서 이번엔 transformers.js 기반의 파이프라인을 만들고, webGPU를 활용해서 모델이 로딩되는 이벤트를 받아 UI에 로딩 진행 상황을 보여주는 기능을 구현해봤다.

마찬가지로 deepseek R1 1.5B 모델을 ONNX 기반으로 돌리는 걸 참고했는데, 기존에 Web Worker를 사용하던 방식 대신에 싱글톤 Controller 패턴으로 직접 제어해보려고 했음. 코드 한 번 보자.

1. 파이프라인 구현 및 모델 로딩 이벤트

1.1 TextGenerationPipeline 클래스

모델과 토크나이저를 싱글톤으로 관리하기 위해 아래와 같이 TextGenerationPipeline 클래스를 정의했다.
이 클래스는 모델 ID, 토크나이저, 그리고 모델을 Promise로 관리하며,
getInstance 메소드를 통해 로딩 상태를 콜백으로 받아올 수 있게 되어 있다.

export class TextGenerationPipeline {
    static model_id = "onnx-community/DeepSeek-R1-Distill-Qwen-1.5B-ONNX";
    static tokenizer: Promise<PreTrainedTokenizer> | null = null;
    static model: Promise<AutoModelForCausalLM> | null = null;

    static async getInstance(progress_callback: ((x: any) => void) | null = null): Promise<[PreTrainedTokenizer, AutoModelForCausalLM]> {
        this.tokenizer ??= AutoTokenizer.from_pretrained(this.model_id, {
            progress_callback: progress_callback || undefined,
        });

        this.model ??= AutoModelForCausalLM.from_pretrained(this.model_id, {
            dtype: "q4f16",
            device: "webgpu",
            progress_callback: progress_callback || undefined,
        });

        return Promise.all([this.tokenizer, this.model]);
    }
}

여기서 주목할 점은 progress_callback 옵션을 통해 모델과 토크나이저가 로딩되는 중간 과정을 실시간으로 받아올 수 있다는 것이다. 이 부분 덕분에 UI에 로딩바나 진행 상황을 표시할 수 있다.

1.2 LLMController 클래스와 이벤트 전파

싱글톤 Controller 패턴을 도입해, 모델 로딩 과정을 이벤트로 전달하기 위해 LLMController를 만들었음.
이 컨트롤러는 이벤트 에미터를 활용해서 아래와 같이 동작함:

export class LLMController {
    async load() {
        // 모델 로딩 시작 알림
        self.postMessage({ type: WORKER_STATUS.STATUS_LOADING });

        // TextGenerationPipeline.getInstance를 호출하면서 진행 상황을 이벤트로 전달
        await TextGenerationPipeline.getInstance((progressEvent) => {
            self.postMessage(progressEvent);
        });

        // 모델 로딩 완료 알림
        self.postMessage({ type: WORKER_STATUS.STATUS_READY });
    }
}

이렇게 구현함으로써, 로딩 시작 시 "STATUS_LOADING" 메시지를 보내고, 로딩 중간에 진행 이벤트들을 받아 UI에 반영한 다음, 최종적으로 "STATUS_READY" 메시지로 모델이 준비되었음을 알리도록 했다.

2. 구현 후기 및 앞으로의 계획

이번 파트에서는 transformers.js를 참고해서 WebGPU 기반의 로컬 LLM 파이프라인을 구성하는 데 집중했다.
싱글톤 Controller 패턴을 활용해 Worker 없이 메인 스레드에서 모델 로드 및 이벤트 처리를 하니까, 디버깅과 관리가 훨씬 수월했지만 메인스레드에 작업이 전부 처리되어 나중에 UI를 구현 했을 때 과연 문제가 없을지 의문이다...

앞으로의 계획:

추론 기능 연동:

모델 로딩이 완료된 후, 실제 텍스트 생성(추론) 기능을 채팅 UI와 연동해 사용자에게 실시간 응답을 제공할 예정이다.

에러 핸들링 및 최적화:

모델 로딩 중 발생할 수 있는 에러를 좀 더 정교하게 처리하고, 메모리 사용 최적화에도 신경 쓸 계획이다.
이번 포스팅은 WEB GPU를 활용한 로컬 LLM 서비스 구현의 두 번째 파트로, 모델 로딩과 UI 연동에 초점을 맞췄다.
앞으로 더욱 다양한 기능과 최적화 방안을 공유할 예정이니, 관심 있으면 계속 지켜봐주길 바란다.