동적 가중치 로딩

동적 가중치 로딩 (Dynamic weight loading)

체크포인트는 종종 모델이 런타임에 기대하는 형식과 일치하지 않는 형식으로 직렬화돼요. 흔한 시나리오는 다음과 같아요:

출처: 문서

본문

  1. 퓨전된 가중치(Fused weights): 체크포인트는 별도의 gate_proj와 up_proj 가중치를 저장하지만, 모델은 효율성을 위해 퓨전된 gate_up_proj를 사용해요.
  2. MoE 전문가 통합: 개별 전문가 가중치(experts.0.weight, experts.1.weight, ...)를 단일 3D 텐서로 쌓아야 해요.
  3. 레거시 명명: 이전 체크포인트는 다른 명명 규칙을 사용해요 (예: LayerNorm.gamma vs LayerNorm.weight).
  4. 복합 모델(Composite models): vision-language 모델에는 각각 자체 체크포인트 규칙이 있는 두 개의 PreTrainedModel 하위 모듈이 포함돼요.
  5. 양자화(Quantization): 가중치가 역직렬화가 필요한 양자화된 형식으로 저장될 수 있어요.

동적 가중치 로딩은 체크포인트 텐서가 로드될 때 스케줄링된 가역(reversible) 연산을 적용해서 이 문제를 해결해요. Transformers는 이를 WeightConverter와 WeightRenaming으로 노출하는데, 이들은 하나 이상의 체크포인트 키가 하나 이상의 모델 파라미터에 어떻게 매핑되는지, 그리고 어떤 조합 가능한 ConversionOps를 매칭된 텐서에 실행해야 하는지 설명해요. 이 접근 방식은 새로운 가중치 레이아웃에 적응하고, 양자화된 mixture-of-experts(MoE)를 지원하며, 텐서 병렬화와 통합돼요.

이 가이드는 WeightConverter를 사용해 텐서를 변환하는 방법을 보여줘요. 변환 매핑은 conversion_mapping.py에 있으며, 등록된 매핑은 model_type 문자열(예: "mixtral") 또는 클래스 이름(예: "LlavaModel")으로 키가 지정돼요.

전체 로딩 파이프라인 (Full loading pipeline)

모든 모델은 동적 가중치 로딩 시스템을 통과해요. 변환 매핑은 그 시스템 안의 선택적 단계로, 모델에 클래스 이름이나 model_type으로 등록된 항목이 있을 때만 활성화돼요.

Checkpoint File → from_pretrained() → convert_and_load_state_dict_in_model()
                                              ↓
                         ┌───────────────────────────────────────────────────────────┐
                         │  For each weight in checkpoint:                           │
                         │  1. Match renamed/processed source key to model parameter │
                         │  2. Shard the weight and send to device (async)           │
                         │  3. Collect tensors with the same source_pattern together │
                         │     (e.g. MoE experts, gate_up_proj)                      │
                         │  4. Apply dequantization/deserialization (if pre-quant)   │
                         │  5. Apply conversion (if defined)                         │
                         │  6. Apply quantization (if enabled and step 4 not used)   │
                         │  7. Set parameter on model                                │
                         └───────────────────────────────────────────────────────────┘
Step When it activates
Dynamic loading Always, for all models
Conversion mapping Only when the model's class or model_type is registered in _MODEL_TO_CONVERSION_PATTERN
TP sharding Only when DistributedConfig(tp_size=N) is used and the model has base_model_tp_plan
Dequantization/deserialization Only when loading a pre-quantized checkpoint
Quantization Only when a quantization config is provided and weights are not pre-quantized

밀집 모델 (Dense models, 예: Llama)

대부분의 밀집 모델에서 체크포인트 형식은 모델 형식과 직접 일치하므로 변환 매핑이 필요 없어요. 일부 모델은 여전히 이름 바꾸기(예: 레거시 명명 규칙)가 필요할 수 있어요. 활성화되면 TP 샤딩은 여전히 적용돼요.

Checkpoint:                             Model:
model.layers.0.self_attn.q_proj.weight  →  model.layers.0.self_attn.q_proj.weight
model.layers.0.self_attn.k_proj.weight  →  model.layers.0.self_attn.k_proj.weight
model.layers.0.mlp.gate_proj.weight     →  model.layers.0.mlp.gate_proj.weight
model.layers.0.mlp.up_proj.weight       →  model.layers.0.mlp.up_proj.weight
model.layers.0.mlp.down_proj.weight     →  model.layers.0.mlp.down_proj.weight

레거시 체크포인트는 모든 모델에 적용되는 내장 이름 바꾸기로 처리되는 이전 명명 규칙을 사용할 수 있어요:

Checkpoint:                          Model:
LayerNorm.gamma              →       LayerNorm.weight
LayerNorm.beta               →       LayerNorm.bias

MoE 모델 (예: Mixtral)

MoE 모델의 경우 체크포인트 형식이 모델 형식과 달라요. 변환 매핑은 별도의 전문가 가중치를 퓨전된 3D 텐서로 변환하고, 변환 후에 TP 샤딩이 적용돼요.

Checkpoint:                              Model:
experts.0.w1.weight  ─┐
experts.1.w1.weight   │ MergeModulelist
...                   ├───────────────→  experts.gate_up_proj (8, hidden, 2*intermediate)
experts.0.w3.weight   │ + Concatenate
experts.1.w3.weight  ─┘

복합 모델 (Composite models, 예: vision-language)

PreTrainedModel은 다른 PreTrainedModel 하위 모듈을 포함할 수 있어요. 각 하위 모델은 클래스 이름이나 model_type으로 등록된 자체 변환 매핑을 가질 수 있어요. 부모 모델이 로드될 때 get_model_conversion_mapping은 하위 모델을 깊이 우선(DFS) 순서로 순회하고, 매핑을 수집하며, scope_prefix를 통해 각 변환을 하위 모듈의 경로로 자동으로 스코프해요.

Composite model:                              Per-submodel mappings (auto-scoped):
LlavaForConditionalGeneration
├── vision_model: SiglipVisionModel    →     SiglipVisionModel mapping (scope="vision_model")
└── language_model: LlamaForCausalLM   →     LlamaForCausalLM mapping  (scope="language_model")

scope_prefix는 하위 모듈의 점 표기 경로("vision_model", "language_model.model" 등)예요. 스코프된 변환은 f"{scope_prefix}."로 시작하는 키에서만 발생하고, 패턴 매칭 전에 접두사가 제거되고 치환 후 다시 붙어요. 따라서 각 하위 모델의 매핑은 루트인 것처럼 하위 모델 기준으로 작성돼요.

아키텍처

시스템은 core_model_loading.py에 정의된 몇 가지 핵심 구성 요소를 중심으로 구축돼요:

Phase 1 — 키별 처리 (체크포인트 키를 반복):

  1. 변환 목록을 한 번 순회해요. 매칭되는 모든 WeightRenaming이 발생하고(예: block_sparse_moe → mlp), 기껏해야 하나의 WeightConverter가 키를 차지(claim)할 수 있어요(예: experts.*.w1.weight).
  2. ThreadPoolExecutor로 샤드(TP)하고 장치로 비동기 전송해요.
  3. 같은 source_pattern을 가진 텐서를 함께 수집해요 (예: 모든 MoE 전문가 가중치, gate + up 프로젝션).

Phase 2 — 매핑별 처리 (수집된 매핑을 반복):

  1. 역양자화/역직렬화 (사전 양자화된 체크포인트 전용).
  2. ConversionOps 체인 적용: Chunk, Concatenate, MergeModulelist, Transpose 등.
  3. 즉시 양자화 (사전 양자화되지 않은 경우).
  4. 모델에 파라미터 설정.

WeightTransform

패턴 매칭과 텐서 수집을 처리하는 기본 클래스예요:

  • 패턴 컴파일: 소스 패턴은 re.search()로 매칭되는 완전한 정규식이에요. * 와일드카드는 인덱싱 가능한 컴포넌트와 매칭하며 모든 매치를 배치 작업을 위해 모아요.
  • 캡처링 그룹 & 백레퍼런스: 소스 패턴의 캡처링 그룹은 타겟 패턴에서 \1, \2, ... 로 참조해서 이름 바꾸기 동안 하위 문자열(예: 레이어 인덱스)을 보존할 수 있어요.
  • 스코핑: scope_prefix(get_model_conversion_mapping이 하위 모델별로 자동 설정)는 변환을 그 경로 아래의 키로 제한해요. 매칭 전에 접두사가 제거되고 치환 후 다시 붙어요.
  • 키 이름 바꾸기: rename_source_key()는 (스코프 처리를 포함한) regex를 적용하고 (renamed_key, source_pattern)을 반환해서 로더가 어떤 컨버터가 키를 차지했는지 알게 해요.
  • 텐서 수집: add_tensor()는 해석된 텐서(또는 Futures)를 source_pattern 아래에 축적해서, 단일 변환(예: 모든 MoE 전문가 가중치)에 필요한 모든 텐서가 연산 체인이 실행되기 전에 함께 배치되도록 해요.
  • 가역성: reverse_transform()는 source ↔ target 패턴을 바꾸고 각 연산을 역전시켜서, 같은 목록을 뒤집으면 저장(saving)을 구동해요.
  • 추적: was_used()는 변환이 로딩 중 어떤 키와 실제로 매칭됐는지 보고해요. 이는 비전단사(non-bijective) 이름 바꾸기(예: 접두사가 이미 있을 수 있는 접두사를 추가하는 PrefixChange)가 저장 시 대칭적으로 다시 적용될 수 있도록 필요해요.

WeightRenaming

WeightRenaming은 텐서 연산 없이 순수한 키 이름 바꾸기를 위한 특수화된 WeightTransform이에요. WeightConverter와 달리 WeightRenaming은 키를 차지하지 않으므로 여러 이름 바꾸기가 자유롭게 연결될 수 있어요. 로드 경로에서는 모든 이름 바꾸기가 컨버터보다 먼저 실행되고, 저장 경로(역전된 목록)에서는 모든 역전된 이름 바꾸기가 역전된 컨버터 이후에 실행돼요 (위 순서 규칙 참고).

# Legacy checkpoint compatibility
WeightRenaming("LayerNorm.gamma", "LayerNorm.weight")

# Module path changes
WeightRenaming(".block_sparse_moe.", ".mlp.")

PrefixChange는 전체 경로 컴포넌트를 떼거나 붙이는 흔한 경우를 위한 WeightRenaming의 상위 수준 래퍼예요. 선택적 model_prefix는 그 네임스페이스 아래의 키로 연산을 스코프해요:

# "model.layers.bad_prefix.weight" → "model.layers.weight"
PrefixChange(prefix_to_remove="bad_prefix", model_prefix="model.layers")
# "layers.0.weight" → "model.layers.0.weight"
PrefixChange(prefix_to_add="model")

WeightConverter

WeightConverter는 수집된 텐서에 작용하는 ConversionOps 체인으로 WeightTransform을 확장해요. 지원되는 네 가지 카디널리티(cardinality)는 다음과 같아요:

Cardinality Source patterns Target patterns Typical operation
one-to-one 1 1 Transpose, PermuteForRope
one-to-many 1 >1 Chunk (e.g. unpack qkv_proj)
many-to-one >1 1 Concatenate, MergeModulelist (e.g. fuse experts)
many-to-many >1 >1 only with operations that explicitly support it (e.g. ErnieFuseAndSplitTextVisionExperts) 
WeightConverter(
    source_patterns=[".experts.*.w1.weight", ".experts.*.w3.weight"],
    target_patterns=".experts.gate_up_proj",
    operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
)

WeightConverter는 또한 가역적이에요. reverse_transform()는 source ↔ target을 바꾸고 각 연산을 reverse_op로 대체하므로, 등록된 변환 매핑은 구조적으로 양방향이에요 (로딩은 목록을 그대로 사용하고, 저장은 뒤집어 사용).

이름 바꾸기와 컨버터 순서

WeightRenaming 항목을 WeightConverter 항목보다 먼저 나열하고, 그 잎(leaf)이 서로 분리되도록 유지하세요: 이름 바꾸기는 컨버터가 소비하는 키를 정규화하지만, 컨버터가 만들어내는 잎은 절대 타겟으로 삼지 않아야 해요. 저장 경로는 이에 의존해 매핑을 두 단계(역전된 컨버터, 그다음 역전된 이름 바꾸기)로 역전해요.

변환 목록은 체크포인트 키당 한 번씩 순서대로 순회돼요. 각 키에 대해:

  • 매칭되는 모든 WeightRenaming이 발생해요. 여러 이름 바꾸기가 연결될 수 있어요.
  • 매칭되는 첫 번째 WeightConverter가 키를 차지해요. 이후의 컨버터는 건너뛰어져요. 이는 텐서가 merge/split 단계를 위해 단일 컨버터로 라우팅되도록 보장하며, 의도가 겹치는 두 컨버터가 있다면 잘못된 구성이라는 뜻이에요.
weight_mapping = [
    WeightRenaming("^old_prefix", "encoder"),              # rename runs always
    WeightConverter(                                       # converter claims the key
        "attn.qkv_proj.weight",
        ["attn.q_proj.weight", "attn.k_proj.weight", "attn.v_proj.weight"],
        operations=[Chunk(dim=0)],
    ),
]
# Load:  "old_prefix.attn.qkv_proj.weight"
#   → WeightRenaming  → "encoder.attn.qkv_proj.weight"
#   → WeightConverter → "encoder.attn.{q,k,v}_proj.weight"
#
# Save: list reversed, each transform inverted:
#   → rev(WeightConverter) repacks QKV → "encoder.attn.qkv_proj.weight"
#   → rev(WeightRenaming)  fixes prefix → "old_prefix.attn.qkv_proj.weight"

변환 연산 (Conversion operations)

WeightConverter 클래스에는 from_pretrained()이 호출될 때 체크포인트 소스 텐서를 모델 타겟 텐서로 변환하기 위해 실행되는 몇 가지 연산이 있어요.

연산은 완전히 가역적이에요. 저장은 변환을 역전시키고 원래 체크포인트를 반환하므로 서로 다른 프레임워크에서 쉽게 작업할 수 있어요.

Operation Reverse
Chunk(dim) Concatenate(dim)
Concatenate(dim) Chunk(dim)
MergeModulelist(dim) SplitModulelist(dim)
SplitModulelist(dim) MergeModulelist(dim)
Transpose(d0, d1) Transpose(d1, d0)
PermuteForRope() PermuteForRope()
Conv3dToLinear(...) LinearToConv3d(...)

Chunk

Chunk 연산은 텐서를 차원을 따라 같은 크기의 부분으로 나눠요. 예를 들어 모델이 Q, K, V를 단일 텐서 대신 세 개의 별도 텐서로 기대할 때 사용해요.

WeightConverter(
    "self_attn.qkv_proj",
    ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"],
    operations=[Chunk(dim=0)],
)

Concatenate

Concatenate 연산은 별도의 텐서를 단일 텐서로 퓨전해요. 예를 들어 모델이 Q, K, V를 별도의 텐서 대신 단일 텐서로 기대할 때 사용해요.

WeightConverter(
    ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"],
    "self_attn.qkv_proj",
    operations=[Concatenate(dim=0)],
)

MergeModulelist

MergeModulelist은 2D 텐서 목록을 단일 3D 텐서로 병합해요. 예를 들어 MergeModulelist를 Concatenate와 조합해서 MoE의 전문가들을 쌓고 하나의 텐서로 포장할 수 있어요.

WeightConverter(
    ["block_sparse_moe.experts.*.w1.weight", "block_sparse_moe.experts.*.w3.weight"],
    "mlp.experts.gate_up_proj",
    operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
)

SplitModulelist

SplitModulelist은 3D 텐서를 다시 2D 텐서 목록으로 분할해요. 예를 들어 전문가 스택을 다시 개별 전문가로 나눌 수 있어요.

WeightConverter(
    "mlp.experts.down_proj",
    "block_sparse_moe.experts.*.w2.weight",
    operations=[SplitModulelist(dim=0)],
)

PermuteForRope

PermuteForRope는 interleaved 형식의 가중치를 sin/cos 형식으로 변환해요. 예를 들어 Chunk를 PermuteForRope와 조합해서 퓨전된 QKV 텐서를 나누고 Q와 K에 sin/cos RoPE 순열을 적용할 수 있어요.

WeightConverter(
    ["model.layers.*.self_attn.qkv_proj.weight"],
    [
        "model.layers.*.self_attn.q_proj.weight",
        "model.layers.*.self_attn.k_proj.weight",
        "model.layers.*.self_attn.v_proj.weight",
    ],
    operations=[Chunk(dim=0), PermuteForRope()],
)

Transpose

Transpose는 텐서의 차원을 바꿔요. 서로 다른 규칙 사이의 가중치 레이아웃을 변환하는 데 유용해요.

WeightConverter(
    source_patterns="mlp.gate.weight",
    target_patterns="mlp.text_moe.gate.weight",
    operations=[Transpose(dim0=0, dim1=1)],
)

연산 체이닝 (Operation chaining)

복잡한 변환을 수행하기 위해 연산을 연결할 수 있어요. 연산은 순서대로 실행되며, 각 연산의 출력이 다음 연산의 입력이 돼요.

예시: Mixtral MoE 변환

WeightConverter(
    source_patterns=[
        ".experts.*.w1.weight",  # gate_proj per expert
        ".experts.*.w3.weight",  # up_proj per expert
    ],
    target_patterns=".experts.gate_up_proj",
    operations=[
        MergeModulelist(dim=0),  # Stack all experts: (n_experts, in, out)
        Concatenate(dim=1),      # Fuse gate+up: (n_experts, in, 2*out)
    ],
)

데이터 흐름:

Input:
  ".experts.*.w1.weight": [tensor_0, tensor_1, ..., tensor_7]  # 8 experts
  ".experts.*.w3.weight": [tensor_0, tensor_1, ..., tensor_7]  # 8 experts

After MergeModulelist(dim=0):
  ".experts.*.w1.weight": (8, 4096, 14336)  # stacked gate
  ".experts.*.w3.weight": (8, 4096, 14336)  # stacked up

After Concatenate(dim=1):
  ".experts.gate_up_proj": (8, 4096, 28672)  # fused gate_up

패턴 매칭 규칙

  • 소스 패턴은 re.search()로 평가되는 완전한 정규식이에요. ^는 스코프된 키(접두사 제거 후)의 시작에 고정돼요.
  • *는 인덱스별 와일드카드로, 올바른 concatenation을 위해 체크포인트 순서를 보존하면서 같은 소스 패턴 아래의 모든 매칭 텐서를 수집해요.
  • 소스 패턴의 캡처링 그룹은 원래 키의 일부(예: 레이어 인덱스)를 유지하기 위해 타겟 패턴에서 \1, \2, ... 로 참조할 수 있어요.
  • 스코프된 변환(scope_prefix 설정)은 f"{scope_prefix}."로 시작하는 키에서만 매칭돼요.

변환 매핑 등록 (Registering a conversion mapping)

변환 목록은 문자열 키 — model_type(예: "mixtral") 또는 클래스 이름(예: "LlavaModel") — 에 대해 등록돼요:

from transformers.conversion_mapping import register_checkpoint_conversion_mapping

register_checkpoint_conversion_mapping(
    "my_model_type",
    [
        WeightRenaming(".old.", ".new."),
        WeightConverter(
            ".experts.*.w1.weight",
            ".experts.gate_proj",
            operations=[MergeModulelist(dim=0)],
        ),
    ],
)

조회 규칙 (Lookup rules)

get_model_conversion_mapping이 PreTrainedModel을 처리할 때, 모든 하위 PreTrainedModel을 DFS 순서(nn.Module.named_modules()를 PreTrainedModel 인스턴스로 필터링한 것)로 방문해요. 각 하위 모델은 두 단계로 해석되는데, 먼저 커스텀 코드 필터가 실행된 다음 등록된 매핑이 조회돼요.

커스텀 코드 모델은 조회가 일어나기 전에 건너뛰어져요. is_custom_code()가 True를 반환하면 하위 모델이 커스텀 코드로 간주되는데, 이는 trust_remote_code=True로 Hub에서 로드된 모델과 Transformers에 없는 모든 모델 클래스를 포함해요. 커스텀 코드 모델은 그 클래스 이름이나 model_type이 네이티브 모델과 충돌할 수 있기 때문에 변환 매핑을 상속받지 않아요. 예를 들어 model_type="mixtral"을 설정하는 커스텀 아키텍처는 Mixtral의 내장 전문가 퓨전 변환을 가져오지 않아요.

[!IMPORTANT] 커스텀 코드 모델에 변환 매핑을 적용하려면 register_checkpoint_conversion_mapping으로 명시적으로 등록하세요. 등록은 클래스 이름이나 model_type을 사용자 등록 매핑 집합에 추가하며, 이는 모델을 커스텀 코드 건너뛰기에서 면제하고 아래의 일반 조회를 거치게 해요.

커스텀 코드 필터를 통과한 각 하위 모델에 대해:

  1. 클래스 이름 조회를 먼저 시도하고, 그다음 model_type을 시도해요. 같은 모듈에 둘 다 등록되어 있으면 클래스 이름 매핑이 이기고, 그 모듈에 대해서는 model_type 매핑이 무시돼요. 이렇게 하면 작업 헤드(예: LlavaForConditionalGeneration)가 공유된 model_type 기준("llava")을 재정의할 수 있어요.
  2. 선택된 매핑은 하위 모듈의 점 표기 경로(루트는 "")로 scope_prefix가 설정돼요.
  3. 조상 기반 중복 제거가 매핑을 유지할지 결정해요:
    • 조상 경로가 이미 같은 식별자(클래스 이름 또는 model_type)를 차지했다면 하위 모듈은 건너뛰어져요 — 조상의 스코프되지 않거나 더 높은 스코프의 매핑이 이미 이 하위 트리를 덮기 때문이에요.
    • 형제(sibling) 만이 그것을 차지했다면 하위 모듈은 자체 scope_prefix로 유지돼요. 각 형제는 자체 스코프된 매핑을 가져요.

클래스 이름과 model_type seen-list는 별도로 추적되는데, 한 가지 미묘한 점이 있어요: 모듈이 클래스 이름으로 매칭되면 그 model_type은 seen-list에 추가되지 않아요. 이는 같은 model_type을 공유하지만 클래스별 매핑이 없는 다른 모듈(예: DetrForSegmentation 아래의 DetrModel)이 model_type 조회를 통해 계속 도달 가능하도록 하기 위함이에요.

클래스 기반 매핑 vs model_type 별칭

두 스타일이 _MODEL_TO_CONVERSION_PATTERN에 공존해요:

_MODEL_TO_CONVERSION_PATTERN = {
    # model_type aliases (lookup by config.model_type)
    "minimax": "mixtral",
    "qwen3_moe": "qwen2_moe",
    "mistral3": "llava",
    # class-name aliases (lookup by type(submodule).__name__)
    "PaliGemmaModel": "LlavaModel",
    "MaskFormerDetrDecoder": "DetrModel",
    ...
}

클래스 이름 키는 model_type이 공유되지만 특정 클래스가 다른 동작을 필요로 할 때 선호돼요.

텐서 병렬화 통합 (Tensor parallelism integration)

동적 로딩 시스템은 src/transformers/integrations/tensor_parallel.py에 정의된 TensorParallelLayer 계층 구조를 통해 텐서 병렬화(TP)와 통합돼요.

TP가 활성화되면 텐서는 이후가 아니라 자재화(materialization) 중에 샤딩돼요. 이는 각 랭크가 필요한 텐서 부분만 로드한다는 뜻이에요.

def spawn_tp_materialize(thread_pool, tensor, sharding_method, tensor_idx, device, dtype):
    def _job():
        return sharding_method.shard_tensor(tensor, tensor_idx=tensor_idx, device=device, dtype=dtype)
    return thread_pool.submit(_job)

사용 가능한 병렬 스타일

Style Weight Shard Dim Description
colwise -2 Column-wise: output features sharded
rowwise -1 Row-wise: input features sharded
packed_colwise -2 For fused weights (gate_up_proj)
packed_rowwise -1 For fused weights
embedding_rowwise 0 Vocabulary parallelism
grouped_gemm 0 Expert parallelism for MoE
sequence_parallel None No weight sharding

포장된 가중치 처리 (Packed weight handling)

gate_up_proj 같은 퓨전된 가중치의 경우 올바르게 샤딩하려면 특별한 주의가 필요해요:

def get_packed_weights(param, empty_param, device_mesh, rank, dim):
    """
    Interleaves gate and up shards correctly.

    Packed tensor: [G0 G1 G2 G3 | U0 U1 U2 U3]

    With TP=2:
    - Rank 0 gets: [G0 G1 | U0 U1]
    - Rank 1 gets: [G2 G3 | U2 U3]
    """

TP 연산은 WeightTransform에 저장되고 변환 연산 후에 적용돼요:

if matched_tp_pattern := tp_plan_alt.search(renamed_key):
    tp_layer = ALL_PARALLEL_STYLES[model.tp_plan[matched_tp_pattern]]
    mapping.distributed_operation = tp_layer(
        device_mesh=device_mesh,
        rank=device_mesh.get_local_rank(),
        empty_param=empty_param.clone()
    )

양자화 통합 (Quantization integration)

양자화는 체크포인트가 이미 양자화되었는지 여부에 따라 두 가지 방식으로 로딩 파이프라인에 연결돼요:

  • 사전 양자화된 체크포인트: 양자화기는 (via get_weight_conversions()) 양자화된 텐서를 역직렬화하는 WeightConverter 인스턴스를 제공해요. 원치 않는 캐스트를 피하기 위해 체크포인트 dtype이 보존돼요.
  • 즉시 양자화(On-the-fly): 양자화기는 변환 연산 후에 적용되는 양자화 연산을 제공해서, 가중치가 로드될 때 양자화해요.

양자화기는 또한 get_model_conversion_mapping 끝에서 update_weight_conversions(...)로 전체 변환 목록을 다시 쓸 수 있어요. 예를 들어 FP8 역양자화기는 모든 기존 컨버터 앞에 Fp8Dequantize 연산을 붙여서, 전문가 병합/concat 연산이 전문가별 구조를 평평하게 만들기 전에 블록별 스케일이 적용되도록 해요.

빠르고 효율적인 모델 로딩

로더는 연산에 필요한 텐서를 알고 자재화를 지연(lazily) 스케줄링하므로 모델 로딩이 더 빠르고 메모리를 덜 사용해요.

로더는 체크포인트를 한 번 스캔해서 패턴 매치를 발견하고 텐서를 수집해요. 그것들을 Future 객체로 저장하고 스레드 풀에 제출해서 GIL을 차단하지 않고 비동기로 로드해요. 파라미터는 스레드가 사용 가능해지는 즉시 로드되기 시작해요.

시스템에서 다른 무거운 프로세스가 실행 중이라면 여러 스레드가 로딩을 가속화하기보다는 오히려 느리게 할 수 있어요. 이 경우 환경 변수 HF_DEACTIVATE_ASYNC_LOAD=1를 설정해서 가중치를 순차적으로 로드해요.

[!NOTE] 비동기 파라미터 로딩의 기본값은 4개의 스레드예요. 이는 다양한 로딩 시나리오와 하드웨어에 걸쳐 최상의 절충안을 제공해요. 작업은 대부분 I/O 바운드지만, 가속기 하드웨어와 로딩에 필요한 dtype에 따라 직렬화된 것과 dtype이 다르면(추가 복사 연산 필요) CPU/GPU 바운드가 될 수 있어요.

비동기 vs 동기 로딩

def spawn_materialize(thread_pool, tensor, device, dtype) -> Future | Callable:
    def _job():
        return _materialize_copy(tensor, device, dtype)

    if thread_pool is not None:
        return thread_pool.submit(_job)  # Async: returns Future
    else:
        return _job  # Sync: returns Callable (deferred execution)

동기 로딩은 다음 경우에 사용돼요:

  • HF_DEACTIVATE_ASYNC_LOAD=1 환경 변수가 설정된 경우.
  • 디스크 오프로딩이 활성화된 경우 (메모리 제약으로 순차 로딩 필요).
  • 즉시 양자화가 활성화된 경우 (워커 스레드가 양자화 단계보다 앞서 실행되는 것을 방지).

자재화 흐름

1. Checkpoint iteration (Phase 1):
   - For each key, walk the transform list once
   - Submit materialization job to ThreadPoolExecutor
   - Job returns Future (async) or Callable (sync)
   - Collect into the matching WeightConverter / WeightRenaming

2. Per-mapping processing (Phase 2, one mapping at a time):
   - materialize_tensors() waits for this mapping's Futures only
   - Apply conversion operations chain (self.operations)
   - Apply quantization operation (if on-the-fly)
   - Set parameters on model
   - Delete realized tensors immediately

3. Cleanup:
   - Thread pool shutdown (with cancel_futures=True for interrupts)

메모리 효율

가중치를 변환할 때 컨버터는 아직 로드되지 않은 모든 필수 텐서가 자재화될 때까지 기다려요. 예를 들어 MergeModulelist 연산은 병합 전에 ModuleList의 모든 가중치가 로드되어야 해요.

텐서를 연결하려면 임시 복사가 필요하므로 MergeModulelist와 Concatenate 같은 연산은 변환 중에 기본 텐서의 2배 메모리가 필요해요. 병합되면 결과 텐서만 메모리에 남아요. 이론상 최악의 경우 메모리 피크는 모델 크기 더하기 가장 큰 MergeModulelist 또는 Concatenate 연산에 필요한 텐서예요.

이 최악의 경우는 다른 모든 파라미터가 요구되는 변환보다 먼저 로드되었을 때만 발생해요. 두 가지 시나리오가 이를 유발해요.

  1. 모든 파라미터가 요구되는 변환에 들어가기 전에 비동기로 로드된 경우 (스레드 풀이 변환 큐보다 빨랐던 경우).
  2. 요구되는 변환이 마지막인 경우.

예를 들어 각 레이어의 전문가에 MergeModulelist를 사용하는 MoE 모델의 경우, 이론상 최악의 경우 메모리 피크는 모델 크기 더하기 한 레이어의 전문가예요.

이런 최악의 시나리오는 흔하지 않아요. 실제 메모리 피크는 모델 크기에 가까이 머무는 경향이 있어요.

가역성 (Reversibility)

시스템은 역변환으로 모델 저장을 지원하므로 왕복 저장/로드가 가능해요. 저장은 두 단계로 실행돼요: 역전된 컨버터 먼저(각 텐서는 기껏해야 하나와 매칭), 그다음 역전된 이름 바꾸기 (위 순서 규칙에 따라).

def revert_weight_conversion(model, state_dict):
    """Applies reverse conversions for saving."""
    weight_conversions = getattr(model, "_weight_conversions", None)

    # Reverse all transforms
    reverse_weight_conversion = [
        conversion.reverse_transform() for conversion in weight_conversions
    ]

    # Apply in reverse
    for first_param_name, reversed_converter in conversion_mapping.items():
        realized_value = reversed_converter.convert(first_param_name, model=model)

로드 시 사용된 변환 목록은 _weight_conversions로 모델에 캐시돼요 (실제로 발생한 항목만 유지되므로 비전단사 이름 바꾸기, 예: PrefixChange, 가 대칭적으로 올바르게 다시 적용돼요). 모델이 from_pretrained 없이 인스턴스화된 경우(따라서 _weight_conversions가 없음) revert_weight_conversion은 get_model_conversion_mapping으로 매핑을 다시 계산하고, 거기서 PrefixChange를 빼요 (원래 체크포인트에 접두사가 있었는지 알 수 없기 때문).

타겟 패턴은 역방향 처리(processing)가 필요한 정규식 요소를 포함할 수 있어요:

def process_target_pattern(pattern: str) -> tuple[str, str | None]:
    """
    - Removes `^` and `$` anchors
    - Removes negative lookahead/lookbehind
    - Detects capturing groups, replaces with \1
    """

실제 예시 (Real examples)

Mixtral 스타일 MoE

체크포인트 형식:

model.layers.0.block_sparse_moe.experts.0.w1.weight  # gate per expert
model.layers.0.block_sparse_moe.experts.0.w2.weight  # down per expert
model.layers.0.block_sparse_moe.experts.0.w3.weight  # up per expert
...
model.layers.0.block_sparse_moe.experts.7.w1.weight

모델 형식:

model.layers.0.mlp.experts.gate_up_proj  # (8, 4096, 28672)
model.layers.0.mlp.experts.down_proj     # (8, 14336, 4096)

변환 매핑 (conversion_mapping.py에서):

"mixtral": [
    WeightRenaming(".block_sparse_moe.", ".mlp."),
    WeightConverter(
        source_patterns=[".experts.*.w1.weight", ".experts.*.w3.weight"],
        target_patterns=".experts.gate_up_proj",
        operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
    ),
    WeightConverter(
        source_patterns=[".experts.*.w2.weight"],
        target_patterns=".experts.down_proj",
        operations=[MergeModulelist(dim=0)],
    ),
],

복합 vision-language 모델 (Gemma3)

Gemma3는 "기본 모델과 다른 헤드 모델에 대한 다른 접두사 로직"의 전형적인 설정을 가져요. 두 항목이 등록돼요 (여기서는 공유된 "llava" / "LlavaModel" 목록을 가리키는 별칭을 통해):

# model_type: applies to Gemma3ForConditionalGeneration / Gemma3ForSequenceClassification
"gemma3": "llava"           # → adds "model." in front of language_model / vision_tower / ...

# class:    applies to the inner Gemma3Model only
"Gemma3Model": "LlavaModel"  # → minimal rename inside the already-prefixed namespace

Gemma3ForConditionalGeneration의 DFS 순회:

  1. 루트 — 클래스 조회는 실패(Gemma3ForConditionalGeneration은 등록되지 않음), model_type 조회는 성공("gemma3") → "llava" 접두사 재작성 변환이 스코프 없이 추가되어 language_model, vision_tower, multi_modal_projector를 model.* 아래에 둠.
  2. 내부 model: Gemma3Model — 클래스 조회 성공(Gemma3Model → LlavaModel) → LlavaModel 매핑만 "model"으로 스코프되어 적용됨. 훨씬 더 넓은 "gemma3"(="llava") 접두사 이름 바꾸기는 여기서 다시 적용되지 않는데, 이것이 바로 원하는 동작이에요: 내부 모델은 접두사가 이미 올바른 네임스페이스에 있기 때문이에요.

같은 쌍 형태는 모든 Llava 계열 VLM(PaliGemma, InternVL, Mistral3, ...)에서 사용돼요. model_type 항목은 헤드 모델의 접두사 수술을 처리하고, 클래스 항목은 내부 기본 모델의 매핑을 최소로 유지해요.

더 깊은 중첩 + 헤드별 재정의 (DETR)

DetrForSegmentation은 두 개의 중첩된 레벨 위에 헤드별 재정의가 쌓인 사례를 보여줘요:

DetrForSegmentation                  (class registered: segmentation-only renames)
├── detr: DetrForObjectDetection     (no mapping; just walked through)
│   └── model: DetrModel             (class registered: shared base transforms)
│       └── backbone, encoder, ...
└── mask_head, bbox_attention        (head-specific weights)

등록된 키당 하나씩 두 개의 매핑이 관련돼요:

"DetrModel":           [WeightRenaming("backbone.conv_encoder", "backbone"), ...]      # shared base
"DetrForSegmentation": [WeightRenaming("mask_head.lay1", "mask_head.conv1.conv"), ...] # head-specific

DFS 순회:

  1. 루트 DetrForSegmentation이 클래스로 매칭 → segmentation 이름 바꾸기가 스코프 없이 추가됨.
  2. detr: DetrForObjectDetection은 등록되지 않음 → 변환 없음; DFS가 그 안으로 계속됨.
  3. detr.model: DetrModel이 클래스로 매칭 → 기본 변환이 scope_prefix="detr.model"로 추가됨.

이것이 DetrForObjectDetection에 자체 매핑이 필요하지 않은 이유예요: 그 하위 트리의 유일한 등록 매핑(DetrModel)이 올바른 경로로 자동 스코프되기 때문이에요.

클래스 키 별칭은 추가 등록 없이 기본 매핑을 재사용해요: "MaskFormerDetrDecoder": "DetrModel"은 MaskFormer 디코더가 자체 클래스 이름 아래에서 같은 변환을 받도록 해요.

커스텀 연산 (ERNIE 4.5 VL MoE)

내장 연산으로 충분하지 않을 때는 커스텀 ConversionOps 서브클래스를 만들 수 있어요. 예를 들어 ERNIE 4.5 VL MoE는 공유 전문가 목록을 텍스트와 비전 모달리티 사이에서 분할해야 하는데, 어떤 내장 연산도 이를 처리하지 못해요. 커스텀 ErnieFuseAndSplitTextVisionExperts 연산은 두 개의 타겟 키에 걸쳐 전문가를 분할하고 다시 쌓아요:

"ernie4_5_vl_moe": [
    WeightRenaming("vision_model", "vision_tower"),
    WeightConverter(
        source_patterns=["experts.*.down_proj.weight"],
        target_patterns=[
            "text_moe.experts.down_proj",
            "vision_moe.experts.down_proj",
        ],
        operations=[ErnieFuseAndSplitTextVisionExperts(stack_dim=0, concat_dim=1)],
    ),
],

커스텀 연산은 왕복 저장/로드를 지원하려면 convert()와 reverse_op 속성을 구현해야 해요.

모델 타입 별칭

많은 모델이 변환 패턴을 공유해요:

_MODEL_TO_CONVERSION_PATTERN = {
    "mixtral": "mixtral",
    "minimax": "mixtral",
    "qwen2_moe": "qwen2_moe",
    "deepseek_v2": "qwen2_moe",
    "deepseek_v3": "qwen2_moe",
    "qwen3_moe": "qwen2_moe",
    "olmoe": "qwen2_moe",
    ...
}

동적 로딩 빌딩 블록 재사용

동적 가중치 로딩은 전체 모델 체크포인트에만 국한되지 않아요. 같은 빌딩 블록으로, 체크포인트 키가 파라미터에 어떻게 매핑되는지 기술하고 타겟 모듈이 존재하도록만 보장하면 어떤 가중치 집합이든 로드할 수 있어요.

전체적으로 계약은 다음과 같아요:

  1. 모델 네임스페이스 준비. 로드하려는 모듈/파라미터가 존재하고 매핑이 타겟으로 삼는 방식으로 이름 지어졌는지 확인해요. 어댑터의 경우 로드 전에 어댑터 모듈이 존재하도록 inject_adapter_in_model(...)을 호출하는 것을 의미해요. 커스텀 헤드나 추가 모듈의 경우 먼저 모델에 인스턴스화해요.
  2. 가중치 매핑 방법 기술. WeightConverter나 WeightRenaming을 사용해 (예를 들어 _build_peft_weight_mapping(...) 같은 헬퍼에서) 변환/이름 바꾸기 목록을 만들어요. 여기서 체크포인트 키를 모델 네임스페이스에 맞게 어떻게 변환, 분할, 병합, 이름 바꿀지 표현해요. 주로 세 가지를 할 수 있어요:
    • 컨버터 목록에 연산 추가: 이들은 어떤 WeightConverter에도 수집되지 않은 모든 가중치에 적용돼요. 일반적으로 WeightRenaming 연산이어야 해요.
    • 각 컨버터의 연산 목록에 연산 추가: Quantization에서 그런데, 어떤 WeightConverter의 연산 목록 뒤에 양자화 연산을 추가하기만 하면 돼요.
    • 연산을 커스텀 연산으로 교체/매핑: peft에서 그런데, 예를 들어 mixtral의 Concatenate 연산을 PeftConcatenate(PEFT에 정의됨)로 교체해요. 이렇게 하면 어댑터 체크포인트를 읽을 때 연결될 가중치들이 수집되어 peft에 맞게 올바르게 형식화돼요.
  3. 로드 + 마무리 + 보고. 코어 로더를 사용해 변환을 수행하고 텐서를 채운 다음 마무리하고 결과를 기록해요. 구체적으로 이 흐름은:
    • LoadStateDictConfig(...) + _load_pretrained_model(...)으로 로드하고 변환.
    • _finalize_load_state_dict(...)으로 meta에서 누락/불일치 텐서를 옮기고 초기화하며 가중치를 바인딩.
    • log_state_dict_report(...)으로 누락/예상 밖/불일치 키(및 변환 오류)를 보고.

이 API들은 커스텀 코드, 커스텀 가중치 형식을 다룰 수 있게 하면서도 transformers API의 가장 높고 효율적인 가중치 로딩, 샤딩, 좋은 사용성을 누릴 수 있도록 공개돼요!

주요 파일 참조

File Purpose
src/transformers/core_model_loading.py Core loading logic, WeightConverter, WeightRenaming, ConversionOps
src/transformers/conversion_mapping.py Built-in mappings and per-submodel composition (get_model_conversion_mapping)
src/transformers/integrations/tensor_parallel.py TP sharding classes and utilities
src/transformers/quantizers/base.py Quantization hooks and base class