Tải trọng số động

Checkpoint thường được serialize ở định dạng không khớp với những gì model kỳ vọng khi chạy. Các kịch bản phổ biến bao gồm:

1. Trọng số đã hợp nhất: Các checkpoint lưu trữ các trọng số gate_proj và up_proj riêng biệt, nhưng mô hình sử dụng gate_up_proj đã hợp nhất để đạt hiệu quả cao hơn.
2. Gộp chuyên gia MoE: Các trọng số chuyên gia riêng lẻ (experts.0.weight, experts.1.weight, …) cần được xếp chồng thành một tensor 3 chiều duy nhất.
3. Đặt tên cũ: Các checkpoint cũ sử dụng các quy ước đặt tên khác nhau (ví dụ: LayerNorm.gamma so với LayerNorm.weight).
4. Mô hình tổng hợp: Một mô hình thị giác-ngôn ngữ chứa hai mô-đun con PreTrainedModel, mỗi mô-đun có quy ước checkpoint riêng.
5. Quantization: Các trọng số có thể được lưu trữ ở định dạng lượng tử hóa, cần được khử chuỗi hóa.

Việc nạp trọng số động giải quyết vấn đề này bằng cách áp dụng các thao tác có lịch trình và có thể đảo ngược lên các tensor checkpoint khi chúng được nạp. Transformers cung cấp tính năng này thông qua WeightConverter và WeightRenaming, mô tả cách một hoặc nhiều khóa checkpoint ánh xạ đến một hoặc nhiều tham số mô hình và các ConversionOps có thể kết hợp nào nên chạy trên các tensor khớp. Cách tiếp cận này thích ứng với các bố cục trọng số mới, hỗ trợ mixture-of-experts (MoE) lượng tử hóa và tích hợp với tensor parallelism.

Hướng dẫn này minh họa cách sử dụng WeightConverter để chuyển đổi các tensor. Các ánh xạ chuyển đổi nằm trong conversion_mapping.py; một ánh xạ đã đăng ký được định danh bằng chuỗi model_type (ví dụ: "mixtral") hoặc tên lớp (ví dụ: "LlavaModel").

Đường ống tải đầy đủ

Tất cả các model đều đi qua hệ thống nạp trọng số động. Bản ánh xạ chuyển đổi là một bước tùy chọn trong hệ thống đó, chỉ được kích hoạt khi model có các mục được đăng ký cho tên lớp của nó hoặc 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                                │
                         └───────────────────────────────────────────────────────────┘

Các mô hình Dense (ví dụ: Llama)

Đối với hầu hết các mô hình dense, định dạng checkpoint khớp trực tiếp với định dạng mô hình, do đó không cần ánh xạ chuyển đổi. Một số mô hình vẫn có thể cần đổi tên (ví dụ: các quy ước đặt tên cũ). Việc chia TP (TP sharding) vẫn được áp dụng khi được bật.

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

Các checkpoint cũ có thể sử dụng các quy ước đặt tên cũ hơn, được xử lý bằng các quy tắc đổi tên tích hợp áp dụng cho tất cả các mô hình:

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

Các mô hình MoE (ví dụ: Mixtral)

Đối với các mô hình MoE, định dạng checkpoint khác với định dạng mô hình. Bản ánh xạ chuyển đổi biến các trọng số chuyên gia riêng lẻ thành các tensor 3D đã hợp nhất, và việc chia TP được áp dụng sau khi chuyển đổi.

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

Các mô hình tổng hợp (ví dụ: thị giác-ngôn ngữ)

Một PreTrainedModel có thể chứa các module con PreTrainedModel khác. Mỗi mô hình con có thể có ánh xạ chuyển đổi riêng, được đăng ký theo tên lớp hoặc model_type của nó. Khi mô hình cha được nạp, get_model_conversion_mapping sẽ duyệt qua các mô hình con theo thứ tự duyệt trước (depth-first), thu thập các ánh xạ của chúng và tự động giới hạn phạm vi mỗi phép biến đổi theo đường dẫn của module con tương ứng thông qua 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 là đường dẫn dạng chấm của module con ("vision_model", "language_model.model", v.v.). Một transform có phạm vi chỉ kích hoạt với các key bắt đầu bằng f"{scope_prefix}."; phần tiền tố sẽ được loại bỏ trước khi khớp mẫu và được gắn lại sau khi thay thế, do đó ánh xạ của mỗi sub-model được viết tương đối so với sub-model, y hệt như thể nó là gốc.

Kiến trúc

Hệ thống được xây dựng xoay quanh một số thành phần chính được định nghĩa trong core_model_loading.py:

Giai đoạn 1 — Xử lý theo từng khóa (lặp qua các khóa checkpoint):

1. Quét danh sách transform một lần. Mọi WeightRenaming khớp sẽ được kích hoạt (ví dụ: block_sparse_moe → mlp), và tối đa một WeightConverter có thể nhận key (ví dụ: experts.*.w1.weight).
2. Chia shard (TP) và gửi đến thiết bị bất đồng bộ thông qua ThreadPoolExecutor.
3. Thu thập các tensor có cùng source_pattern lại với nhau (ví dụ: tất cả trọng số chuyên gia MoE, các phép chiếu gate + up).

Giai đoạn 2 — Xử lý theo từng ánh xạ (lặp qua các ánh xạ đã thu thập):

1. Khử lượng tử hóa/giải chuỗi (chỉ dành cho checkpoint đã lượng tử hóa trước).
2. Áp dụng chuỗi ConversionOps: Chunk, Concatenate, MergeModulelist, Transpose, v.v.
3. Quantize trực tiếp (nếu chưa được lượng tử hóa trước).
4. Đặt tham số trên mô hình.

WeightTransform

Lớp cơ sở xử lý khớp mẫu và thu thập tensor:

• Biên dịch mẫu: Các mẫu nguồn là các biểu thức chính quy đầy đủ được khớp bằng re.search(). Ký tự đại diện * khớp với bất kỳ thành phần nào có thể lập chỉ mục và nhóm tất cả các kết quả khớp lại với nhau cho các thao tác theo lô.
• Nhóm bắt & tham chiếu ngược: Các nhóm bắt trong mẫu nguồn có thể được tham chiếu dưới dạng \1, \2, … trong mẫu đích để giữ nguyên các chuỗi con (ví dụ: chỉ số lớp) qua quá trình đổi tên.
• Phạm vi: scope_prefix (được đặt tự động cho từng mô hình con bởi get_model_conversion_mapping) giới hạn phép biến đổi đối với các khóa nằm dưới đường dẫn đó. Tiền tố sẽ bị loại bỏ trước khi khớp và được gắn lại sau khi thay thế.
• Đổi tên khóa chính: rename_source_key() áp dụng biểu thức chính quy (có xử lý phạm vi) và trả về (renamed_key, source_pattern) để trình nạp biết trình chuyển đổi nào, nếu có, đã nhận khóa đó.
• Tập hợp tensor: add_tensor() tích lũy các tensor đã được phân giải (hoặc Future) dưới source_pattern của chúng để tất cả các tensor cần thiết cho một phép chuyển đổi duy nhất (ví dụ: tất cả trọng số chuyên gia MoE) được nhóm lại với nhau trước khi chuỗi thao tác được thực thi.
• Tính khả nghịch: reverse_transform() hoán đổi các mẫu nguồn ↔ đích và đảo ngược từng thao tác, do đó danh sách đảo ngược sẽ được dùng để lưu trữ.
• Theo dõi: was_used() báo cáo xem transform có thực sự khớp với bất kỳ key nào trong quá trình tải hay không; điều này là bắt buộc để các đổi tên không song ánh (ví dụ: PrefixChange thêm một tiền tố có thể đã tồn tại) có thể được áp dụng lại đối xứng khi lưu.

WeightRenaming

WeightRenaming là một WeightTransform chuyên dụng cho việc đổi tên khóa thuần túy mà không có thao tác tensor. Khác với WeightConverter, một WeightRenaming không claim khóa, do đó nhiều lần đổi tên có thể được nối chuỗi tự do. Trên đường dẫn load, tất cả các lần đổi tên được thực hiện trước converter; trên đường dẫn save (danh sách đảo ngược), tất cả các lần đổi tên đảo ngược được thực hiện sau converter đảo ngược — xem quy tắc thứ tự ở trên.

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

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

PrefixChange là một lớp bọc cấp cao hơn của WeightRenaming cho trường hợp phổ biến là xóa hoặc thêm một thành phần đường dẫn hoàn chỉnh. Tham số tùy chọn model_prefix giới hạn phạm vi thao tác đối với các khóa nằm dưới 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 kế thừa WeightTransform với một chuỗi các ConversionOps tác động lên các tensor đã thu thập. Bốn loại cardinality được hỗ trợ là:

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

WeightConverter cũng có thể đảo ngược: reverse_transform() hoán đổi nguồn ↔ đích và thay thế mỗi thao tác bằng reverse_op của nó, do đó một ánh xạ chuyển đổi đã đăng ký là hai chiều theo thiết kế (tải dùng danh sách như vậy, lưu dùng nó theo chiều ngược lại).

Xếp thứ tự các lệnh đổi tên và bộ chuyển đổi

Liệt kê các mục WeightRenaming trước các mục WeightConverter, và giữ cho các lá của chúng không giao nhau: việc đổi tên chuẩn hóa các khóa mà bộ chuyển đổi sử dụng, nhưng không bao giờ nhắm vào một lá do bộ chuyển đổi tạo ra. Lối lưu trữ dựa vào điều này để đảo ngược ánh xạ trong hai giai đoạn (đảo bộ chuyển đổi, sau đó đảo việc đổi tên).

Danh sách transform được duyệt lần lượt theo thứ tự, một lần cho mỗi khóa checkpoint. Với mỗi khóa:

• Mọi WeightRenaming khớp đều được kích hoạt; nhiều lần đổi tên có thể được nối tiếp nhau.
• claims của WeightConverter đầu tiên khớp với key. Các converter tiếp theo sẽ bị bỏ qua — điều này đảm bảo tensor được định tuyến đến một converter duy nhất cho bước merge/split, và cũng là lý do tại sao hai converter có ý định chồng chéo sẽ là một cấu hình sai.

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"

Các thao tác chuyển đổi

Lớp WeightConverter có một số thao tác được thực hiện khi from_pretrained() được gọi để chuyển đổi các tensor nguồn của checkpoint thành các tensor đích của mô hình.

Các thao tác hoàn toàn có thể đảo ngược. Việc lưu sẽ đảo ngược các phép chuyển đổi và trả về checkpoint gốc để bạn có thể dễ dàng làm việc xuyên suốt các framework khác nhau.

Chunk

Thao tác Chunk chia một tensor thành các phần bằng nhau dọc theo một chiều. Ví dụ, nếu một model kỳ vọng Q, K và V là ba tensor riêng biệt thay vì một tensor duy nhất.

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

Nối

Thao tác Concatenate hợp nhất các tensor riêng lẻ thành một tensor duy nhất. Ví dụ, khi một model kỳ vọng Q, K và V dưới dạng một tensor duy nhất thay vì các tensor riêng biệt.

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

MergeModulelist

MergeModulelist gộp một danh sách các tensor 2D thành một tensor 3D duy nhất. Ví dụ, bạn có thể kết hợp MergeModulelist với Concatenate để xếp chồng các chuyên gia trong một MoE và đóng gói chúng vào một tensor duy nhất.

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 tách một tensor 3D thành một danh sách các tensor 2D. Ví dụ, bạn có thể tách một chồng các chuyên gia (experts) thành các chuyên gia riêng lẻ.

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

PermuteForRope

PermuteForRope chuyển trọng số từ định dạng xen kẽ sang định dạng sin/cos. Ví dụ, bạn có thể kết hợp Chunk với PermuteForRope để tách một tensor QKV đã hợp nhất và áp dụng phép hoán vị RoPE sin/cos cho Q và 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()],
)

Chuyển vị

Transpose hoán đổi các chiều của một tensor. Hữu ích để chuyển đổi bố cục trọng số giữa các quy ước khác nhau.

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

Chuỗi thao tác

Các thao tác có thể được nối tiếp để thực hiện các phép biến đổi phức tạp. Các thao tác được thực hiện theo thứ tự, với đầu ra của mỗi thao tác trở thành đầu vào của thao tác tiếp theo.

Ví dụ: Chuyển đổi 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)
    ],
)

Lưu chuyển dữ liệu:

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

Quy tắc khớp mẫu

• Các mẫu nguồn là các biểu thức chính quy đầy đủ được đánh giá bằng re.search(). ^ neo vào đầu khóa scoped (sau khi loại bỏ scope_prefix).
• * là ký tự đại diện theo từng chỉ số, thu thập tất cả các tensor khớp dưới cùng một mẫu nguồn (giữ nguyên thứ tự checkpoint để ghép nối chính xác).
• Các nhóm bắt (capturing groups) trong mẫu nguồn có thể được tham chiếu dưới dạng \1, \2, … trong mẫu đích để giữ lại các phần của khóa gốc (ví dụ: chỉ số lớp).
• Các biến đổi có phạm vi (scope_prefix được đặt) chỉ khớp với các khóa bắt đầu bằng f"{scope_prefix}.".

Đăng ký ánh xạ chuyển đổi

Một danh sách chuyển đổi được đăng ký theo một khóa chuỗi — hoặc là model_type (ví dụ "mixtral") hoặc là tên lớp (ví dụ "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)],
        ),
    ],
)

Quy tắc tra cứu

Khi get_model_conversion_mapping xử lý một PreTrainedModel, mọi PreTrainedModel con đều được truy cập theo thứ tự DFS (nn.Module.named_modules() được lọc để chỉ giữ lại các instance của PreTrainedModel). Mỗi mô hình con được giải quyết qua hai giai đoạn: bộ lọc mã tùy chỉnh được chạy trước, sau đó tra cứu ánh xạ đã đăng ký.

Các mô hình mã tùy chỉnh bị bỏ qua trước khi bất kỳ thao tác tra cứu nào diễn ra. Một mô hình con được tính là mã tùy chỉnh khi is_custom_code() trả về True, bao gồm các mô hình được tải từ Hub với trust_remote_code=True và bất kỳ lớp mô hình nào không có trong Transformers. Các mô hình mã tùy chỉnh không kế thừa ánh xạ chuyển đổi vì tên lớp hoặc model_type của chúng có thể trùng với một mô hình bản địa. Ví dụ, một kiến trúc tùy chỉnh đặt model_type="mixtral" sẽ không nhận được các phép biến đổi chuyên gia hợp nhất có sẵn của Mixtral.

Đăng ký một model mã tùy chỉnh một cách tường minh bằng register_checkpoint_conversion_mapping để áp dụng ánh xạ chuyển đổi. Việc đăng ký sẽ thêm tên lớp hoặc model_type vào tập hợp các ánh xạ do người dùng đăng ký, qua đó miễn trừ model khỏi bước bỏ qua mã tùy chỉnh và định tuyến nó qua quy trình tra cứu thông thường bên dưới.

Với mỗi mô hình con vượt qua bộ lọc mã tùy chỉnh:

1. Tra cứu theo tên lớp được thực hiện trước, sau đó đến model_type. Nếu cả hai đều được đăng ký cho cùng một module, bản ánh tên lớp sẽ được ưu tiên và phần model_type sẽ bị bỏ qua cho module đó — điều này cho phép một task head (ví dụ LlavaForConditionalGeneration) ghi đè lên baseline model_type chung ("llava").
2. Mapping được chọn có scope_prefix được đặt thành đường dẫn dạng chấm của module con ("" cho gốc).
3. Ancestor-based deduplication quyết định có giữ lại ánh xạ hay không:
   • Nếu một đường dẫn cha đã khai báo cùng một định danh (tên lớp hoặc model_type), thì module con sẽ bị bỏ qua — ánh xạ không phạm vi hoặc phạm vi cao hơn của phần cha đã bao phủ cây con này.
   • Nếu chỉ có một anh chị em đã yêu cầu nó, mô-đun con sẽ được giữ lại với scope_prefix riêng của nó. Mỗi anh chị em sẽ có ánh xạ phạm vi riêng.

Các danh sách đã thấy cho tên lớp và model_type được theo dõi riêng biệt, với một điểm tinh tế: khi một module được khớp theo tên lớp của nó, model_type của module đó sẽ không được thêm vào danh sách đã thấy. Điều này nhằm đảm bảo các module khác chia sẻ cùng một model_type nhưng không có ánh xạ cụ thể cho lớp (ví dụ: DetrModel dưới DetrForSegmentation) vẫn có thể được truy cập thông qua tra cứu model_type.

Ánh xạ dựa trên lớp so với alias model_type

Cả hai phong cách cùng tồn tại trong _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",
    ...
}

Các khóa theo tên lớp được ưu tiên khi model_type được dùng chung nhưng một lớp cụ thể cần hành vi khác.

Tích hợp song song tensor

Hệ thống nạp động được tích hợp với song song tensor (TP) thông qua hệ thống phân cấp TensorParallelLayer được định nghĩa trong src/transformers/integrations/tensor_parallel.py.

Khi TP được bật, các tensor được chia nhỏ trong quá trình vật liệu hóa, chứ không phải sau đó. Điều này có nghĩa là mỗi rank chỉ tải phần tensor mà nó cần.

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)

Các kiểu song song khả dụng

Xử lý trọng số đóng gói

Đối với các trọng số đã hợp nhất như gate_up_proj, cần đặc biệt chú ý để chia nhỏ (shard) chính xác:

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

Thao tác TP được lưu trong WeightTransform và được áp dụng sau các thao tác chuyển đổi:

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

Tích hợp lượng tử hóa

Quantization tích hợp vào quy trình tải theo hai cách, tùy thuộc vào việc checkpoint đã được lượng tử hóa hay chưa:

• Checkpoint đã lượng tử hóa trước: Bộ lượng tử hóa cung cấp các instance WeightConverter (thông qua get_weight_conversions()) để giải chuỗi hóa các tensor đã lượng tử hóa. Kiểu dữ liệu của checkpoint được giữ nguyên để tránh các phép ép kiểu không mong muốn.
• Trích lượng hóa theo thời gian thực: Bộ trích lượng hóa cung cấp một thao tác trích lượng hóa được áp dụng sau các thao tác chuyển đổi, trích lượng hóa trọng số khi chúng được nạp.

Bộ lượng tử hóa cũng có thể viết lại toàn bộ danh sách chuyển đổi ở cuối get_model_conversion_mapping thông qua update_weight_conversions(...) — chẳng hạn, bộ khử lượng tử FP8 chèn một op Fp8Dequantize vào đầu mọi bộ chuyển đổi hiện có để các thang đo theo khối được áp dụng trước khi các op ghép chuyên gia / concat làm phẳng cấu trúc theo chuyên gia.

Tải mô hình nhanh và hiệu quả

Việc tải một model nhanh hơn và sử dụng ít bộ nhớ hơn vì trình tải biết các tensor nào là cần thiết cho các thao tác và lên lịch vật liệu hóa chúng một cách lười biếng.

Loader quét checkpoint một lần để phát hiện các khớp mẫu và thu thập các tensor. Nó lưu chúng dưới dạng các đối tượng Future và gửi chúng đến một nhóm luồng để tải bất đồng bộ mà không chặn GIL. Một tham số bắt đầu tải ngay khi có một luồng khả dụng cho nó.

Nếu hệ thống của bạn đang chạy các tiến trình nặng khác, nhiều luồng có thể làm chậm quá trình tải thay vì tăng tốc. Trong trường hợp này, hãy đặt biến môi trường HF_DEACTIVATE_ASYNC_LOAD=1 để tải trọng số theo tuần tự.

Mặc định là 4 luồng để tải tham số bất đồng bộ. Điều này mang lại sự cân bằng tốt nhất giữa các kịch bản tải và phần cứng. Công việc chủ yếu bị giới hạn bởi I/O, nhưng tùy thuộc vào phần cứng bộ tăng tốc và dtype cần thiết khi tải, nó có thể trở nên bị giới hạn bởi CPU/GPU nếu dtype khác với loại đã được lưu trữ (điều này yêu cầu một thao tác sao chép bổ sung).

Tải bất đồng bộ so với tải đồng bộ

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)

Việc tải đồng bộ được sử dụng khi:

• biến môi trường HF_DEACTIVATE_ASYNC_LOAD=1 được đặt.
• Việc offload xuống ổ đĩa đã được bật (ràng buộc bộ nhớ yêu cầu nạp tuần tự).
• Lượng tử hóa theo thời gian thực được bật (tránh việc các luồng worker chạy trước bước lượng tử hóa).

Luồng vật liệu hóa

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)

Hiệu quả bộ nhớ

Khi chuyển đổi một trọng số, trình chuyển đổi sẽ chờ cho đến khi tất cả các tensor cần thiết được tạo ra nếu chúng chưa được nạp. Ví dụ, thao tác MergeModulelist yêu cầu tất cả các trọng số trong ModuleList phải được nạp trước khi gộp.

Việc nối các tensor đòi hỏi một bản sao tạm thời, do đó các thao tác như MergeModulelist và Concatenate cần gấp đôi bộ nhớ so với các tensor gốc trong quá trình chuyển đổi. Sau khi hợp nhất, chỉ tensor kết quả còn lại trong bộ nhớ. Đỉnh bộ nhớ lý thuyết trong trường hợp xấu nhất là kích thước mô hình cộng với các tensor cần thiết cho thao tác MergeModulelist hoặc Concatenate lớn nhất.

Trường hợp xấu nhất này chỉ xảy ra khi tất cả các tham số khác đã được tải trước khi quá trình chuyển đổi đòi hỏi nhiều tài nguyên được thực thi. Hai kịch bản kích hoạt điều này.

1. Tất cả các tham số được tải bất đồng bộ trước khi bước vào quá trình chuyển đổi đòi hỏi nhiều tài nguyên (bộ luồng nhanh hơn hàng đợi chuyển đổi).
2. Việc chuyển đổi đòi hỏi cao nhất là việc cuối cùng.

Ví dụ, một mô hình MoE sử dụng MergeModulelist cho các chuyên gia trên mỗi lớp, đỉnh bộ nhớ trong trường hợp xấu nhất về mặt lý thuyết là kích thước mô hình cộng với các chuyên gia trên một lớp.

Các kịch bản xấu nhất này không phổ biến. Đỉnh bộ nhớ thực tế thường duy trì ở mức gần với kích thước mô hình.

Tính khả nghịch

Hệ thống hỗ trợ lưu trữ các mô hình với các phép biến đổi ngược, cho phép lưu/tải lại qua lại (round-trip). Việc lưu diễn ra theo hai giai đoạn: trước tiên là các bộ chuyển đổi ngược (mỗi tensor khớp với tối đa một bộ chuyển đổi), sau đó là các lệnh đổi tên ngược, theo quy tắc sắp xếp ở trên.

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)

Danh sách các transform được sử dụng khi tải được lưu đệm trên model dưới dạng _weight_conversions (chỉ giữ các mục thực sự được kích hoạt, do đó các đổi tên không song ánh như PrefixChange được áp dụng lại đối xứng một cách chính xác). Khi model được khởi tạo mà không dùng from_pretrained (và do đó không có _weight_conversions), revert_weight_conversion sẽ chuyển sang tính lại ánh xạ thông qua get_model_conversion_mapping và loại bỏ mọi PrefixChange khỏi đó (chúng ta không thể biết được checkpoint gốc có tiền tố hay không).

Các mẫu đích có thể chứa các phần tử regex cần được xử lý cho hướng ngược lại:

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

Ví dụ thực tế

MoE theo phong cách Mixtral

Định dạng 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

Định dạng mô hình:

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

Bản đồ chuyển đổi (từ 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)],
    ),
],

Mô hình thị giác-ngôn ngữ tổng hợp (Gemma3)

Gemma3 có cấu hình chuẩn “logic tiền tố khác nhau cho mô hình đầu so với mô hình cơ sở”. Hai mục được đăng ký (ở đây thông qua các alias trỏ đến các danh sách "llava" / "LlavaModel" dùng chung):

# 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

DFS walk cho Gemma3ForConditionalGeneration:

1. Root — tra cứu class không khớp (Gemma3ForConditionalGeneration không được đăng ký), tra cứu model_type khớp ("gemma3") → các phép biến đổi viết lại tiền tố "llava" được thêm không có phạm vi, đặt language_model, vision_tower, multi_modal_projector dưới model.*.
2. Inner model: Gemma3Model — lớp tra cứu khớp (Gemma3Model → LlavaModel) → chỉ ánh xạ LlavaModel được áp dụng, giới hạn trong "model". Việc đổi tên tiền tố rộng hơn nhiều của "gemma3" (="llava") không được áp dụng lại ở đây, đây chính là điều bạn muốn: mô hình bên trong nằm trong một không gian tên mà tiền tố đã đúng.

Cùng một hình dạng cặp được sử dụng bởi mọi VLM thuộc họ Llava (PaliGemma, InternVL, Mistral3, …). Mục model_type xử lý việc chỉnh sửa tiền tố của mô hình head; mục class giữ ánh xạ của mô hình cơ sở bên trong ở mức tối thiểu.

Lồng ghép sâu hơn với ghi đè cụ thể cho đầu (DETR)

DetrForSegmentation cho thấy một lớp ghi đè cụ thể cho head được xếp chồng lên hai mức lồng nhau:

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)

Có hai ánh xạ liên quan (một cho mỗi khóa đã đăng ký):

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

Đi theo DFS:

1. Root DetrForSegmentation khớp theo lớp → các tên đặt lại cho phân đoạn được thêm không phạm vi.
2. detr: DetrForObjectDetection không được đăng ký → không thêm transforms nào; DFS tiếp tục đi vào nó.
3. detr.model: DetrModel khớp theo lớp → các biến đổi cơ sở được thêm với scope_prefix="detr.model".

Đây là lý do tại sao DetrForObjectDetection không cần ánh xạ riêng: ánh xạ duy nhất được đăng ký trong cây con của nó (DetrModel) sẽ tự động được giới hạn phạm vi vào đường dẫn đúng.

Biệt danh theo khóa lớp tái sử dụng ánh xạ cơ bản mà không cần đăng ký thêm: "MaskFormerDetrDecoder": "DetrModel" cho phép bộ giải mã MaskFormer áp dụng cùng các phép biến đổi dưới tên lớp của chính nó.

Các thao tác tùy chỉnh (ERNIE 4.5 VL MoE)

Khi các thao tác tích hợp không đủ, bạn có thể tạo một lớp con ConversionOps tùy chỉnh. Ví dụ, ERNIE 4.5 VL MoE cần tách một danh sách chuyên gia dùng chung giữa các chế độ đầu vào văn bản và thị giác — điều mà không có thao tác tích hợp đơn lẻ nào xử lý được. Thao tác tùy chỉnh ErnieFuseAndSplitTextVisionExperts sẽ tách và xếp lại các chuyên gia trên hai khóa đích:

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

Các op tùy chỉnh phải triển khai convert() và thuộc tính reverse_op để hỗ trợ lưu/tải vòng tròn.

Biệt danh kiểu mô hình

Nhiều mô hình chia sẻ các mẫu chuyển đổi:

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

Tái sử dụng các khối xây dựng của việc tải động

Việc tải trọng số động không bị giới hạn ở các checkpoint của toàn bộ mô hình. Các thành phần xây dựng tương tự cho phép bạn tải bất kỳ bộ trọng số nào, miễn là bạn có thể mô tả cách các khóa trong checkpoint ánh xạ đến các tham số và đảm bảo các module đích tồn tại.

Nhìn chung, hợp đồng có dạng như sau:

1. Chuẩn bị không gian tên của mô hình. Đảm bảo các module/parameter bạn muốn nạp có sẵn và được đặt tên theo cách mà ánh xạ của bạn sẽ nhắm đến. Đối với adapter, điều này có nghĩa là gọi inject_adapter_in_model(...) để các module adapter tồn tại trước khi nạp. Đối với các head tùy chỉnh hoặc module bổ sung, hãy khởi tạo chúng trên mô hình trước.
2. Describe how to map weights. Xây dựng danh sách chuyển đổi/đổi tên (ví dụ, trong một hàm trợ giúp như _build_peft_weight_mapping(...)) sử dụng WeightConverter hoặc WeightRenaming. Đây là nơi bạn mô tả cách các khóa checkpoint nên được chuyển đổi, tách, gộp hoặc đổi tên để khớp với không gian tên của mô hình. Bạn có thể thực hiện chủ yếu 3 việc:
   • thêm các thao tác vào danh sách bộ chuyển đổi: chúng sẽ được áp dụng cho tất cả các trọng số, ngoại trừ những trọng số đã được thu thập trong bất kỳ WeightConverter nào. Nhìn chung, đây nên là các thao tác WeightRenaming
   • thêm các thao tác vào danh sách thao tác của từng bộ chuyển đổi: đây là điều xảy ra với Quantization, trong đó chúng ta chỉ cần thêm một thao tác lượng tử hóa sau danh sách thao tác của bất kỳ WeightConverter nào.
   • thay thế các thao tác / map bằng các thao tác tùy chỉnh của bạn: đây là điều xảy ra với peft. Chúng ta thay thế thao tác Concatenate của, ví dụ, mixtral, thành PeftConcatenate (được định nghĩa trong PEFT). Nhờ đó, khi checkpoint adapter được đọc, các trọng số cần được nối sẽ được thu thập và được định dạng đúng cho peft
3. Load + finalize + report. Sử dụng trình nạp lõi để thực hiện chuyển đổi và nạp tensor, sau đó hoàn tất và ghi log kết quả. Cụ thể, luồng này là:
   • LoadStateDictConfig(...) + _load_pretrained_model(...) để tải và chuyển đổi.
   • _finalize_load_state_dict(...) để di chuyển các tensor bị thiếu/không khớp khỏi meta, khởi tạo chúng và liên kết trọng số.
   • log_state_dict_report(...) để báo cáo các khóa bị thiếu/bất ngờ/không khớp (và các lỗi chuyển đổi).

Các API này được công khai để bạn có thể xử lý mã tùy chỉnh, định dạng trọng số tùy chỉnh, đồng thời đảm bảo bạn tận dụng được việc tải trọng số nhanh và hiệu quả nhất, cùng với việc chia nhỏ (sharding) và trải nghiệm sử dụng tốt của API transformers!

Tham chiếu các tệp chính