Trainer Chưng cất Kiến thức Tổng quát

Tổng quan

Generalized Knowledge Distillation (GKD) được đề xuất trong On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes bởi Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist và Olivier Bachem.

Tóm tắt từ bài báo như sau:

Nén mô hình (Knowledge distillation - KD) được sử dụng rộng rãi để giảm chi phí suy luận và dung lượng bộ nhớ của mô hình giáo viên (teacher model) bằng cách huấn luyện một mô hình học sinh (student model) nhỏ hơn. Tuy nhiên, các phương pháp KD hiện tại cho các mô hình chuỗi tự hồi quy (auto-regressive sequence models) gặp phải vấn đề lệch phân phối giữa các chuỗi đầu ra được xem trong quá trình huấn luyện và các chuỗi do mô hình học sinh tạo ra trong quá trình suy luận. Để giải quyết vấn đề này, chúng tôi giới thiệu Nén kiến thức tổng quát (Generalized Knowledge Distillation - GKD). Thay vì chỉ dựa vào một tập hợp cố định các chuỗi đầu ra, GKD huấn luyện mô hình học sinh trên các chuỗi đầu ra do chính nó tạo ra bằng cách tận dụng phản hồi từ mô hình giáo viên trên các chuỗi đó. Khác với các tiếp cận KD có giám sát, GKD cũng cung cấp sự linh hoạt để sử dụng các hàm mất mát (loss functions) thay thế giữa mô hình học sinh và mô hình giáo viên, điều này có thể hữu ích khi mô hình học sinh thiếu khả năng diễn đạt (expressivity) để bắt chước phân phối của mô hình giáo viên. Hơn nữa, GKD tạo điều kiện thuận lợi cho việc tích hợp liền mạch quá trình nén kiến thức với tinh chỉnh bằng học tăng cường (RL fine-tuning, RLHF). Chúng tôi chứng minh hiệu quả của GKD trong việc nén các mô hình ngôn ngữ tự hồi quy trên các tác vụ tóm tắt, dịch thuật và suy luận toán học, cũng như nén kiến thức không phụ thuộc tác vụ (task-agnostic distillation) cho tinh chỉnh theo hướng dẫn (instruction-tuning).

Các khía cạnh chính của GKD là:

Nó giải quyết sự lệch phân phối giữa huấn luyện và suy luận trong các mô hình chuỗi tự hồi quy bằng cách huấn luyện mô hình học sinh trên các chuỗi đầu ra do chính nó tạo ra.
GKD cho phép linh hoạt trong việc lựa chọn các phép đo phân kỳ khác nhau giữa mô hình học sinh và mô hình giáo viên thông qua phân kỳ Jensen-Shannon tổng quát (JSD), điều này có thể hữu ích khi mô hình học sinh không đủ khả năng để mô phỏng hoàn toàn mô hình giáo viên.

Phương pháp hậu huấn luyện này được đóng góp bởi Kashif Rasul và Lewis Tunstall.

Mẹo sử dụng

experimental.gkd.GKDTrainer là một lớp bao quanh lớp SFTTrainer, nhận một đối số mô hình giáo viên. Nó cần ba tham số được đặt thông qua experimental.gkd.GKDConfig, cụ thể là:

lmbda: kiểm soát tỷ lệ dữ liệu của student, tức là phần trăm các đầu ra do student tạo ra theo chính sách (on-policy). Khi lmbda=0.0, loss thu gọn về JSD có giám sát, trong đó student được huấn luyện với xác suất ở cấp độ token của teacher. Khi lmbda=1.0, loss thu gọn về JSD on-policy, trong đó student tạo ra các chuỗi đầu ra và nhận phản hồi cụ thể cho từng token trên các chuỗi này từ teacher. Với các giá trị nằm trong khoảng [0, 1], quá trình sẽ ngẫu nhiên giữa hai chế độ dựa trên giá trị lmbda cho mỗi batch.
seq_kd: kiểm soát việc có thực hiện Sequence-Level KD hay không (có thể xem là supervised FT trên out do teacher sinh ra). Khi seq_kd=True và lmbda=0.0, loss sẽ rút gọn về supervised JSD, trong đó teacher sinh ra các chuỗi output và student nhận phản hồi cụ thể cho từng token trên các chuỗi này từ teacher.
beta: kiểm soát sự nội suy trong phân kỳ Jensen-Shannon tổng quát. Khi beta=0.0, loss xấp xỉ phân kỳ KL thuận, trong khi với beta=1.0, loss xấp xỉ phân kỳ KL nghịch. Với các giá trị nằm trong khoảng [0, 1], nó nội suy giữa hai phân kỳ này.

Các tác giả nhận thấy rằng dữ liệu on-policy (giá trị lmbda cao) cho kết quả tốt hơn và giá trị tối ưu của beta thay đổi tùy thuộc vào tác vụ và phương pháp đánh giá.

Hãy đảm bảo rằng attn_implementation="kernels-community/flash-attn2" khi huấn luyện mô hình Gemma. Nếu không, bạn sẽ gặp các giá trị NaN trong logits do kỹ thuật soft capping được áp dụng bởi kiến trúc này.

API cơ bản như sau:

from datasets import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl.experimental.gkd import GKDConfig, GKDTrainer

NUM_DUMMY_SAMPLES = 100

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
# The model to optimise
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
# The teacher model to calculate the KL divergence against
teacher_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-1.5B-Instruct")

train_dataset = Dataset.from_dict(
    {
        "messages": [
            [
                {"role": "user", "content": "Hi, how are you?"},
                {"role": "assistant", "content": "I'm great thanks"},
            ]
        ]
        * NUM_DUMMY_SAMPLES
    }
)
eval_dataset = Dataset.from_dict(
    {
        "messages": [
            [
                {"role": "user", "content": "What colour is the sky?"},
                {"role": "assistant", "content": "The sky is blue"},
            ]
        ]
        * NUM_DUMMY_SAMPLES
    }
)

training_args = GKDConfig(output_dir="gkd-model", per_device_train_batch_size=1)
trainer = GKDTrainer(
    model=model,
    teacher_model=teacher_model,
    args=training_args,
    processing_class=tokenizer,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)
trainer.train()

Kiểu dữ liệu tập dữ liệu dự kiến

Bộ dữ liệu nên được định dạng dưới dạng một danh sách các “messages”, trong đó mỗi message là một danh sách các từ điển có các khóa sau:

role: là system, assistant hoặc user
content: nội dung tin nhắn

GKDTrainer

class trl.experimental.gkd.GKDTrainer

< source >

( model: typing.Union[transformers.modeling_utils.PreTrainedModel, torch.nn.Module, str, NoneType] = Noneteacher_model: typing.Union[transformers.modeling_utils.PreTrainedModel, torch.nn.Module, str] = Noneargs: trl.experimental.gkd.gkd_config.GKDConfig | None = Nonedata_collator: collections.abc.Callable[[list[typing.Any]], dict[str, typing.Any]] | None = Nonetrain_dataset: datasets.arrow_dataset.Dataset | None = Noneeval_dataset: datasets.arrow_dataset.Dataset | dict[str, datasets.arrow_dataset.Dataset] | None = Noneprocessing_class: transformers.tokenization_utils_base.PreTrainedTokenizerBase | transformers.image_processing_utils.BaseImageProcessor | transformers.feature_extraction_utils.FeatureExtractionMixin | transformers.processing_utils.ProcessorMixin | None = Nonecompute_metrics: collections.abc.Callable[[transformers.trainer_utils.EvalPrediction], dict] | None = Nonecallbacks: list[transformers.trainer_callback.TrainerCallback] | None = Noneoptimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None)preprocess_logits_for_metrics: collections.abc.Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = Nonepeft_config: PeftConfig | None = Noneformatting_func: collections.abc.Callable | None = None )

Tham số

model (PreTrainedModel hoặc torch.nn.Module hoặc str, không bắt buộc) — Mô hình cần được huấn luyện, hoặc chuỗi định danh của mô hình cần được khởi tạo từ một mô hình được huấn luyện trước.
teacher_model (PreTrainedModel hoặc torch.nn.Module hoặc str, tuỳ chọn) — Mô hình teacher cho kỹ thuật chưng cất tri thức, hoặc chuỗi định danh của mô hình cần được khởi tạo từ một mô hình được huấn luyện trước.
args (experimental.gkd.GKDConfig, tùy chọn) — Các đối số huấn luyện.
data_collator (DataCollator, tùy chọn) — Data collator để gom mẫu từ dataset thành batch. Mặc định là experimental.utils.DataCollatorForChatML sử dụng processing_class.
train_dataset (Dataset, tùy chọn) — Dataset để huấn luyện.
eval_dataset (Dataset hoặc dict của Dataset, tuỳ chọn) — Dataset để đánh giá.
processing_class (PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin hoặc ProcessorMixin, tùy chọn) — Lớp để xử lý dữ liệu.
compute_metrics (Callable, tùy chọn) — Hàm để tính các chỉ số đánh giá. Phải nhận vào một EvalPrediction và trả về một từ điển từ chuỗi sang số float.
callbacks (list của TrainerCallback, tùy chọn) — Các callback được sử dụng trong quá trình huấn luyện.
optimizers (tuple của torch.optim.Optimizer và torch.optim.lr_scheduler.LambdaLR, tuỳ chọn, mặc định là (None, None)) — tuple chứa optimizer và learning rate scheduler sẽ được sử dụng cho quá trình huấn luyện.
preprocess_logits_for_metrics (Callable, tùy chọn) — Hàm để tiền xử lý logits trước khi tính toán các chỉ số. Phải nhận vào logits và labels và trả về logits sẽ được sử dụng để tính toán chỉ số.
peft_config (PeftConfig, tùy chọn) — Cấu hình PEFT để sử dụng PEFT cho việc huấn luyện. Nếu là None, PEFT sẽ không được sử dụng. Nếu được cung cấp, model sẽ được bọc với bộ thích ứng PEFT được chỉ định.
formatting_func (Callable, tùy chọn) — Hàm để định dạng bộ dữ liệu. Phải nhận vào một ví dụ và trả về một ví dụ.

Trainer cho Chưng cất Kiến thức Tổng quát (GKD) của các mô hình ngôn ngữ.

Để biết thêm chi tiết về GKD, hãy xem bài báo: On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes.

train

< source >

Tham số

resume_from_checkpoint (str hoặc bool, tùy chọn) — Nếu là str, đường dẫn cục bộ đến một checkpoint đã lưu bởi một phiên Trainer trước đó. Nếu là bool và bằng True, tải checkpoint cuối cùng trong args.output_dir do một phiên Trainer trước đó lưu lại. Nếu có, quá trình huấn luyện sẽ tiếp tục từ trạng thái model/optimizer/scheduler được tải ở đây.
trial (optuna.Trial hoặc dict[str, Any], tùy chọn) — Lần chạy thử nghiệm hoặc từ điển siêu tham số cho việc tìm kiếm siêu tham số.
ignore_keys_for_eval (list[str], tùy chọn) — Một danh sách các khóa trong đầu ra của mô hình của bạn (nếu đó là một từ điển) cần được bỏ qua khi thu thập các dự đoán để đánh giá trong quá trình huấn luyện.

Trả về

~trainer_utils.TrainOutput

Đối tượng chứa số bước toàn cục, tổn thất huấn luyện và các chỉ số.

Điểm vào chính cho quá trình huấn luyện.

save_model

< source >

( output_dir: str | None = None_internal_call: bool = False )

Sẽ lưu mô hình, để bạn có thể tải lại bằng from_pretrained().

Chỉ lưu từ tiến trình chính.

push_to_hub

< source >

( commit_message: str | None = 'End of training'blocking: bool = Truetoken: str | None = Nonerevision: str | None = None**kwargs )

Tham số

commit_message (str, tùy chọn, mặc định là "End of training") — Thông điệp để commit khi đẩy lên.
blocking (bool, tùy chọn, mặc định là True) — Xác định xem hàm có chỉ trả về khi git push đã hoàn tất hay không.
token (str, tùy chọn, mặc định là None) — Token có quyền ghi để ghi đè lên các args gốc của Trainer.
revision (str, tùy chọn) — Phiên bản git để commit từ. Mặc định là đầu của nhánh “main”.
kwargs (dict[str, Any], tùy chọn) — Các đối số từ khóa bổ sung được truyền đến ~Trainer.create_model_card.

Tải self.model và self.processing_class lên 🤗 model hub tại repo self.args.hub_model_id.

GKDConfig

lớp trl.experimental.gkd.GKDConfig

< source >

( output_dir: str | None = Noneper_device_train_batch_size: int = 8num_train_epochs: float = 3.0max_steps: int = -1learning_rate: float = 2e-05lr_scheduler_type: transformers.trainer_utils.SchedulerType | str = 'linear'lr_scheduler_kwargs: dict | str | None = Nonewarmup_steps: float = 0optim: transformers.training_args.OptimizerNames | str = 'adamw_torch_fused'optim_args: str | None = Noneweight_decay: float = 0.0adam_beta1: float = 0.9adam_beta2: float = 0.999adam_epsilon: float = 1e-08optim_target_modules: None | str | list[str] = Nonegradient_accumulation_steps: int = 1average_tokens_across_devices: bool = Truemax_grad_norm: float = 1.0label_smoothing_factor: float = 0.0bf16: bool | None = Nonefp16: bool = Falsebf16_full_eval: bool = Falsefp16_full_eval: bool = Falsetf32: bool | None = Nonegradient_checkpointing: bool = Truegradient_checkpointing_kwargs: dict[str, typing.Any] | str | None = Nonetorch_compile: bool = Falsetorch_compile_backend: str | None = Nonetorch_compile_mode: str | None = Noneuse_liger_kernel: bool = Falseliger_kernel_config: dict[str, bool] | None = Noneuse_cache: bool = Falseneftune_noise_alpha: float | None = Nonetorch_empty_cache_steps: int | None = Noneauto_find_batch_size: bool = Falselogging_strategy: transformers.trainer_utils.IntervalStrategy | str = 'steps'logging_steps: float = 10logging_first_step: bool = Falselog_on_each_node: bool = Truelogging_nan_inf_filter: bool = Trueinclude_num_input_tokens_seen: str | bool = 'no'log_level: str = 'passive'log_level_replica: str = 'warning'disable_tqdm: bool | None = Nonereport_to: None | str | list[str] = 'none'run_name: str | None = Noneproject: str = 'huggingface'trackio_space_id: str | None = Nonetrackio_bucket_id: str | None = Nonetrackio_static_space_id: typing.Union[str, NoneType, typing.Literal[False]] = Noneeval_strategy: transformers.trainer_utils.IntervalStrategy | str = 'no'eval_steps: float | None = Noneeval_delay: float = 0per_device_eval_batch_size: int = 8prediction_loss_only: bool = Falseeval_on_start: bool = Falseeval_do_concat_batches: bool = Trueeval_use_gather_object: bool = Falseeval_accumulation_steps: int | None = Noneinclude_for_metrics: list[str] = <factory>batch_eval_metrics: bool = Falsesave_only_model: bool = Falsesave_strategy: transformers.trainer_utils.SaveStrategy | str = 'steps'save_steps: float = 500save_on_each_node: bool = Falsesave_total_limit: int | None = Noneenable_jit_checkpoint: bool = Falsepush_to_hub: bool = Falsehub_token: str | None = Nonehub_private_repo: bool | None = Nonehub_model_id: str | None = Nonehub_strategy: transformers.trainer_utils.HubStrategy | str = 'every_save'hub_always_push: bool = Falsehub_revision: str | None = Noneload_best_model_at_end: bool = Falsemetric_for_best_model: str | None = Nonegreater_is_better: bool | None = Noneignore_data_skip: bool = Falserestore_callback_states_from_checkpoint: bool = Falsefull_determinism: bool = Falseseed: int = 42data_seed: int | None = Noneuse_cpu: bool = Falseaccelerator_config: dict | str | None = Noneparallelism_config: accelerate.parallelism_config.ParallelismConfig | None = Nonedataloader_drop_last: bool = Falsedataloader_num_workers: int = 0dataloader_pin_memory: bool = Truedataloader_persistent_workers: bool = Falsedataloader_prefetch_factor: int | None = Nonedataloader_multiprocessing_context: str | None = Nonedataloader_in_order: bool = Trueremove_unused_columns: bool = Truelabel_names: list[str] | None = Nonetrain_sampling_strategy: str = 'random'length_column_name: str = 'length'ddp_find_unused_parameters: bool | None = Noneddp_bucket_cap_mb: int | None = Noneddp_broadcast_buffers: bool | None = Noneddp_static_graph: bool | None = Noneddp_backend: str | None = Noneddp_timeout: int = 1800fsdp: str | None = Nonefsdp_config: dict[str, typing.Any] | str | None = Nonedeepspeed: dict | str | None = Nonedebug: str | list[transformers.debug_utils.DebugOption] = ''skip_memory_metrics: bool = Truedo_train: bool = Falsedo_eval: bool = Falsedo_predict: bool = Falseresume_from_checkpoint: str | None = Nonelocal_rank: int = -1model_init_kwargs: dict[str, typing.Any] | str | None = Nonetrust_remote_code: bool = Falserouter_aux_loss_coef: float | None = Nonechat_template_path: str | None = Nonedataset_text_field: str = 'text'dataset_kwargs: dict[str, typing.Any] | str | None = Nonedataset_num_proc: int | None = Noneeos_token: str | None = Nonemax_length: int | None = 1024truncation_mode: str = 'keep_start'shuffle_dataset: bool = Falsepacking: bool = Falsepacking_strategy: str = 'bfd'padding_free: bool = Falsepad_to_multiple_of: int | str | None = 'auto'eval_packing: bool | None = Nonecompletion_only_loss: bool | None = Noneassistant_only_loss: bool = Falseloss_type: str = 'nll'activation_offloading: bool = Falsepad_token: str | None = Nonetemperature: float = 0.9lmbda: float = 0.5beta: float = 0.5max_new_tokens: int = 128teacher_model_name_or_path: str | None = Noneteacher_model_init_kwargs: dict[str, typing.Any] | str | None = Nonedisable_dropout: bool = Trueseq_kd: bool = False )

Tham số

temperature (float, tùy chọn, mặc định là 0.9) — Nhiệt độ cho quá trình lấy mẫu. Nhiệt độ càng cao, các phần hoàn thành càng ngẫu nhiên.
lmbda (float, tùy chọn, mặc định là 0.5) — Tham số Lambda kiểm soát tỷ lệ dữ liệu của học sinh (tức là tỷ lệ các đầu ra do học sinh tạo ra theo chính sách).
beta (float, tuỳ chọn, mặc định là 0.5) — Hệ số nội suy giữa 0.0 và 1.0 của hàm mất mát Generalized Jensen-Shannon Divergence. Khi beta là 0.0, hàm mất mát là KL divergence. Khi beta là 1.0, hàm mất mát là Inverse KL Divergence.
max_new_tokens (int, tùy chọn, mặc định là 128) — Số lượng token tối đa được tạo ra cho mỗi lần hoàn thành.
teacher_model_name_or_path (str, tùy chọn) — Tên hoặc đường dẫn của mô hình teacher. Nếu là None, mô hình teacher sẽ giống với mô hình đang được huấn luyện.
teacher_model_init_kwargs (dict[str, Any], tùy chọn) — Các đối số từ khóa để truyền cho AutoModelForCausalLM.from_pretrained khi khởi tạo mô hình teacher từ một chuỗi.
disable_dropout (bool, tùy chọn, mặc định là True) — Có vô hiệu hóa dropout trong mô hình hay không.
seq_kd (bool, tùy chọn, mặc định là False) — Tham số Seq_kd kiểm soát việc có thực hiện Sequence-Level KD (có thể xem là FT giám sát trên đầu ra do teacher tạo ra) hay không.

Lớp cấu hình cho experimental.gkd.GKDTrainer.

Lớp này chỉ bao gồm các tham số đặc thù cho việc huấn luyện GKD. Để xem danh sách đầy đủ các đối số huấn luyện, hãy tham khảo tài liệu của TrainingArguments và SFTConfig.

Update on GitHub

TRL

Trainer Chưng cất Kiến thức Tổng quát

Tổng quan

Mẹo sử dụng

Kiểu dữ liệu tập dữ liệu dự kiến

GKDTrainer

class trl.experimental.gkd.GKDTrainer

train

save_model

push_to_hub

GKDConfig

lớp trl.experimental.gkd.GKDConfig