Trình tối ưu hóa và trình lập lịch

Optimizer cập nhật trọng số mô hình trong quá trình huấn luyện. Scheduler bao quanh optimizer và điều chỉnh tốc độ học ở mỗi bước huấn luyện. Trainer tạo cả hai khi nó gọi create_optimizer_and_scheduler().

                                    ┌────────────┐         ┌──────────────┐
                                    │ Optimizer  │         │  Scheduler   │
                                    │ (adamw_torch_fused)◄─│  (linear)    │
                                    │            │         │              │
                                    │ param_groups         |              |
                                    │  └ lr       ◄────────┤              |
                                    │  └ weight_decay      │              │
                                    └──────┬─────┘         └──────────────┘
                                           │                      
  ┌──── EACH TRAINING STEP ───────────────────────────────────────────┐
  │                                        │                          │
  │   model(batch)                         │                          │
  │       │                                │                          │
  │       ▼                                │                          │
  │     loss ──► loss.backward() ──► param.grad                       │
  │                                        │                          │
  │                          ┌─────────────┘                          │
  │                          ▼                                        │
  │              optimizer.step()                                     │
  │                          │                                        │
  │                          ▼                                        │
  │                   param.data updated                              │
  │                          │                                        │
  │                          ▼                                        │
  │              lr_scheduler.step()  ──► recalculates lr             │
  │                          │            writes to optimizer         │
  │                          ▼            .param_groups['lr']         │
  │              model.zero_grad()                                    │
  │                                                                   │
  └───────────────────────────────────────────────────────────────────┘

Cấu hình hành vi của optimizer và scheduler, chẳng hạn như lr_scheduler_type() và optim(), trong TrainingArguments. Các giá trị mặc định (optimizer adamw_torch và scheduler warmup linear) là điểm khởi đầu tốt cho hầu hết các lần fine-tuning.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    # Optimizer
    optim="adamw_torch",          # or "adamw_torch_fused", "adafactor", "sgd", etc.
    learning_rate=2e-5,
    weight_decay=0.01,
    adam_beta1=0.9,
    adam_beta2=0.999,
    adam_epsilon=1e-8,
    # Scheduler
    lr_scheduler_type="cosine",   # "linear", "cosine", "constant_with_warmup", etc.
    warmup_steps=500,
    lr_scheduler_kwargs={"num_cycles": 3},  # scheduler-specific extras
)

Trình lập lịch dựa trên chỉ số

Một số trình lập lịch thích ứng với động lực huấn luyện thay vì tuân theo một lịch trình cố định.

GreedyLR cập nhật tốc độ học từ kết quả đánh giá. Nó tăng tốc độ học bằng cách chia cho factor khi chỉ số tiếp tục cải thiện, và giảm tốc độ học bằng cách nhân với factor khi chỉ số không cải thiện. Khi tốc độ học dừng lại ở min_lr và không cải thiện sau reset_start bước, GreedyLR sẽ đặt lại về trạng thái ban đầu và bắt đầu một chu kỳ mới.

GreedyLR yêu cầu đánh giá trong quá trình huấn luyện. Đặt eval_strategy thành "steps" hoặc "epoch".

args = TrainingArguments(
+   lr_scheduler_type="greedy",
+   lr_scheduler_kwargs={"patience": 10, "factor": 0.95, "min_lr": 1e-5},
+   eval_strategy="steps",
+   eval_steps=200,
    ...  # remaining args from the TrainingArguments intro config
)

Chế độ mặc định mode="min" hoạt động với loss. Nếu bạn đang theo dõi một chỉ số mà giá trị cao hơn là tốt hơn, chẳng hạn như accuracy, hãy truyền "mode": "max" trong lr_scheduler_kwargs.

Xem lớp GreedyLR để biết danh sách đầy đủ các tham số có thể cấu hình.

Tích hợp bộ tối ưu

Transformers tích hợp các optimizer bên thứ ba cho các kịch bản huấn luyện chuyên biệt.

Trình tối ưu	Cài đặt	`optim="value"`	Mô tả
APOLLO	`apollo-torch`	`apollo_adamw`	Hiệu quả bộ nhớ cho toàn bộ tham số thông qua các phép chiếu ngẫu nhiên; hạng 1 là đủ
GrokAdamW	`grokadamw`	`grokadamw`	Định hướng khái quát hóa chậm (grokking)
LOMO / AdaLomo	`lomo-optim`	`lomo` / `adalomo`	Kết hợp bước gradient + cập nhật cho việc tinh chỉnh tham số đầy đủ với bộ nhớ thấp
Lịch miễn phí	`schedulefree`	`schedule_free_adamw`, `schedule_free_radam`, `schedule_free_sgd`	Loại bỏ việc giảm dần LR; kết hợp với `lr_scheduler_type="constant"`
GaLore	`galore-torch`	`galore_adamw`, `galore_adafactor`, `galore_adamw_8bit`	Học toàn bộ tham số thông qua chiếu hạng thấp của gradient
StableAdamW	`torch-optimi`	`stable_adamw`	Cắt bỏ cập nhật AdamW + AdaFactor; không cần cắt bỏ gradient

APOLLO

GrokAdamW

LOMO

Lịch miễn phí

StableAdamW

GaLore

pip install apollo-torch

Approximated Gradient Scaling for Memory Efficient LLM Optimization (APOLLO) là bộ tối ưu hóa tiết kiệm bộ nhớ cho việc học toàn bộ tham số trong quá trình huấn luyện trước và tinh chỉnh. Nó đạt hiệu suất tương đương AdamW với chi phí bộ nhớ tương tự SGD bằng cách sử dụng các phép chiếu ngẫu nhiên rẻ thay vì SVD. Để tiết kiệm bộ nhớ tối đa, hãy sử dụng APOLLO-Mini, một biến thể rank-1.

Sử dụng tham số optim_target_modules để chỉ định các lớp cần huấn luyện.

args = TrainingArguments(
+   optim="apollo_adamw",
+   optim_target_modules=[r".*.attn.*", r".*.mlp.*"],
    ...  # remaining args from the TrainingArguments intro config
)

Truyền các siêu tham số bổ sung thông qua optim_args.

Đặt scale thành n/r, trong đó n là chiều không gian gốc và r là chiều không gian hạng thấp. Việc điều chỉnh tốc độ học trong khi giữ scale ở giá trị mặc định sẽ đạt được hiệu ứng tương tự.

parameter	mô tả	APOLLO	APOLLO-Mini
rank	bậc của không gian con phụ trợ cho việc điều chỉnh tỷ lệ gradient	256	1
scale_type	cách các hệ số mở rộng được áp dụng	`channel` (thang theo từng kênh)	`tensor` (thang đo theo từng tensor)
scale	điều chỉnh các bản cập nhật gradient để ổn định quá trình huấn luyện	1.0	128
update_proj_gap	các bước trước khi cập nhật các ma trận chiếu	200	200
proj	kiểu chiếu	`random`	`random`

Bật APOLLO-Mini với cấu hình rank-1.

args = TrainingArguments(
    optim="apollo_adamw",
    optim_target_modules=[r".*.attn.*", r".*.mlp.*"],
    optim_args="proj=random,rank=1,scale=128.0,scale_type=tensor,update_proj_gap=200",
    ...  # remaining args from the TrainingArguments intro config
)

Tùy chỉnh bộ tối ưu hóa và trình lập lịch

Tạo một optimizer và scheduler tùy chỉnh để sử dụng một optimizer chưa được tích hợp, điều chỉnh tốc độ học theo từng lớp, hoặc áp dụng logic tùy chỉnh.

Truyền một class và kwargs

~Trainer.optimizer_cls_and_kwargs chấp nhận một lớp optimizer tùy chỉnh trong khi giao phó việc nhóm tham số và đặt thiết bị cho Trainer.

Trainer trì hoãn việc tạo optimizer cho đến khi create_optimizer() được thực thi, do đó mô hình đã nằm trên thiết bị đúng.

import torch

trainer = Trainer(
    ...
    optimizer_cls_and_kwargs=(
        torch.optim.SGD,
        {"momentum": 0.9, "nesterov": True}
    ),
)

Truyền các instance đã được xây dựng sẵn

Truyền một optimizer và scheduler đã định nghĩa sẵn vào ~Trainer.optimizers. Trainer sẽ bỏ qua create_optimizer() và create_scheduler() khi các instance đã được xây dựng sẵn được cung cấp. Nếu bạn không truyền scheduler, Trainer sẽ tự động tạo một.

Xây dựng optimizer sau khi đã đặt model lên đúng thiết bị. Các tham số được xác định tại thời điểm khởi tạo, trước khi Trainer di chuyển model. Trong huấn luyện phân tán, việc thiết bị không khớp có thể âm thầm gây ra hành vi sai.

import torch
from transformers import Trainer, get_cosine_schedule_with_warmup

optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
scheduler = get_cosine_schedule_with_warmup(
    optimizer, num_warmup_steps=500, num_training_steps=10_000
)

trainer = Trainer(
    ...
    optimizers=(optimizer, scheduler),
)

Các instance được xây dựng sẵn bỏ qua create_optimizer() và create_scheduler(), vì vậy bạn cần tự chỉ định các nhóm tham số của riêng mình.

Ghi đè các phương thức của optimizer và scheduler

Định nghĩa lớp con của create_optimizer() và create_scheduler() để kiểm soát hoàn toàn. Cả hai phương thức này đều chạy trong quá trình train().

Ghi đè create_scheduler() để sử dụng một scheduler như OneCycleLR không có sẵn trong SchedulerType.

Với mỗi phương thức, hãy đảm bảo gán cho self và trả về nó.

import torch
from transformers import Trainer

class MyTrainer(Trainer):

    def create_scheduler(self, num_training_steps, optimizer=None):
        optimizer = optimizer or self.optimizer
        self.lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
            optimizer,
            max_lr=0.1,
            total_steps=num_training_steps,
        )
        return self.lr_scheduler

Bạn không cần ghi đè create_optimizer() nếu optimizer mặc định hoạt động tốt. Mở rộng một phương thức bằng super() dễ dàng hơn so với thay thế hoàn toàn. Ví dụ, thêm một nhóm tham số bổ sung trong khi giữ nguyên mọi thứ khác.

class MyTrainer(Trainer):
    def create_optimizer(self, model=None):
        super().create_optimizer(model)  # builds the default two param groups
        # add extra param group
        self.optimizer.add_param_group({
            "params": self.model.classifier.parameters(),
            "lr": self.args.learning_rate * 10,
        })
        return self.optimizer

Update on GitHub