Mô hình tùy chỉnh

Một số kỹ thuật tinh chỉnh, chẳng hạn như prompt tuning, chỉ áp dụng cho mô hình ngôn ngữ. Điều đó có nghĩa là trong 🤗 PEFT, ta giả định rằng một mô hình 🤗 Transformers đang được sử dụng. Tuy nhiên, các kỹ thuật tinh chỉnh khác - như LoRA - không bị giới hạn ở một loại mô hình cụ thể.

Trong hướng dẫn này, chúng ta sẽ xem cách áp dụng LoRA cho một mạng cảm nhận đa lớp (multilayer perceptron), một mô hình thị giác máy tính từ thư viện timm, hoặc một kiến trúc 🤗 Transformers mới.

Mạng nơ-ron cảm nhận nhiều lớp

Giả sử chúng ta muốn tinh chỉnh một mạng cảm nhận đa lớp (multilayer perceptron) bằng LoRA. Đây là định nghĩa:

from torch import nn


class MLP(nn.Module):
    def __init__(self, num_units_hidden=2000):
        super().__init__()
        self.seq = nn.Sequential(
            nn.Linear(20, num_units_hidden),
            nn.ReLU(),
            nn.Linear(num_units_hidden, num_units_hidden),
            nn.ReLU(),
            nn.Linear(num_units_hidden, 2),
            nn.LogSoftmax(dim=-1),
        )

    def forward(self, X):
        return self.seq(X)

Đây là một mạng cảm nhận đa tầng đơn giản với một lớp đầu vào, một lớp ẩn và một lớp đầu ra.

Đối với ví dụ đơn giản này, chúng ta chọn một số lượng đơn vị ẩn cực kỳ lớn để làm nổi bật lợi ích về hiệu quả từ PEFT, nhưng những lợi ích này cũng phù hợp với các ví dụ thực tế hơn.

Mô hình này có một vài lớp tuyến tính có thể được tinh chỉnh bằng LoRA. Khi làm việc với các mô hình 🤗 Transformers phổ biến, PEFT sẽ biết các lớp nào cần áp dụng LoRA, nhưng trong trường hợp này, người dùng sẽ phải tự chọn các lớp. Để xác định tên của các lớp cần tinh chỉnh:

print([(n, type(m)) for n, m in MLP().named_modules()])

Điều này nên in ra:

[('', __main__.MLP),
 ('seq', torch.nn.modules.container.Sequential),
 ('seq.0', torch.nn.modules.linear.Linear),
 ('seq.1', torch.nn.modules.activation.ReLU),
 ('seq.2', torch.nn.modules.linear.Linear),
 ('seq.3', torch.nn.modules.activation.ReLU),
 ('seq.4', torch.nn.modules.linear.Linear),
 ('seq.5', torch.nn.modules.activation.LogSoftmax)]

Giả sử chúng ta muốn áp dụng LoRA vào lớp đầu vào và lớp ẩn, đó là 'seq.0' và 'seq.2'. Hơn nữa, giả sử chúng ta muốn cập nhật lớp đầu ra mà không dùng LoRA, đó sẽ là 'seq.4'. Cấu hình tương ứng sẽ là:

from peft import LoraConfig

config = LoraConfig(
    target_modules=["seq.0", "seq.2"],
    modules_to_save=["seq.4"],
)

Như vậy, chúng ta có thể tạo mô hình PEFT của mình và kiểm tra tỷ lệ tham số được huấn luyện:

from peft import get_peft_model

model = MLP()
peft_model = get_peft_model(model, config)
peft_model.print_healthcheck()

Cuối cùng, chúng ta có thể sử dụng bất kỳ framework huấn luyện nào mình thích, hoặc tự viết vòng lặp fit, để huấn luyện peft_model.

Để xem một ví dụ hoàn chỉnh, hãy xem notebook này.

mô hình timm

Thư viện timm chứa một số lượng lớn các mô hình thị giác máy tính được huấn luyện trước. Các mô hình này cũng có thể được tinh chỉnh với PEFT. Hãy cùng xem cách thức hoạt động của nó trong thực tế.

Để bắt đầu, hãy đảm bảo rằng timm đã được cài đặt trong môi trường Python:

python -m pip install -U timm

Tiếp theo, chúng ta tải một mô hình timm cho tác vụ phân loại ảnh:

import timm

num_classes = ...
model_id = "timm/poolformer_m36.sail_in1k"
model = timm.create_model(model_id, pretrained=True, num_classes=num_classes)

Một lần nữa, chúng ta cần đưa ra quyết định về việc áp dụng LoRA cho những lớp nào. Vì LoRA hỗ trợ các lớp conv 2D, và vì đây là một thành phần cấu trúc chính của mô hình này, chúng ta nên áp dụng LoRA cho các lớp conv 2D. Để xác định tên của các lớp đó, hãy xem xét tất cả các tên lớp:

print([(n, type(m)) for n, m in model.named_modules()])

Điều này sẽ in ra một danh sách rất dài, chúng ta chỉ hiển thị vài mục đầu tiên:

[('', timm.models.metaformer.MetaFormer),
 ('stem', timm.models.metaformer.Stem),
 ('stem.conv', torch.nn.modules.conv.Conv2d),
 ('stem.norm', torch.nn.modules.linear.Identity),
 ('stages', torch.nn.modules.container.Sequential),
 ('stages.0', timm.models.metaformer.MetaFormerStage),
 ('stages.0.downsample', torch.nn.modules.linear.Identity),
 ('stages.0.blocks', torch.nn.modules.container.Sequential),
 ('stages.0.blocks.0', timm.models.metaformer.MetaFormerBlock),
 ('stages.0.blocks.0.norm1', timm.layers.norm.GroupNorm1),
 ('stages.0.blocks.0.token_mixer', timm.models.metaformer.Pooling),
 ('stages.0.blocks.0.token_mixer.pool', torch.nn.modules.pooling.AvgPool2d),
 ('stages.0.blocks.0.drop_path1', torch.nn.modules.linear.Identity),
 ('stages.0.blocks.0.layer_scale1', timm.models.metaformer.Scale),
 ('stages.0.blocks.0.res_scale1', torch.nn.modules.linear.Identity),
 ('stages.0.blocks.0.norm2', timm.layers.norm.GroupNorm1),
 ('stages.0.blocks.0.mlp', timm.layers.mlp.Mlp),
 ('stages.0.blocks.0.mlp.fc1', torch.nn.modules.conv.Conv2d),
 ('stages.0.blocks.0.mlp.act', torch.nn.modules.activation.GELU),
 ('stages.0.blocks.0.mlp.drop1', torch.nn.modules.dropout.Dropout),
 ('stages.0.blocks.0.mlp.norm', torch.nn.modules.linear.Identity),
 ('stages.0.blocks.0.mlp.fc2', torch.nn.modules.conv.Conv2d),
 ('stages.0.blocks.0.mlp.drop2', torch.nn.modules.dropout.Dropout),
 ('stages.0.blocks.0.drop_path2', torch.nn.modules.linear.Identity),
 ('stages.0.blocks.0.layer_scale2', timm.models.metaformer.Scale),
 ('stages.0.blocks.0.res_scale2', torch.nn.modules.linear.Identity),
 ('stages.0.blocks.1', timm.models.metaformer.MetaFormerBlock),
 ('stages.0.blocks.1.norm1', timm.layers.norm.GroupNorm1),
 ('stages.0.blocks.1.token_mixer', timm.models.metaformer.Pooling),
 ('stages.0.blocks.1.token_mixer.pool', torch.nn.modules.pooling.AvgPool2d),
 ...
 ('head.global_pool.flatten', torch.nn.modules.linear.Identity),
 ('head.norm', timm.layers.norm.LayerNorm2d),
 ('head.flatten', torch.nn.modules.flatten.Flatten),
 ('head.drop', torch.nn.modules.linear.Identity),
 ('head.fc', torch.nn.modules.linear.Linear)]
 ]

Khi xem xét kỹ hơn, chúng ta thấy rằng các lớp conv 2D có tên như "stages.0.blocks.0.mlp.fc1" và "stages.0.blocks.0.mlp.fc2". Làm thế nào để khớp chính xác với những tên lớp đó? Bạn có thể viết một biểu thức chính quy để khớp với tên lớp. Trong trường hợp của chúng ta, regex r".*\.mlp\.fc\d" sẽ làm được việc này.

Hơn nữa, giống như trong ví dụ đầu tiên, chúng ta cần đảm bảo rằng lớp đầu ra, trong trường hợp này là phần phân loại (classification head), cũng được cập nhật. Nhìn vào cuối danh sách được in ở trên, chúng ta có thể thấy nó có tên là 'head.fc'. Với điều đó trong đầu, đây là cấu hình LoRA của chúng ta:

config = LoraConfig(target_modules=r".*\.mlp\.fc\d", modules_to_save=["head.fc"])

Sau đó, chúng ta chỉ cần tạo mô hình PEFT bằng cách truyền mô hình cơ sở và cấu hình của chúng ta vào get_peft_model:

peft_model = get_peft_model(model, config)
peft_model.print_healthcheck()

Điều này cho thấy chúng ta chỉ cần huấn luyện dưới 2% tổng số tham số, một mức tăng hiệu quả rất lớn.

Để xem một ví dụ hoàn chỉnh, hãy xem notebook này.

Các kiến trúc transformers mới

Khi các kiến trúc transformers phổ biến mới được phát hành, chúng tôi cố gắng hết sức để nhanh chóng thêm chúng vào PEFT. Nếu bạn gặp một model transformers không được hỗ trợ ngay từ đầu, đừng lo, nó rất có thể vẫn hoạt động nếu config được thiết lập đúng. Cụ thể, bạn phải xác định các lớp cần được thích ứng và thiết lập chúng chính xác khi khởi tạo lớp config tương ứng, ví dụ LoraConfig. Dưới đây là một số mẹo để hỗ trợ bạn trong việc này.

Bước đầu tiên, bạn nên kiểm tra các mô hình hiện có để tìm cảm hứng. Bạn có thể tìm thấy chúng trong constants.py trong kho PEFT. Thường thì bạn sẽ tìm thấy một kiến trúc tương tự sử dụng cùng các tên. Ví dụ, nếu kiến trúc mô hình mới là một biến thể của mô hình “mistral” và bạn muốn áp dụng LoRA, bạn có thể thấy rằng mục cho “mistral” trong TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING chứa ["q_proj", "v_proj"]. Điều này cho bạn biết rằng đối với các mô hình “mistral”, target_modules cho LoRA nên là ["q_proj", "v_proj"]:

from peft import LoraConfig, get_peft_model

my_mistral_model = ...
config = LoraConfig(
    target_modules=["q_proj", "v_proj"],
    ...,  # other LoRA arguments
)
peft_model = get_peft_model(my_mistral_model, config)

Nếu điều đó không hữu ích, hãy kiểm tra các module hiện có trong kiến trúc mô hình của bạn bằng phương thức named_modules và cố gắng xác định các lớp attention, đặc biệt là các lớp key, query và value. Chúng thường có tên như c_attn, query, q_proj, v.v. Lớp key không phải lúc nào cũng được điều chỉnh, và lý tưởng nhất là bạn nên kiểm tra xem việc bao gồm nó có mang lại hiệu suất tốt hơn hay không.

Ngoài ra, các lớp tuyến tính thường là mục tiêu để thích ứng (ví dụ trong bài báo QLoRA, các tác giả đề xuất cũng nên thích ứng chúng). Tên của chúng thường chứa các chuỗi fc hoặc dense.

Nếu bạn muốn thêm một mô hình mới vào PEFT, hãy tạo một mục trong constants.py và mở một pull request trên repository. Đừng quên cập nhật cả README.

Xác minh các tham số và lớp

Bạn có thể kiểm tra xem mình đã áp dụng đúng phương pháp PEFT cho mô hình của mình hay không bằng một vài cách.

• Chạy print_healthcheck() để kiểm tra tỷ lệ tham số có thể huấn luyện và trạng thái adapter. Nếu tỷ lệ tham số thấp hơn hoặc cao hơn dự kiến, hãy kiểm tra repr của model bằng cách in model ra. Thao tác này hiển thị tên của tất cả các loại layer trong model. Đảm bảo rằng chỉ các layer đích dự kiến mới được thay thế bằng các layer adapter. Ví dụ, nếu LoRA được áp dụng cho các layer nn.Linear, thì bạn chỉ nên thấy các layer lora.Linear được sử dụng.

peft_model.print_healthcheck()

• Một cách khác để xem các lớp đã thích ứng là sử dụng thuộc tính targeted_module_names để liệt kê tên của từng module đã được thích ứng.

print(peft_model.targeted_module_names)

Các loại module không được hỗ trợ

Các phương pháp như LoRA chỉ hoạt động nếu các module đích được PEFT hỗ trợ. Ví dụ, bạn có thể áp dụng LoRA cho các lớp nn.Linear và nn.Conv2d, nhưng không thể áp dụng cho, chẳng hạn, nn.LSTM. Nếu bạn phát hiện một lớp mà bạn muốn áp dụng PEFT nhưng không được hỗ trợ, bạn có thể:

• định nghĩa ánh xạ tùy chỉnh để điều phối động các module tùy chỉnh trong LoRA
• mở một issue và yêu cầu tính năng đó, nơi các nhà bảo trì sẽ triển khai nó hoặc hướng dẫn bạn cách tự triển khai nếu nhu cầu đối với loại module này đủ cao

Hỗ trợ thử nghiệm cho việc điều phối động các module tùy chỉnh trong LoRA

Tính năng này mang tính thử nghiệm và có thể thay đổi tùy thuộc vào phản hồi từ cộng đồng. Chúng tôi sẽ giới thiệu một API công khai và ổn định nếu có nhu cầu đáng kể về tính năng này.

PEFT hỗ trợ một API thử nghiệm cho các loại module tùy chỉnh dành cho LoRA. Giả sử bạn có một cài đặt LoRA cho LSTMs. Thông thường, bạn sẽ không thể chỉ định cho PEFT sử dụng nó, ngay cả khi về mặt lý thuyết nó có thể hoạt động với PEFT. Tuy nhiên, điều này là có thể nhờ vào việc phân phát động (dynamic dispatch) các lớp tùy chỉnh.

API thử nghiệm hiện có dạng như sau:

class MyLoraLSTMLayer:
    ...

base_model = ...  # load the base model that uses LSTMs

# add the LSTM layer names to target_modules
config = LoraConfig(..., target_modules=["lstm"])
# define a mapping from base layer type to LoRA layer type
custom_module_mapping = {nn.LSTM: MyLoraLSTMLayer}
# register the new mapping
config._register_custom_module(custom_module_mapping)
# after registration, create the PEFT model
peft_model = get_peft_model(base_model, config)
# do training

Khi bạn gọi get_peft_model(), bạn sẽ thấy một cảnh báo vì PEFT không nhận diện được kiểu module được nhắm tới. Trong trường hợp này, bạn có thể bỏ qua cảnh báo này.

Khi cung cấp một ánh xạ tùy chỉnh, PEFT trước tiên kiểm tra các lớp của mô hình cơ sở đối chiếu với ánh xạ tùy chỉnh và chuyển đến kiểu lớp LoRA tùy chỉnh nếu có sự khớp. Nếu không có sự khớp, PEFT sẽ kiểm tra các kiểu lớp LoRA tích hợp sẵn để tìm sự khớp.

Do đó, tính năng này cũng có thể được sử dụng để ghi đè lên logic dispatch hiện có, ví dụ như khi bạn muốn sử dụng lớp LoRA riêng của mình cho nn.Linear thay vì dùng lớp do PEFT cung cấp.

Khi tạo module LoRA tùy chỉnh của bạn, vui lòng tuân theo các quy tắc giống như các module LoRA hiện có. Một số ràng buộc quan trọng cần xem xét:

• Module tùy chỉnh phải kế thừa từ nn.Module và peft.tuners.lora.layer.LoraLayer.
• Phương thức __init__ của module tùy chỉnh phải có các đối số theo vị trí base_layer và adapter_name. Sau đó, có các **kwargs bổ sung mà bạn có thể tự do sử dụng hoặc bỏ qua.
• Các tham số có thể học được nên được lưu trữ trong một nn.ModuleDict hoặc nn.ParameterDict, trong đó khóa tương ứng với tên của adapter cụ thể (hãy nhớ rằng một mô hình có thể có nhiều hơn một adapter cùng một lúc).
• Tên của các thuộc tính tham số có thể học này nên bắt đầu bằng "lora_", ví dụ self.lora_new_param = ....
• Một số phương thức là tùy chọn, ví dụ bạn chỉ cần triển khai merge và unmerge nếu muốn hỗ trợ gộp trọng số.

Hiện tại, thông tin về module tùy chỉnh không được lưu lại khi bạn lưu mô hình. Khi tải mô hình, bạn phải đăng ký lại các module tùy chỉnh.

# saving works as always and includes the parameters of the custom modules
peft_model.save_pretrained(<model-path>)

# loading the model later:
base_model = ...
# load the LoRA config that you saved earlier
config = LoraConfig.from_pretrained(<model-path>)
# register the custom module again, the same way as the first time
custom_module_mapping = {nn.LSTM: MyLoraLSTMLayer}
config._register_custom_module(custom_module_mapping)
# pass the config instance to from_pretrained:
peft_model = PeftModel.from_pretrained(model, tmp_path / "lora-custom-module", config=config)

Nếu bạn sử dụng tính năng này và thấy nó hữu ích, hoặc nếu bạn gặp phải các vấn đề, hãy cho chúng tôi biết bằng cách tạo một issue hoặc một discussion trên GitHub. Điều này cho phép chúng tôi ước lượng nhu cầu đối với tính năng này và thêm một API công khai nếu nhu cầu đủ cao.