from_config 모델 로딩이 30초 걸린 이유, 어차피 덮어쓸 난수 초기화

로컬 추론 서버에 분류 모델을 띄우다가, 체크포인트 하나를 로드하는 데 30초가 걸리는 걸 발견했습니다. 가중치는 이미 디스크 캐시에 있는데도요.

모델은 3~4억 파라미터짜리 인코더(ModernBERT 계열) 위에 자체 헤드를 얹은 구조였고, 로딩 코드는 흔한 방식이었습니다.

model = AutoModel.from_config(config)
model.load_state_dict(load_file("model.safetensors"), strict=True)

30초는 어디에 쓰였을까?

cProfile 을 돌려 보니 30.8초 중 22.8초가 post_init() 안의 _init_weights, 정확히는 trunc_normal_ 이었습니다.

from_config 는 새 모델을 학습하려는 것으로 보고 모든 Linear 와 Embedding 을 난수로 초기화합니다. 그런데 바로 다음 줄의 load_state_dict(strict=True) 가 그 값을 전부 덮어쓰죠. 수억 개의 절단 정규분포 난수를 CPU 로 뽑아서 곧바로 버리는 셈입니다.

from_pretrained 는 이 초기화를 알아서 건너뜁니다. 하지만 from_config 로 뼈대를 만들고 직접 로드하는 경로는 그렇지 않습니다.

초기화를 끄고 뼈대를 만들면 됩니다

어차피 전부 덮어쓸 거라면 초기화를 끄면 됩니다. transformers 5.x 에는 no_init_weights() 컨텍스트가 있습니다.

from transformers import AutoConfig, AutoModel
from transformers.initialization import no_init_weights   # transformers 5.x
from safetensors.torch import load_file
 
cfg = AutoConfig.from_pretrained(encoder_dir)
with no_init_weights():                       # 난수 초기화 생략
    model = AutoModel.from_config(cfg)
model.load_state_dict(load_file(weights_path), strict=True)   # 전부 덮어쓰기, strict 필수

로드 시간은 이렇게 줄었습니다.

같은 입력에 대한 확률 dict 도 초기화를 건 쪽과 == 로 비트 단위까지 같았습니다.

건너뛰어도 안전할까요?

조건은 하나입니다. strict=True 로 모든 파라미터가 덮어써져야 합니다. 누락된 키가 있으면 그 파라미터는 초기화되지 않은 메모리 값으로 남기 때문입니다.

다행히 strict=True 가 누락을 에러로 잡아 주니, 같이 쓰면 안전합니다. RoPE 의 inv_freq 처럼 __init__ 에서 계산되는 비영속 버퍼는 초기화 함수와 상관이 없어서 영향이 없습니다.

라이브러리 코드를 고칠 수 없다면 호출하는 쪽에서 로드 함수 전체를 with no_init_weights(): 로 감싸도 됩니다. 이 경우 자체 헤드의 nn.Linear 초기화까지 같이 생략됩니다.

transformers 밖이라면 PyTorch 쪽 방법이 같은 효과를 냅니다. torch.device("meta") 위에서 모듈을 만들고 load_state_dict(..., assign=True) 를 쓰거나, torch.nn.utils.skip_init 을 쓰면 됩니다.

로딩이 느리면 디스크부터 의심했는데, 이번엔 CPU 가 쓸데없는 난수를 뽑고 있었네요. 메모 겸 남겨 둡니다.

참고

• https://pytorch.org/tutorials/prototype/skip_param_init.html
• https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.load_state_dict
• https://huggingface.co/docs/transformers/main_classes/model