Transformers documentation

Carregamento dinâmico de pesos

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Carregamento dinâmico de pesos

Os checkpoints são frequentemente serializados em um formato que não corresponde ao que o modelo espera em tempo de execução. Cenários comuns incluem:

  1. Pesos fundidos: Os checkpoints armazenam pesos separados de gate_proj e up_proj, mas o modelo utiliza um gate_up_proj fundido para eficiência.
  2. Consolidação de especialistas MoE: Os pesos individuais dos especialistas (experts.0.weight, experts.1.weight, …) precisam ser empilhados em um único tensor 3D.
  3. Nomeação legada: Checkpoints antigos usam convenções de nomenclatura diferentes (por exemplo, LayerNorm.gamma vs LayerNorm.weight).
  4. Modelos compostos: Um modelo de visão-linguagem contém dois submódulos PreTrainedModel, cada um com sua própria convenção de checkpoint.
  5. Quantização: Os pesos podem ser armazenados em formatos quantizados que exigem desserialização.

O carregamento dinâmico de pesos aborda isso aplicando operações agendadas e reversíveis aos tensores do checkpoint à medida que são carregados. O Transformers expõe isso por meio de WeightConverter e WeightRenaming, que descrevem como uma ou mais chaves de checkpoint mapeiam para um ou mais parâmetros do modelo e quais ConversionOps componíveis devem ser executados nos tensores correspondentes. Essa abordagem se adapta a novos layouts de pesos, suporta misturas de especialistas quantizadas (MoEs) e se integra à paralelização de tensores.

Este guia demonstra como usar WeightConverter para converter tensores. Os mapeamentos de conversão estão em conversion_mapping.py; um mapeamento registrado é indexado por uma string model_type (por exemplo, "mixtral") ou por um nome de classe (por exemplo, "LlavaModel").

Pipeline completo de carregamento

Todos os modelos passam pelo sistema de carregamento dinâmico de pesos. O mapeamento de conversão é uma etapa opcional dentro desse sistema que só é ativada quando o modelo possui entradas registradas para o nome da sua classe ou para o 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                                │
                         └───────────────────────────────────────────────────────────┘
EtapaQuando ele é ativado
Carregamento dinâmicoSempre, para todos os modelos
Mapeamento de conversãoSomente quando a classe do modelo ou model_type está registrada em _MODEL_TO_CONVERSION_PATTERN
Sharding de TPSomente quando DistributedConfig(tp_size=N) é usado e o modelo possui base_model_tp_plan
Desquantização/desserializaçãoSomente ao carregar um checkpoint pré-quantizado
QuantizaçãoSomente quando uma configuração de quantização é fornecida e os pesos não estão pré-quantizados

Modelos densos (por exemplo, Llama)

Para a maioria dos modelos densos, o formato do checkpoint corresponde diretamente ao formato do modelo, portanto, nenhum mapeamento de conversão é necessário. Alguns modelos ainda podem exigir renomeação (por exemplo, convenções de nomenclatura legadas). O sharding de TP ainda se aplica quando habilitado.

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

Os checkpoints legados podem usar convenções de nomenclatura mais antigas, que são tratadas por renomeações embutidas aplicadas a todos os modelos:

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

Modelos MoE (por exemplo, Mixtral)

Para modelos MoE, o formato do checkpoint difere do formato do modelo. O mapeamento de conversão transforma pesos de especialistas separados em tensores 3D fundidos, e o sharding de TP é aplicado após a conversão.

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  ─┘

Modelos compostos (por exemplo, visão-linguagem)

Uma PreTrainedModel pode conter outros submódulos PreTrainedModel. Cada submodelo pode ter seu próprio mapeamento de conversão, registrado em relação ao nome da sua classe ou ao seu model_type. Quando o modelo pai é carregado, get_model_conversion_mapping percorre os submodelos em ordem pré-ordem (depth-first), coleta seus mapeamentos e escopa automaticamente cada transformação para o caminho de seu submódulo por meio de 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 é o caminho pontilhado do submódulo ("vision_model", "language_model.model", etc.). Uma transformação com escopo só é acionada para chaves que começam com f"{scope_prefix}."; o prefixo é removido antes da correspondência de padrões e reanexado após a substituição, de modo que o mapeamento de cada submodelo é escrito relativo ao submodelo, exatamente como se ele fosse a raiz.

Arquitetura

O sistema é construído em torno de vários componentes principais definidos em core_model_loading.py:

Fase 1 — Processamento por chave (iteração sobre as chaves do checkpoint):

  1. Percorra a lista de transformações uma única vez. Cada WeightRenaming que corresponder será acionado (por exemplo, block_sparse_moe → mlp), e no máximo um WeightConverter pode reivindicar a chave (por exemplo, experts.*.w1.weight).
  2. Shard (TP) e enviar para o dispositivo de forma assíncrona via ThreadPoolExecutor.
  3. Coletar os tensors com o mesmo source_pattern juntos (por exemplo, todos os pesos de especialistas MoE, projeções gate + up).

Fase 2 — Processamento por mapeamento (itera sobre os mapeamentos coletados):

  1. Desquantizar/desserializar (somente checkpoints pré-quantizados).
  2. Aplicar a cadeia ConversionOps: Chunk, Concatenate, MergeModulelist, Transpose, etc.
  3. Quantize em tempo de execução (se não estiver pré-quantizado).
  4. Defina o parâmetro Set parameter no modelo.

WeightTransform

A classe base que trata do correspondência de padrões e da coleta de tensores:

  • Compilação de padrões: Os padrões de origem são expressões regulares completas, correspondidas com re.search(). O curinga * corresponde a qualquer componente indexável e agrupa todas as correspondências para operações em lote.
  • Grupos de captura e referências cruzadas: Os grupos de captura nos padrões de origem podem ser referenciados como \1, \2, … nos padrões de destino para preservar substrings (por exemplo, índices de camadas) durante a renomeação.
  • Escopo: scope_prefix (definido automaticamente para cada submodelo por get_model_conversion_mapping) restringe a transformação às chaves sob esse caminho. O prefixo é removido antes da correspondência e reanexado após a substituição.
  • Renomeação de chaves: rename_source_key() aplica a expressão regular (com tratamento de escopo) e retorna (renamed_key, source_pattern) para que o loader saiba qual conversor, se houver, reivindicou a chave.
  • Coleção de tensores: add_tensor() acumula tensores resolvidos (ou Futures) sob seu source_pattern para que todos os tensores necessários para uma única conversão (por exemplo, todos os pesos de especialistas MoE) sejam agrupados antes da execução da cadeia de operações.
  • Reversibilidade: reverse_transform() troca os padrões de origem ↔ destino e inverte cada operação, de modo que a mesma lista invertida dirige o salvamento.
  • Rastreamento: was_used() relata se a transformação correspondeu a alguma chave durante o carregamento; isso é necessário para que renomeações não-bijeções (por exemplo, PrefixChange adicionando um prefixo que pode já estar presente) possam ser reaplicadas simetricamente na hora de salvar.

WeightRenaming

WeightRenaming é um WeightTransform especializado em renomeações puras de chaves, sem operações em tensores. Diferentemente de WeightConverter, um WeightRenaming não claim a chave, permitindo que várias renomeações sejam encadeadas livremente. No caminho de carregamento, todas as renomeações são executadas antes do conversor; no caminho de salvamento (lista invertida), todas as renomeações invertidas são executadas após o conversor invertido — consulte a regra de ordenação acima.

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

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

PrefixChange é um wrapper de nível superior para WeightRenaming, voltado ao caso comum de remover ou adicionar um componente inteiro de caminho. O parâmetro opcional model_prefix limita a operação às chaves sob esse namespace:

# "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 estende WeightTransform com uma cadeia de ConversionOps que atuam nos tensores coletados. As quatro cardinalidades suportadas são:

CardinalidadePadrões de origemPadrões de destinoOperação típica
um para um11Transpose, PermuteForRope
um-para-vários1>1Chunk (por exemplo, descompactar qkv_proj)
muitos-para-um>11Concatenate, MergeModulelist (por exemplo, fundir especialistas)
muitos-para-muitos>1>1somente com operações que explicitamente o suportam (por exemplo, ErnieFuseAndSplitTextVisionExperts)
WeightConverter(
    source_patterns=[".experts.*.w1.weight", ".experts.*.w3.weight"],
    target_patterns=".experts.gate_up_proj",
    operations=[MergeModulelist(dim=0), Concatenate(dim=1)],
)

Um WeightConverter também é reversível: reverse_transform() troca origem ↔ destino e substitui cada operação pelo seu reverse_op, de modo que um mapeamento de conversão registrado é bidirecional por construção (o carregamento usa a lista como está, e o salvamento a usa invertida).

Ordem de renomeações e conversores

Liste as entradas de WeightRenaming antes das entradas de WeightConverter e mantenha suas folhas disjuntas: os renames normalizam as chaves que os conversores consomem, mas nunca devem visar uma folha produzida por um conversor. O caminho de salvamento depende disso para inverter o mapeamento em duas fases (reverter conversores e, em seguida, reverter renames).

A lista de transformações é percorrida em ordem, uma vez por cada chave de checkpoint. Para cada chave:

  • Cada WeightRenaming que corresponde é acionado; várias renomeações podem ser encadeadas.
  • O primeiro WeightConverter que corresponde às claims da chave. Os conversores subsequentes são ignorados — isso garante que o tensor seja roteado para um único conversor na etapa de mesclagem/divisão, e é a razão pela qual dois conversores com intenção sobreposta configurariam um erro de configuração.
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"

Operações de conversão

A classe WeightConverter possui várias operações que são executadas quando from_pretrained() é chamada para transformar os tensores de origem do checkpoint nos tensores de destino do modelo.

As operações são totalmente reversíveis. A operação de salvamento reverte as conversões e retorna o checkpoint original, permitindo que você trabalhe facilmente entre diferentes frameworks.

OperaçãoReverter
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

A operação Chunk divide um tensor em partes iguais ao longo de uma dimensão. Por exemplo, quando um modelo espera Q, K e V como três tensores separados, em vez de um único tensor.

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

Concatenar

A operação Concatenate funde tensores separados em um único tensor. Por exemplo, quando um modelo espera Q, K e V como um único tensor, em vez de tensores separados.

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

MergeModulelist

MergeModulelist mescla uma lista de tensores 2D em um único tensor 3D. Por exemplo, você pode combinar MergeModulelist com Concatenate para empilhar os especialistas em um MoE e empacotá-los em um único tensor.

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 divide um tensor 3D de volta em uma lista de tensores 2D. Por exemplo, você pode dividir uma pilha de especialistas de volta em especialistas individuais.

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

PermuteForRope

PermuteForRope converte pesos do formato intercalado para usar o formato sin/cos. Por exemplo, você pode compor Chunk com PermuteForRope para dividir um tensor QKV fundido e aplicar a permutação RoPE sin/cos a Q e K.

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()],
)

Transpor

Transpose troca as dimensões de um tensor. Útil para converter layouts de pesos entre diferentes convenções.

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

Encadeamento de operações

As operações podem ser encadeadas para realizar transformações complexas. As operações são executadas em ordem, com a saída de cada operação se tornando a entrada da operação seguinte.

Exemplo: conversão do 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)
    ],
)

Fluxo de dados:

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

Regras de correspondência de padrões

  • Os padrões de origem são expressões regulares completas avaliadas com re.search(). ^ ancora ao início da chave com escopo (após a remoção de scope_prefix).
  • * é um curinga por índice que coleta todos os tensors correspondentes sob o mesmo padrão de origem (preservando a ordem do checkpoint para a concatenação correta).
  • Grupos capturados em padrões de origem podem ser referenciados como \1, \2, … em padrões de destino para manter partes da chave original (por exemplo, índices de camadas).
  • Transformações com escopo (scope_prefix definido) correspondem apenas a chaves que começam com f"{scope_prefix}.".

Registrando um mapeamento de conversão

Uma lista de conversão é registrada em relação a uma chave de string — seja um model_type (por exemplo, "mixtral") ou um nome de classe (por exemplo, "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)],
        ),
    ],
)

Regras de consulta

Quando get_model_conversion_mapping processa um PreTrainedModel, cada sub-PreTrainedModel é visitado em ordem DFS (nn.Module.named_modules() filtrado para instâncias de PreTrainedModel). Cada submodelo é resolvido em duas etapas: primeiro, um filtro de código personalizado é executado; em seguida, o mapeamento registrado é consultado.

Modelos de código personalizado são ignorados antes de qualquer consulta. Um submodelo é considerado de código personalizado quando is_custom_code() retorna True, o que abrange modelos carregados do Hub com trust_remote_code=True e qualquer classe de modelo que não esteja no Transformers. Modelos de código personalizado não herdam um mapeamento de conversão, pois o nome da classe ou o model_type podem conflitar com um modelo nativo. Por exemplo, uma arquitetura personalizada que define model_type="mixtral" não aplicará as transformações de fusão de especialistas embutidas do Mixtral.

Registre explicitamente um modelo de código personalizado com register_checkpoint_conversion_mapping para aplicar um mapeamento de conversão. O registro adiciona o nome da classe ou model_type ao conjunto de mapeamentos registrados pelo usuário, o que isenta o modelo do skip de código personalizado e o encaminha pela busca normal descrita abaixo.

Para cada submodelo que passa no filtro de código personalizado:

  1. A busca pelo nome da classe é tentada primeiro, em seguida o model_type. Se ambos estiverem registrados para o mesmo módulo, o mapeamento pelo nome da classe tem prioridade e o model_type é ignorado para aquele módulo — isso permite que uma cabeça de tarefa (por exemplo, LlavaForConditionalGeneration) sobrescreva a linha de base compartilhada de model_type ("llava").
  2. O mapeamento selecionado tem scope_prefix definido como o caminho pontilhado do submódulo ("" para a raiz).
  3. Ancestor-based deduplication determina se o mapeamento deve ser mantido:
    • Se um caminho ancestral já reivindicou o mesmo identificador (nome da classe ou model_type), o submódulo é ignorado — o mapeamento sem escopo ou com escopo mais amplo do ancestral já cobre essa subárvore.
    • Se apenas um irmão o reivindicou, o submódulo é mantido com seu próprio scope_prefix. Cada irmão recebe seu próprio mapeamento com escopo.

As listas de verificação de nomes de classe e de model_type são rastreadas separadamente, com uma sutileza: quando um módulo é correspondido por seu nome de classe, seu model_type não é adicionado à lista de verificação. Isso permite que outros módulos que compartilhem o mesmo model_type, mas que não tenham um mapeamento específico de classe (por exemplo, DetrModel sob DetrForSegmentation), permaneçam acessíveis por meio da consulta de model_type.

Mapeamentos baseados em classes vs. aliases de model_type

Ambos os estilos coexistem em _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",
    ...
}

Chaves de nome de classe são preferidas quando o model_type é compartilhado, mas uma classe específica precisa de um comportamento diferente.

Integração com paralelismo de tensor

O sistema de carregamento dinâmico se integra à paralelização de tensores (TP) por meio da hierarquia TensorParallelLayer definida em src/transformers/integrations/tensor_parallel.py.

Quando o TP está habilitado, os tensores são fragmentados durante a materialização, e não depois. Isso significa que cada rank carrega apenas a parte do tensor de que precisa.

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)

Estilos de paralelismo disponíveis

EstiloDimensão do Fragmento de PesoDescrição
colwise-2Por coluna: recursos de saída fragmentados
rowwise-1Linha a linha: recursos de entrada fragmentados
packed_colwise-2Para pesos fundidos (gate_up_proj)
packed_rowwise-1Para pesos fundidos
embedding_rowwise0Paralelismo de vocabulário
grouped_gemm0Paralelismo de especialistas para MoE
sequence_parallelNoneSem sharding de pesos

Tratamento de pesos empacotados

Para pesos fundidos, como gate_up_proj, é necessário tomar cuidado especial para fazer o sharding corretamente:

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]
    """

A operação de TP é armazenada no WeightTransform e aplicada após as operações de conversão:

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()
    )

Integração de quantização

A quantização se integra ao pipeline de carregamento de duas maneiras, dependendo de o checkpoint já estar quantizado:

  • Checkpoints pré-quantizados: O quantizador fornece instâncias de WeightConverter (via get_weight_conversions()) que desserializam tensores quantizados. Os dtypes dos checkpoints são preservados para evitar casts indesejados.
  • Quantização em tempo real: O quantizador fornece uma operação de quantização que é aplicada após as operações de conversão, quantizando os pesos conforme são carregados.

O quantizador também pode reescrever a lista de conversão inteira no final de get_model_conversion_mapping por meio de update_weight_conversions(...) — por exemplo, o dequantizador FP8 insere uma op Fp8Dequantize no início de cada conversor existente para que as escalas por bloco sejam aplicadas antes de qualquer op de mesclagem de especialistas / concatenação aplanar a estrutura por especialista.

Carregamento de modelos rápido e eficiente

Carregar um modelo é mais rápido e usa menos memória, pois o loader sabe quais tensores são necessários para as operações e agenda sua materialização de forma preguiçosa.

O loader varre o checkpoint uma única vez para descobrir correspondências de padrões e coletar tensores. Ele os armazena como objetos Future e os envia para uma thread pool para carregamento assíncrono sem bloquear o GIL. Um parâmetro inicia o carregamento assim que uma thread fica disponível para ele.

Se o seu sistema executa outros processos pesados, múltiplas threads podem desacelerar o carregamento em vez de acelerá-lo. Nesse caso, defina a variável de ambiente HF_DEACTIVATE_ASYNC_LOAD=1 para carregar os pesos sequencialmente.

O padrão é 4 threads para o carregamento assíncrono de parâmetros. Isso oferece o melhor equilíbrio entre os cenários de carregamento e o hardware. O trabalho é majoritariamente limitado por I/O, mas, dependendo do hardware do acelerador e do dtype necessário no carregamento, pode se tornar limitado por CPU/GPU se o dtype for diferente do serializado (isso requer uma operação de cópia adicional).

Carregamento assíncrono vs síncrono

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)

O carregamento síncrono é usado quando:

  • a variável de ambiente HF_DEACTIVATE_ASYNC_LOAD=1 está definida.
  • O offloading para o disco está habilitado (as restrições de memória exigem carregamento sequencial).
  • A quantização em tempo de execução está habilitada (evita que as threads de trabalho fiquem à frente da etapa de quantização).

Fluxo de materialização

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)

Eficiência de memória

Ao converter um peso, o conversor aguarda que todos os tensores necessários sejam materializados, caso ainda não tenham sido carregados. Por exemplo, a operação MergeModulelist exige que todos os pesos em ModuleList estejam carregados antes da fusão.

A concatenação de tensores requer uma cópia temporária, portanto, operações como MergeModulelist e Concatenate precisam de 2x a memória dos tensores subjacentes durante a conversão. Após a fusão, apenas o tensor resultante permanece na memória. O pico de memória no pior caso teórico é o tamanho do modelo mais os tensores necessários para a maior operação de MergeModulelist ou Concatenate.

Esse pior caso só ocorre quando todos os outros parâmetros já foram carregados antes da conversão exigente ser executada. Dois cenários disparam isso.

  1. Todos os parâmetros foram carregados de forma assíncrona antes de entrar na conversão exigente (a pool de threads foi mais rápida que a fila de conversão).
  2. A conversão mais exigente é a última.

Por exemplo, em um modelo MoE que usa MergeModulelist para os especialistas em cada camada, o pico de memória teoricamente mais alto é o tamanho do modelo mais os especialistas de uma camada.

Esses cenários de pior caso são incomuns. O pico de memória real tende a permanecer próximo ao tamanho do modelo.

Reversibilidade

O sistema permite salvar modelos com as transformações inversas, habilitando o salvamento e carregamento de ida e volta. O salvamento ocorre em duas fases: primeiro os conversores invertidos (cada tensor corresponde a no máximo um), e depois os renomeamentos invertidos, conforme a regra de ordenação acima.

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)

A lista de transforms usados no momento do carregamento é armazenada em cache no modelo como _weight_conversions (somente as entradas que foram efetivamente aplicadas são mantidas, de modo que renomeações não-bijeções, como PrefixChange, são reaplicadas corretamente de forma simétrica). Quando o modelo foi instanciado sem from_pretrained (e, portanto, não possui _weight_conversions), revert_weight_conversion recorre à recomputação do mapeamento via get_model_conversion_mapping e remove qualquer PrefixChange dele (não é possível determinar se o checkpoint original continha o prefixo).

Os padrões de destino podem conter elementos de regex que precisam de processamento para a direção reversa:

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

Exemplos reais

MoE no estilo Mixtral

Formato do checkpoint:

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

Formato do modelo:

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

Mapeamento de conversão (de 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)],
    ),
],

Modelo composto de visão e linguagem (Gemma3)

O Gemma3 possui a configuração canônica de “lógica de prefixo diferente para o modelo de head do que para o modelo base”. Duas entradas são registradas (aqui, por meio de aliases que apontam para as listas compartilhadas "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

Percorrido DFS para Gemma3ForConditionalGeneration:

  1. Root — falha na busca de classe (Gemma3ForConditionalGeneration não está registrado), acerto na busca de model_type ("gemma3") → as transformações de reescrita do prefixo "llava" são adicionadas sem escopo, colocando language_model, vision_tower e multi_modal_projector sob model.*.
  2. Modelo interno model: Gemma3Model — consultas de classe correspondentes (Gemma3Model → LlavaModel) → apenas o mapeamento de LlavaModel é aplicado, limitado a "model". As renomeações de prefixo muito mais amplas de "gemma3" (="llava") não são reaplicadas aqui, o que é exatamente o que se deseja: o modelo interno está em um namespace onde o prefixo já está correto.

O mesmo formato de par é utilizado por todos os VLMs da família Llava (PaliGemma, InternVL, Mistral3, …). A entrada de model_type trata a cirurgia de prefixo do modelo de head; a entrada de classe mantém o mapeamento do modelo base interno minimalista.

Aninhamento mais profundo com uma substituição específica para a cabeça (DETR)

DetrForSegmentation exibe uma sobrescrita específica de head sobreposta a dois níveis aninhados:

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)

Dois mapeamentos estão envolvidos (um por chave registrada):

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

Percorrido em DFS:

  1. A raiz DetrForSegmentation corresponde por classe → renomeações de segmentação adicionadas unscoped.
  2. detr: DetrForObjectDetection não está registrado → nenhuma transformação foi adicionada; a DFS continua para dentro dele.
  3. detr.model: DetrModel corresponde por classe → transformações base adicionadas com scope_prefix="detr.model".

É por isso que DetrForObjectDetection não precisa de um mapeamento próprio: o único mapeamento registrado em sua subárvore (DetrModel) é automaticamente limitado ao caminho correto.

Os aliases com chave de classe reutilizam o mapeamento base sem registro adicional: "MaskFormerDetrDecoder": "DetrModel" faz com que um decoder MaskFormer utilize as mesmas transformações sob o nome da própria classe.

Operações personalizadas (ERNIE 4.5 VL MoE)

Quando as operações embutidas não são suficientes, você pode criar um subclasse personalizado de ConversionOps. Por exemplo, o ERNIE 4.5 VL MoE precisa dividir uma lista de especialistas compartilhados entre as modalidades de texto e visão — algo que nenhuma operação embutida individual trata. A operação personalizada ErnieFuseAndSplitTextVisionExperts divide e reempilha os especialistas em duas chaves de destino:

"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)],
    ),
],

Ops personalizadas devem implementar convert() e a propriedade reverse_op para suportar o salvamento e a carga de ida e volta.

Aliases de tipo de modelo

Muitos modelos compartilham padrões de conversão:

_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",
    ...
}

Reutilizando os blocos de construção do carregamento dinâmico

O carregamento dinâmico de pesos não se limita a checkpoints de modelos completos. Os mesmos blocos de construção permitem carregar qualquer conjunto de pesos, desde que você possa descrever como as chaves do checkpoint mapeiam para os parâmetros e garantir que os módulos de destino existam.

Em alto nível, o contrato é o seguinte:

  1. Prepare o namespace do modelo. Verifique se os módulos/parâmetros que você deseja carregar estão presentes e nomeados da forma como seu mapeamento os identificará. No caso de adaptadores, isso significa chamar inject_adapter_in_model(...) para que os módulos do adaptador existam antes do carregamento. Para cabeçalhos personalizados ou módulos extras, instancie-os no modelo primeiro.
  2. Describe how to map weights. Crie uma lista de conversão/renomeação (por exemplo, em um helper como _build_peft_weight_mapping(...)) usando WeightConverter ou WeightRenaming. Este é o local onde você expressa como as chaves do checkpoint devem ser convertidas, divididas, mescladas ou renomeadas para corresponder ao namespace do seu modelo. Você pode fazer basicamente 3 coisas:
    • adicione operações à lista de conversores: elas serão aplicadas a todos os pesos, exceto aos coletados em qualquer WeightConverter. Em geral, devem ser operações de WeightRenaming
    • adicione operações à lista de operações de cada conversor: é isso que acontece para Quantization, onde simplesmente adicionamos uma operação de quantização após a lista de operações de qualquer WeightConverter.
    • substituir as operações de / map pelas suas operações personalizadas: é isso que acontece com o peft. Substituímos a operação Concatenate de, por exemplo, o mixtral por PeftConcatenate (que é definida no PEFT). Dessa forma, quando o checkpoint do adaptador é lido, os pesos a serem concatenados são coletados e formatados adequadamente para o peft
  3. Load + finalize + report. Use o carregador principal para executar a conversão e preencher os tensores, em seguida, finalize e registre os resultados. Concretamente, esse fluxo é:
    • LoadStateDictConfig(...) + _load_pretrained_model(...) para carregar e converter.
    • _finalize_load_state_dict(...) para mover quaisquer tensores ausentes ou incompatíveis de meta, inicializá-los e vincular os pesos.
    • log_state_dict_report(...) para relatar chaves ausentes/inesperadas/desconformes (e erros de conversão).

Essas APIs são expostas para permitir que você gerencie código personalizado e formatos de pesos personalizados, mas também para garantir que você se beneficie do carregamento de pesos mais rápido e eficiente, do sharding e da boa qualidade de vida da API do transformers!

Referência dos arquivos principais

ArquivoPropósito
src/transformers/core_model_loading.pyLógica central de carregamento, WeightConverter, WeightRenaming, ConversionOps
src/transformers/conversion_mapping.pyMapeamentos integrados e composição por submodelo (get_model_conversion_mapping)
src/transformers/integrations/tensor_parallel.pyClasses e utilitários de sharding de TP
src/transformers/quantizers/base.pyHooks de quantização e classe base
Update on GitHub