diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 3aa566e943f..2f4217b3582 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -105,8 +105,6 @@ title: Async Distillation - local: async_grpo_trainer title: Asynchronous GRPO - - local: bco_trainer - title: BCO - local: bema_for_reference_model title: BEMA for Reference Model - local: cpo_trainer diff --git a/docs/source/bco_trainer.md b/docs/source/bco_trainer.md deleted file mode 100644 index ddf3bcd1b79..00000000000 --- a/docs/source/bco_trainer.md +++ /dev/null @@ -1,111 +0,0 @@ -# BCO Trainer - -[![model badge](https://img.shields.io/badge/All_models-BCO-blue)](https://huggingface.co/models?other=bco,trl) - -TRL supports the Binary Classifier Optimization (BCO). -The [BCO](https://huggingface.co/papers/2404.04656) authors train a binary classifier whose logit serves as a reward so that the classifier maps {prompt, chosen completion} pairs to 1 and {prompt, rejected completion} pairs to 0. - -## Expected dataset type - -The [`experimental.bco.BCOTrainer`] requires an [unpaired preference dataset](dataset_formats#unpaired-preference). -The [`experimental.bco.BCOTrainer`] supports both [conversational](dataset_formats#conversational) and [standard](dataset_formats#standard) dataset formats. When provided with a conversational dataset, the trainer will automatically apply the chat template to the dataset. - -## Expected model format - -The BCO trainer expects a model of `AutoModelForCausalLM`. - -## Using the `BCOTrainer` - -At a high level we need to initialize the `BCOTrainer` with a `model` we wish to train and a reference `ref_model` which we will use to calculate the implicit rewards of the preferred and rejected response. - -The `beta` refers to the hyperparameter of the implicit reward, and the dataset contains the 3 entries listed above. Note that the `model` and `ref_model` need to have the same architecture (ie decoder only or encoder-decoder). - -```python -from datasets import load_dataset -from transformers import AutoModelForCausalLM, AutoTokenizer -from trl.experimental.bco import BCOConfig, BCOTrainer - -model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct") -model_ref = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct") -tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct") -train_dataset = load_dataset("trl-lib/ultrafeedback-gpt-3.5-turbo-helpfulness", split="train") - -training_args = BCOConfig( - beta=0.1, -) - -bco_trainer = BCOTrainer( - model, - model_ref, - args=training_args, - train_dataset=train_dataset, - processing_class=tokenizer, -) -``` - -After this one can then call: - -```python -bco_trainer.train() -``` - -## Underlying Distribution matching (UDM) - -In practical scenarios, the thumbs-up and thumbs-down datasets are likely to have divergent underlying distributions of prompts. -Consider an LLM deployed for user feedback: if the model excels in writing tasks but underperforms in coding, the thumbs-up dataset will be dominated by writing-related prompts, while the thumbs-down dataset will contain mostly coding-related prompts. -If the prompts in your desired and undesired datasets differ a lot, it is useful to enable UDM. - -Choose an embedding model and tokenizer: - -```python -embedding_model = AutoModel.from_pretrained(your_model_id) -embedding_tokenizer = AutoTokenizer.from_pretrained(your_model_id) - -# customize this function depending on your embedding model -def embed_prompt(input_ids, attention_mask, model): - outputs = model(input_ids=input_ids, attention_mask=attention_mask) - return outputs.last_hidden_state.mean(dim=1) - -embedding_model = Accelerator().prepare_model(embedding_model) -embedding_func = partial(embed_prompt, model=embedding_model) -``` - -Set `prompt_sample_size` to define how many prompts are selected to train the UDM classifier and start the training with the provided embedding function: - -```python -training_args = BCOConfig( - beta=0.1, - prompt_sample_size=512, -) - -bco_trainer = BCOTrainer( - model, - model_ref, - args=training_args, - train_dataset=train_dataset, - processing_class=tokenizer, - embedding_func=embedding_func, - embedding_tokenizer=embedding_tokenizer, -) - -bco_trainer.train() -``` - -### For Mixture of Experts Models: Enabling the auxiliary loss - -MOEs are the most efficient if the load is about equally distributed between experts. -To ensure that we train MOEs similarly during preference-tuning, it is beneficial to add the auxiliary loss from the load balancer to the final loss. - -This option is enabled by setting `output_router_logits=True` in the model config (e.g. MixtralConfig). -To scale how much the auxiliary loss contributes to the total loss, use the hyperparameter `router_aux_loss_coef=...` (default: 0.001). - -## BCOTrainer - -[[autodoc]] experimental.bco.BCOTrainer - - train - - save_model - - push_to_hub - -## BCOConfig - -[[autodoc]] experimental.bco.BCOConfig diff --git a/docs/source/dataset_formats.md b/docs/source/dataset_formats.md index 00f3e8cda18..b3239612e5d 100644 --- a/docs/source/dataset_formats.md +++ b/docs/source/dataset_formats.md @@ -412,7 +412,6 @@ Choosing the right dataset type depends on the task you are working on and the s | [`RewardTrainer`] | [Preference (implicit prompt recommended)](#preference) | | [`RLOOTrainer`] | [Prompt-only](#prompt-only) | | [`SFTTrainer`] | [Language modeling](#language-modeling) or [Prompt-completion](#prompt-completion) | -| [`experimental.bco.BCOTrainer`] | [Unpaired preference](#unpaired-preference) or [Preference (explicit prompt recommended)](#preference) | | [`experimental.cpo.CPOTrainer`] | [Preference (explicit prompt recommended)](#preference) | | [`experimental.gkd.GKDTrainer`] | [Prompt-completion](#prompt-completion) | | [`experimental.nash_md.NashMDTrainer`] | [Prompt-only](#prompt-only) | diff --git a/docs/source/dpo_trainer.md b/docs/source/dpo_trainer.md index 219ec1d0e0e..ecf06960b72 100644 --- a/docs/source/dpo_trainer.md +++ b/docs/source/dpo_trainer.md @@ -114,7 +114,7 @@ Several formulations of the objective have been proposed in the literature. Init | `"exo_pair"` | The [EXO](https://huggingface.co/papers/2402.00856) authors propose reverse-KL preference optimization. `label_smoothing` must be strictly greater than `0.0`; a recommended value is `1e-3` (see Eq. 16 for the simplified pairwise variant). The full method uses `K>2` SFT completions and approaches PPO as `K` grows. | | `"nca_pair"` | The [NCA](https://huggingface.co/papers/2402.05369) authors shows that NCA optimizes the absolute likelihood for each response rather than the relative likelihood. | | `"robust"` | The [Robust DPO](https://huggingface.co/papers/2403.00409) authors propose an unbiased DPO loss under noisy preferences. Use `label_smoothing` in [`DPOConfig`] to model label-flip probability; valid values are in the range `[0.0, 0.5)`. | -| `"bco_pair"` | The [BCO](https://huggingface.co/papers/2404.04656) authors train a binary classifier whose logit serves as a reward so that the classifier maps {prompt, chosen completion} pairs to 1 and {prompt, rejected completion} pairs to 0. For unpaired data, we recommend the dedicated [`experimental.bco.BCOTrainer`]. | +| `"bco_pair"` | The [BCO](https://huggingface.co/papers/2404.04656) authors train a binary classifier whose logit serves as a reward so that the classifier maps {prompt, chosen completion} pairs to 1 and {prompt, rejected completion} pairs to 0. | | `"sppo_hard"` | The [SPPO](https://huggingface.co/papers/2405.00675) authors claim that SPPO is capable of solving the Nash equilibrium iteratively by pushing the chosen rewards to be as large as 1/2 and the rejected rewards to be as small as -1/2 and can alleviate data sparsity issues. The implementation approximates this algorithm by employing hard label probabilities, assigning 1 to the winner and 0 to the loser. | | `"aot"` or `loss_type="aot_unpaired"` | The [AOT](https://huggingface.co/papers/2406.05882) authors propose Distributional Preference Alignment via Optimal Transport. `loss_type="aot"` is for paired data; `loss_type="aot_unpaired"` is for unpaired data. Both enforce stochastic dominance via sorted quantiles; larger per-GPU batch sizes help. | | `"apo_zero"` or `loss_type="apo_down"` | The [APO](https://huggingface.co/papers/2408.06266) method introduces an anchored objective. `apo_zero` boosts winners and downweights losers (useful when the model underperforms the winners). `apo_down` downweights both, with stronger pressure on losers (useful when the model already outperforms winners). | diff --git a/docs/source/index.md b/docs/source/index.md index 16e57f10f6e..4ea49df6286 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -71,7 +71,6 @@ Below is an overview of TRL trainers, organized by maturity and method type. #### Offline methods - [BEMA for Reference Model](bema_for_reference_model) -- [`BCOTrainer`](bco_trainer) - [`CPOTrainer`](cpo_trainer) - [`ORPOTrainer`](orpo_trainer) - [`TPOTrainer`](tpo_trainer) diff --git a/docs/source/paper_index.md b/docs/source/paper_index.md index edd2b8e1d6e..993e31fa00c 100644 --- a/docs/source/paper_index.md +++ b/docs/source/paper_index.md @@ -954,8 +954,6 @@ training_args = DPOConfig( ) ``` -For the unpaired version, the user should utilize [`experimental.bco.BCOConfig`] and [`experimental.bco.BCOTrainer`]. - ### Learn Your Reference Model for Real Good Alignment **📜 Paper**: https://huggingface.co/papers/2404.09656 diff --git a/pyproject.toml b/pyproject.toml index ee7a263ff3e..45ebbb45268 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,10 +44,6 @@ Homepage = "https://github.com/huggingface/trl" trl = "trl.cli:main" [project.optional-dependencies] -bco = [ - "scikit-learn", - "joblib" -] deepspeed = [ "deepspeed>=0.18.6", "transformers!=5.1.0", # see transformers#43780 @@ -78,9 +74,6 @@ quality = [ quantization = [ "bitsandbytes" ] -scikit = [ - "scikit-learn" -] test = [ "pytest-cov", "pytest-datadir>=1.7.0", # lazy datadirs @@ -101,9 +94,6 @@ vlm = [ "num2words==0.5.14" ] dev = [ - # bco - "scikit-learn", - "joblib", # deepspeed "deepspeed>=0.18.6", # kernels: transformers renamed the "hub-kernels" extra to "kernels" in 5.1.0 @@ -120,7 +110,6 @@ dev = [ "hf-doc-builder", # quantization "bitsandbytes", - # scikit: included in bco # test "pytest-cov", "pytest-datadir>=1.7.0", # lazy datadirs diff --git a/tests/experimental/test_bco_trainer.py b/tests/experimental/test_bco_trainer.py deleted file mode 100644 index 6b39722f6fd..00000000000 --- a/tests/experimental/test_bco_trainer.py +++ /dev/null @@ -1,535 +0,0 @@ -# Copyright 2020-2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from functools import partial - -import pytest -import torch -from accelerate import Accelerator -from datasets import DatasetDict, load_dataset -from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer -from transformers.utils import is_peft_available - -from trl.experimental.bco import BCOConfig, BCOTrainer -from trl.experimental.bco.bco_trainer import _process_tokens, _tokenize - -from ..testing_utils import TrlTestCase, require_no_wandb, require_peft, require_sklearn - - -if is_peft_available(): - from peft import LoraConfig - - -@pytest.mark.low_priority -class TestBCOTrainer(TrlTestCase): - @require_sklearn - def test_trust_remote_code(self): - dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train") - model_id = "trl-internal-testing/tiny-RemoteForCausalLM" - tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) - - with pytest.raises(ValueError, match="custom code"): - BCOTrainer( - model=model_id, - args=BCOConfig(output_dir=self.tmp_dir, report_to="none"), - processing_class=tokenizer, - train_dataset=dataset, - ) - - trainer = BCOTrainer( - model=model_id, - args=BCOConfig(output_dir=self.tmp_dir, report_to="none", trust_remote_code=True), - processing_class=tokenizer, - train_dataset=dataset, - ) - assert type(trainer.model).__name__ == "RemoteForCausalLM" - - @pytest.mark.parametrize( - "config_name", - [ - "standard_preference", - "standard_implicit_prompt_preference", - "standard_unpaired_preference", - "conversational_preference", - "conversational_implicit_prompt_preference", - "conversational_unpaired_preference", - ], - ) - @require_sklearn - def test_train(self, config_name): - model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5" - model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32") - ref_model = AutoModelForCausalLM.from_pretrained(model_id) - tokenizer = AutoTokenizer.from_pretrained(model_id) - - dataset = load_dataset("trl-internal-testing/zen", config_name, split="train") - - training_args = BCOConfig( - output_dir=self.tmp_dir, - remove_unused_columns=False, # warning raised if not set to False - learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates - report_to="none", - ) - - trainer = BCOTrainer( - model=model, - ref_model=ref_model, - args=training_args, - processing_class=tokenizer, - train_dataset=dataset, - ) - - previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()} - - trainer.train() - - assert trainer.state.log_history[-1]["train_loss"] is not None - - # Check that the params have changed - for n, param in previous_trainable_params.items(): - new_param = trainer.model.get_parameter(n) - if param.sum() != 0: # ignore 0 biases - assert not torch.equal(param.cpu(), new_param.cpu()) - - @pytest.mark.parametrize( - "eval_dataset_type", - [ - "dataset", - "dataset_dict", - "none", - ], - ) - @require_sklearn - def test_init_with_eval_dataset(self, eval_dataset_type): - # BCO tokenizes the eval dataset at init by calling `.map()` directly on it, so streaming (iterable) and - # plain dict-of-datasets eval datasets are not yet supported; only `Dataset` and `DatasetDict` are tested. - model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5" - model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32") - ref_model = AutoModelForCausalLM.from_pretrained(model_id) - tokenizer = AutoTokenizer.from_pretrained(model_id) - - train_dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train") - - if eval_dataset_type == "none": - eval_dataset = None - else: - eval_split = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="test") - if eval_dataset_type == "dataset": - eval_dataset = eval_split - else: # "dataset_dict" - eval_dataset = DatasetDict({"data1": eval_split, "data2": eval_split}) - - training_args = BCOConfig(output_dir=self.tmp_dir, remove_unused_columns=False, report_to="none") - trainer = BCOTrainer( - model=model, - ref_model=ref_model, - args=training_args, - processing_class=tokenizer, - train_dataset=train_dataset, - eval_dataset=eval_dataset, - ) - - if eval_dataset_type == "none": - assert trainer.eval_dataset is None - elif isinstance(trainer.eval_dataset, dict): - assert set(trainer.eval_dataset.keys()) == {"data1", "data2"} - # Each split was tokenized independently. - assert "prompt_input_ids" in next(iter(trainer.eval_dataset["data1"])) - assert "prompt_input_ids" in next(iter(trainer.eval_dataset["data2"])) - else: - assert "prompt_input_ids" in next(iter(trainer.eval_dataset)) - - @require_sklearn - def test_train_processing_class_autoloaded(self): - # processing_class is documented as optional: when omitted it should be - # auto-loaded from the model, consistent with DPOTrainer / RewardTrainer. - model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5" - model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32") - ref_model = AutoModelForCausalLM.from_pretrained(model_id) - dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train") - training_args = BCOConfig( - output_dir=self.tmp_dir, - remove_unused_columns=False, - learning_rate=0.1, - report_to="none", - ) - trainer = BCOTrainer( - model=model, - ref_model=ref_model, - args=training_args, - train_dataset=dataset, - ) - assert trainer.processing_class is not None - trainer.train() - assert trainer.state.log_history[-1]["train_loss"] is not None - - @require_sklearn - def test_train_with_precompute(self): - model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5" - model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32") - ref_model = AutoModelForCausalLM.from_pretrained(model_id) - tokenizer = AutoTokenizer.from_pretrained(model_id) - - dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train") - - training_args = BCOConfig( - output_dir=self.tmp_dir, - remove_unused_columns=False, # warning raised if not set to False - learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates - precompute_ref_log_probs=True, - report_to="none", - ) - - trainer = BCOTrainer( - model=model, - ref_model=ref_model, - args=training_args, - processing_class=tokenizer, - train_dataset=dataset, - ) - - previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()} - - trainer.train() - - assert trainer.state.log_history[-1]["train_loss"] is not None - - # Check that the params have changed - for n, param in previous_trainable_params.items(): - new_param = trainer.model.get_parameter(n) - if param.sum() != 0: # ignore 0 biases - assert not torch.equal(param.cpu(), new_param.cpu()) - - @require_sklearn - def test_train_eval(self): - model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5" - model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32") - ref_model = AutoModelForCausalLM.from_pretrained(model_id) - tokenizer = AutoTokenizer.from_pretrained(model_id) - - dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference") - - training_args = BCOConfig( - output_dir=self.tmp_dir, - remove_unused_columns=False, # warning raised if not set to False - eval_strategy="steps", - eval_steps=3, - report_to="none", - ) - - trainer = BCOTrainer( - model=model, - ref_model=ref_model, - args=training_args, - processing_class=tokenizer, - train_dataset=dataset["train"], - eval_dataset=dataset["test"], - ) - - trainer.train() - - @require_sklearn - def test_init_with_ref_model_is_model(self): - model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5" - model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32") - tokenizer = AutoTokenizer.from_pretrained(model_id) - - dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train") - - training_args = BCOConfig( - output_dir=self.tmp_dir, - remove_unused_columns=False, # warning raised if not set to False - report_to="none", - ) - - with pytest.raises(ValueError): - BCOTrainer( - model=model, - ref_model=model, # ref_model can't be the same as model - args=training_args, - processing_class=tokenizer, - train_dataset=dataset, - ) - - @require_sklearn - def test_tokenize_and_process_tokens(self): - model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5" - model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32") - ref_model = AutoModelForCausalLM.from_pretrained(model_id) - tokenizer = AutoTokenizer.from_pretrained(model_id) - - dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train") - - training_args = BCOConfig( - output_dir=self.tmp_dir, - remove_unused_columns=False, # warning raised if not set to False - report_to="none", - ) - - trainer = BCOTrainer( - model=model, - ref_model=ref_model, - args=training_args, - processing_class=tokenizer, - train_dataset=dataset, - ) - - tokenized_dataset = dataset.map( - _tokenize, - fn_kwargs={"tokenizer": trainer.processing_class}, - batched=True, - batch_size=2, - ) - assert tokenized_dataset["prompt"][:] == dataset["prompt"][:] - assert tokenized_dataset["completion"][:] == dataset["completion"][:] - assert tokenized_dataset["label"][:] == dataset["label"][:] - assert tokenized_dataset["prompt_input_ids"][0] == [46518, 374, 2664, 1091] - assert tokenized_dataset["prompt_attention_mask"][0] == [1, 1, 1, 1] - assert tokenized_dataset["answer_input_ids"][0] == [27261, 13] - assert tokenized_dataset["answer_attention_mask"][0] == [1, 1] - - fn_kwargs = { - "prefix": "", - "is_encoder_decoder": trainer.is_encoder_decoder, - "tokenizer": trainer.processing_class, - "max_length": trainer.max_length, - } - processed_dataset = tokenized_dataset.map(_process_tokens, fn_kwargs=fn_kwargs) - assert processed_dataset["prompt"][:] == dataset["prompt"][:] - assert processed_dataset["completion"][:] == dataset["completion"][:] - assert processed_dataset["label"][:] == dataset["label"][:] - assert processed_dataset["prompt_input_ids"][0] == [46518, 374, 2664, 1091] - assert processed_dataset["prompt_attention_mask"][0] == [1, 1, 1, 1] - assert processed_dataset["completion_input_ids"][0] == [46518, 374, 2664, 1091, 27261, 13, 151645] - assert processed_dataset["completion_attention_mask"][0] == [1, 1, 1, 1, 1, 1, 1] - assert processed_dataset["completion_labels"][0] == [-100, -100, -100, -100, 27261, 13, 151645] - - @require_sklearn - def test_train_without_providing_ref_model(self): - model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5" - model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32") - tokenizer = AutoTokenizer.from_pretrained(model_id) - - dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train") - - training_args = BCOConfig( - output_dir=self.tmp_dir, - remove_unused_columns=False, # warning raised if not set to False - learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates - report_to="none", - ) - - trainer = BCOTrainer( - model=model, - args=training_args, - processing_class=tokenizer, - train_dataset=dataset, - ) - - previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()} - - trainer.train() - - assert trainer.state.log_history[-1]["train_loss"] is not None - - # Check that the params have changed - for n, param in previous_trainable_params.items(): - new_param = trainer.model.get_parameter(n) - if param.sum() != 0: # ignore 0 biases - assert not torch.equal(param.cpu(), new_param.cpu()) - - @require_sklearn - def test_train_udm(self): - model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5" - model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32") - tokenizer = AutoTokenizer.from_pretrained(model_id) - - # Get embedding model - embedding_model_id = "trl-internal-testing/tiny-BartModel" - embedding_model = AutoModel.from_pretrained(embedding_model_id) - embedding_tokenizer = AutoTokenizer.from_pretrained(embedding_model_id) - - def embed_prompt(input_ids, attention_mask, model): - outputs = model(input_ids=input_ids, attention_mask=attention_mask) - - return outputs.last_hidden_state.mean(dim=1) - - embedding_model = Accelerator().prepare_model(embedding_model) - embedding_func = partial(embed_prompt, model=embedding_model) - - dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train") - - training_args = BCOConfig( - output_dir=self.tmp_dir, - remove_unused_columns=False, # warning raised if not set to False - learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates - report_to="none", - ) - - trainer = BCOTrainer( - model=model, - args=training_args, - processing_class=tokenizer, - train_dataset=dataset, - embedding_func=embedding_func, - embedding_tokenizer=embedding_tokenizer, - ) - - previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()} - - trainer.train() - - assert trainer.state.log_history[-1]["train_loss"] is not None - - # Check that the params have changed - for n, param in previous_trainable_params.items(): - new_param = trainer.model.get_parameter(n) - if param.sum() != 0: # ignore 0 biases - assert not torch.equal(param.cpu(), new_param.cpu()) - - @require_sklearn - @require_peft - def test_train_without_providing_ref_model_with_lora(self): - model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5" - model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32") - lora_config = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, task_type="CAUSAL_LM") - tokenizer = AutoTokenizer.from_pretrained(model_id) - - dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train") - - training_args = BCOConfig( - output_dir=self.tmp_dir, - remove_unused_columns=False, # warning raised if not set to False - learning_rate=0.1, # use higher lr because gradients are tiny and default lr can stall updates - report_to="none", - ) - - trainer = BCOTrainer( - model=model, - args=training_args, - processing_class=tokenizer, - train_dataset=dataset, - peft_config=lora_config, - ) - - previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()} - - trainer.train() - - assert trainer.state.log_history[-1]["train_loss"] is not None - - # Check that the params have changed - for n, param in previous_trainable_params.items(): - if "lora" in n: - new_param = trainer.model.get_parameter(n) - if param.sum() != 0: # ignore 0 biases - assert not torch.equal(param.cpu(), new_param.cpu()) - - @require_sklearn - @require_no_wandb - def test_generate_during_eval_no_wandb(self): - model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5" - model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32") - tokenizer = AutoTokenizer.from_pretrained(model_id) - - dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference") - - training_args = BCOConfig( - output_dir=self.tmp_dir, - remove_unused_columns=False, # warning raised if not set to False - eval_strategy="steps", - eval_steps=3, - generate_during_eval=True, - report_to="none", - ) - - with pytest.raises( - ValueError, - match="`generate_during_eval=True` requires Weights and Biases or Comet to be installed." - " Please install `wandb` or `comet-ml` to resolve.", - ): - BCOTrainer( - model=model, - args=training_args, - processing_class=tokenizer, - train_dataset=dataset["train"], - eval_dataset=dataset["test"], - ) - - @require_sklearn - @require_peft - def test_lora_train_and_save(self): - model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5" - model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32") - lora_config = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, task_type="CAUSAL_LM") - tokenizer = AutoTokenizer.from_pretrained(model_id) - - dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference", split="train") - - training_args = BCOConfig( - output_dir=self.tmp_dir, - remove_unused_columns=False, # warning raised if not set to False - report_to="none", - ) - - trainer = BCOTrainer( - model=model, - args=training_args, - processing_class=tokenizer, - train_dataset=dataset, - peft_config=lora_config, - ) - - # train the model - trainer.train() - - # save peft adapter - trainer.save_model() - - # assert that the model is loaded without giving OSError - AutoModelForCausalLM.from_pretrained(self.tmp_dir) - - @require_sklearn - def test_compute_metrics(self): - model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5" - model = AutoModelForCausalLM.from_pretrained(model_id, dtype="float32") - ref_model = AutoModelForCausalLM.from_pretrained(model_id) - tokenizer = AutoTokenizer.from_pretrained(model_id) - - dataset = load_dataset("trl-internal-testing/zen", "standard_unpaired_preference") - - def dummy_compute_metrics(*args, **kwargs): - return {"test": 0.0} - - training_args = BCOConfig( - output_dir=self.tmp_dir, - remove_unused_columns=False, # warning raised if not set to False - eval_strategy="steps", - eval_steps=3, - report_to="none", - ) - - trainer = BCOTrainer( - model=model, - ref_model=ref_model, - args=training_args, - processing_class=tokenizer, - train_dataset=dataset["train"], - eval_dataset=dataset["test"], - compute_metrics=dummy_compute_metrics, - ) - - trainer.train() - - assert trainer.state.log_history[-2]["eval_test"] == 0.0 diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 97800a34c90..a9394698311 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -23,7 +23,7 @@ import torch import torch.nn as nn from packaging.version import Version -from transformers import is_bitsandbytes_available, is_comet_available, is_sklearn_available, is_wandb_available +from transformers import is_bitsandbytes_available, is_comet_available, is_wandb_available from transformers.testing_utils import backend_device_count, torch_device from transformers.utils import ( is_kernels_available, @@ -39,7 +39,6 @@ from trl.import_utils import ( is_harbor_available, is_jmespath_available, - is_joblib_available, is_liger_kernel_available, is_math_verify_available, is_mergekit_available, @@ -73,9 +72,6 @@ reason="test requires jmespath for response parsing on transformers below 5.13.0", ) require_rich = pytest.mark.skipif(not is_rich_available(), reason="test requires rich") -require_sklearn = pytest.mark.skipif( - not (is_sklearn_available() and is_joblib_available()), reason="test requires sklearn" -) require_torch_accelerator = pytest.mark.skipif( torch_device is None or torch_device == "cpu", reason="test requires accelerator" ) diff --git a/trl/experimental/bco/__init__.py b/trl/experimental/bco/__init__.py deleted file mode 100644 index 3a6871ca60e..00000000000 --- a/trl/experimental/bco/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# Copyright 2020-2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from .bco_config import BCOConfig -from .bco_trainer import BCOTrainer diff --git a/trl/experimental/bco/bco_config.py b/trl/experimental/bco/bco_config.py deleted file mode 100644 index 79c54ab63ed..00000000000 --- a/trl/experimental/bco/bco_config.py +++ /dev/null @@ -1,163 +0,0 @@ -# Copyright 2020-2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from dataclasses import dataclass, field -from typing import Any - -from ...trainer.base_config import _BaseConfig - - -@dataclass -class BCOConfig(_BaseConfig): - r""" - Configuration class for the [`experimental.bco.BCOTrainer`]. - - This class includes only the parameters that are specific to BCO training. For a full list of training arguments, - please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may - differ from those in [`~transformers.TrainingArguments`]. - - Using [`~transformers.HfArgumentParser`] we can turn this class into - [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the - command line. - - Parameters: - max_length (`int` or `None`, *optional*, defaults to `1024`): - Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want - to use the default data collator. - max_completion_length (`int`, *optional*): - Maximum length of the completion. This argument is required if you want to use the default data collator - and your model is an encoder-decoder. - beta (`float`, *optional*, defaults to `0.1`): - Parameter controlling the deviation from the reference model. Higher β means less deviation from the - reference model. - disable_dropout (`bool`, *optional*, defaults to `True`): - Whether to disable dropout in the model and reference model. - generate_during_eval (`bool`, *optional*, defaults to `False`): - If `True`, generates and logs completions from both the model and the reference model to W&B or Comet - during evaluation. - is_encoder_decoder (`bool`, *optional*): - When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument, - you need to specify if the model returned by the callable is an encoder-decoder model. - precompute_ref_log_probs (`bool`, *optional*, defaults to `False`): - Whether to precompute reference model log probabilities for training and evaluation datasets. This is - useful when training without the reference model to reduce the total GPU memory needed. - model_init_kwargs (`dict[str, Any]`, *optional*): - Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model and - reference model from strings. The `revision` value is also used when loading the processing class. - trust_remote_code (`bool`, *optional*, defaults to `False`): - Whether to allow loading models that ship custom Python code from the Hub. Forwarded to - [`~transformers.AutoModelForCausalLM.from_pretrained`] for both the model and reference model. - dataset_num_proc (`int`, *optional*): - Number of processes to use for processing the dataset. - prompt_sample_size (`int`, *optional*, defaults to `1024`): - Number of prompts that are fed to density ratio classifier. - min_density_ratio (`float`, *optional*, defaults to `0.5`): - Minimum value of the density ratio. The estimated density ratio is clamped to this value. - max_density_ratio (`float`, *optional*, defaults to `10.0`): - Maximum value of the density ratio. The estimated density ratio is clamped to this value. - - > [!NOTE] - > These parameters have default values different from [`~transformers.TrainingArguments`]: - > - `logging_steps`: Defaults to `10` instead of `500`. - > - `gradient_checkpointing`: Defaults to `True` instead of `False`. - > - `bf16`: Defaults to `True` if `fp16` is not set, instead of `False`. - > - `learning_rate`: Defaults to `5e-7` instead of `5e-5`. - """ - - _VALID_DICT_FIELDS = _BaseConfig._VALID_DICT_FIELDS + ["model_init_kwargs"] - - # Parameters whose default values are overridden from TrainingArguments - learning_rate: float = field( - default=5e-7, - metadata={"help": "The initial learning rate for AdamW."}, - ) - - max_length: int | None = field( - default=1024, - metadata={ - "help": "Maximum length of the sequences (prompt + completion) in the batch. " - "This argument is required if you want to use the default data collator." - }, - ) - max_completion_length: int | None = field( - default=None, - metadata={ - "help": "Maximum length of the completion. This argument is required if you want to use the " - "default data collator and your model is an encoder-decoder." - }, - ) - beta: float = field( - default=0.1, - metadata={ - "help": "Parameter controlling the deviation from the reference model. " - "Higher β means less deviation from the reference model." - }, - ) - disable_dropout: bool = field( - default=True, - metadata={"help": "Whether to disable dropout in the model and reference model."}, - ) - generate_during_eval: bool = field( - default=False, - metadata={ - "help": "If `True`, generates and logs completions from both the model and the reference model " - "to W&B during evaluation." - }, - ) - is_encoder_decoder: bool | None = field( - default=None, - metadata={ - "help": "When using the `model_init` argument (callable) to instantiate the model instead of the " - "`model` argument, you need to specify if the model returned by the callable is an " - "encoder-decoder model." - }, - ) - precompute_ref_log_probs: bool = field( - default=False, - metadata={ - "help": "Whether to precompute reference model log probabilities for training and evaluation datasets. " - "This is useful when training without the reference model to reduce the total GPU memory " - "needed." - }, - ) - model_init_kwargs: dict[str, Any] | str | None = field( - default=None, - metadata={ - "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the " - "model from a string. The `revision` value is also used when loading the processing class." - }, - ) - trust_remote_code: bool = field( - default=False, - metadata={ - "help": "Whether to allow loading models that ship custom Python code from the Hub. Forwarded to " - "`AutoModelForCausalLM.from_pretrained` for both the model and reference model." - }, - ) - dataset_num_proc: int | None = field( - default=None, - metadata={"help": "Number of processes to use for processing the dataset."}, - ) - prompt_sample_size: int = field( - default=1024, - metadata={"help": "Number of prompts that are fed to density ratio classifier."}, - ) - min_density_ratio: float = field( - default=0.5, - metadata={"help": "Minimum value of the density ratio. The estimated density ratio is clamped to this value."}, - ) - max_density_ratio: float = field( - default=10.0, - metadata={"help": "Maximum value of the density ratio. The estimated density ratio is clamped to this value."}, - ) diff --git a/trl/experimental/bco/bco_trainer.py b/trl/experimental/bco/bco_trainer.py deleted file mode 100644 index 61f4b44023c..00000000000 --- a/trl/experimental/bco/bco_trainer.py +++ /dev/null @@ -1,1597 +0,0 @@ -# Copyright 2020-2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import dataclasses -import inspect -import json -import os -import random -import textwrap -from collections import defaultdict -from collections.abc import Callable -from contextlib import contextmanager, nullcontext -from dataclasses import dataclass -from operator import itemgetter -from pathlib import Path -from typing import Any, Literal - -import numpy as np -import pandas as pd -import torch -import torch.nn as nn -import torch.nn.functional as F -import transformers -from accelerate import Accelerator, PartialState -from accelerate.logging import get_logger -from accelerate.utils import is_peft_model, tqdm -from datasets import Dataset -from packaging.version import Version -from torch import autocast -from torch.utils.data import DataLoader, SequentialSampler -from transformers import ( - AutoModelForCausalLM, - AutoTokenizer, - BaseImageProcessor, - DataCollator, - FeatureExtractionMixin, - PreTrainedModel, - PreTrainedTokenizerBase, - ProcessorMixin, - TrainerCallback, - TrainingArguments, - is_comet_available, - is_sklearn_available, - is_wandb_available, -) -from transformers.trainer_utils import EvalLoopOutput, has_length -from transformers.utils import is_peft_available - -from ...data_utils import maybe_apply_chat_template, maybe_extract_prompt, maybe_unpair_preference_dataset -from ...import_utils import is_joblib_available -from ...models.utils import prepare_deepspeed -from ...trainer.base_trainer import _BaseTrainer -from ...trainer.utils import ( - disable_dropout_in_model, - get_config_model_id, - log_table_to_comet_experiment, - selective_log_softmax, -) -from ..utils import DPODataCollatorWithPadding, create_reference_model, pad_to_length, peft_module_casting_to_bf16 -from .bco_config import BCOConfig - - -if is_joblib_available(): - import joblib - - -if is_peft_available(): - from peft import PeftConfig, get_peft_model, prepare_model_for_kbit_training - - -if is_sklearn_available(): - from sklearn.linear_model import LogisticRegression - - -if is_wandb_available(): - import wandb - - -logger = get_logger(__name__) - -RUNNING_NAME = "running.json" -CLF_NAME = "clf.pkl" - - -@torch.no_grad() -def get_global_statistics( - accelerator, xs: torch.Tensor, mask=None, device="cpu" -) -> tuple[torch.Tensor, torch.Tensor, int]: - """ - Computes element-wise mean and variance of the tensor across processes. Reference: - https://github.com/OpenLMLab/MOSS-RLHF/blob/40b91eb2f2b71b16919addede0341d2bef70825d/utils.py#L57C1-L73C75 - """ - xs = xs.to(accelerator.device) - sum_and_count = torch.tensor([xs.sum(), (xs.numel() if mask is None else mask.sum())], device=xs.device) - sum_and_count = accelerator.reduce(sum_and_count) - global_sum, count = sum_and_count - global_mean = global_sum / count - - sum_var = torch.sum(((xs - global_mean) ** 2).mul(1 if mask is None else mask)) - sum_var = accelerator.reduce(sum_var) - global_var = sum_var / count - - return global_mean.to(device), global_var.to(device), count.item() - - -@dataclass -class RunningMoments: - """ - Calculates the running mean and standard deviation of a data stream. Reference: - https://github.com/OpenLMLab/MOSS-RLHF/blob/40b91eb2f2b71b16919addede0341d2bef70825d/utils.py#L75 - """ - - accelerator: Accelerator - mean: float = 0 - std: float = 1 - var: float = 1 - count: float = 1e-24 - - @torch.no_grad() - def update(self, xs: torch.Tensor) -> tuple[float, float]: - """ - Updates running moments from batch's moments computed across ranks - """ - if self.accelerator.use_distributed: - xs_mean, xs_var, xs_count = get_global_statistics(self.accelerator, xs) - else: - xs_count = xs.numel() - xs_var, xs_mean = torch.var_mean(xs, unbiased=False) - xs_mean, xs_var = xs_mean.float(), xs_var.float() - - delta = xs_mean - self.mean - tot_count = self.count + xs_count - - new_sum = xs_var * xs_count - # correct old_sum deviation accounting for the new mean - old_sum = self.var * self.count + delta**2 * self.count * xs_count / tot_count - tot_sum = old_sum + new_sum - - self.mean += (delta * xs_count / tot_count).item() - new_var = tot_sum / tot_count - self.std = (new_var * tot_count / (tot_count - 1)).float().sqrt().item() - self.var = new_var.item() - self.count = tot_count - - return xs_mean.item(), (xs_var * xs_count / (xs_count - 1)).float().sqrt().item() - - def save_to_json(self, json_path: str): - """Save the content of this instance in JSON format inside `json_path`.""" - # save everything except accelerator - if self.accelerator.is_main_process: - save_dict = dataclasses.asdict(self, dict_factory=lambda x: {k: v for (k, v) in x if k != "accelerator"}) - json_string = json.dumps(save_dict, indent=2, sort_keys=True) + "\n" - with open(json_path, "w", encoding="utf-8") as f: - f.write(json_string) - - @classmethod - def load_from_json(cls, accelerator: Accelerator, json_path: str): - """Create an instance from the content of `json_path`.""" - # load everything except accelerator - with open(json_path, encoding="utf-8") as f: - text = f.read() - return cls(accelerator=accelerator, **json.loads(text)) - - -def _tokenize( - batch: dict[str, list[Any]], - tokenizer: PreTrainedTokenizerBase, - embedding_tokenizer: PreTrainedTokenizerBase | None = None, -) -> dict[str, list[Any]]: - """Tokenize a batch from a BCO specific dataset.""" - prompt_tokenized = tokenizer(batch["prompt"], add_special_tokens=False) - prompt_input_ids = prompt_tokenized["input_ids"] - prompt_attention_mask = prompt_tokenized["attention_mask"] - prompt_and_completion = [ - prompt + completion for prompt, completion in zip(batch["prompt"], batch["completion"], strict=True) - ] - full_tokenized = tokenizer(prompt_and_completion, add_special_tokens=False) - full_input_ids = full_tokenized["input_ids"] - full_attention_mask = full_tokenized["attention_mask"] - - answer_input_ids = [f[len(p) :] for f, p in zip(full_input_ids, prompt_input_ids, strict=True)] - answer_attention_mask = [f[len(p) :] for f, p in zip(full_attention_mask, prompt_attention_mask, strict=True)] - - # Concat tokens to form `enc(a) + enc(a + b)[len(enc(a)):]` - full_concat_input_ids = [np.concatenate([p, a]) for p, a in zip(prompt_input_ids, answer_input_ids, strict=True)] - # Prepare input tokens for token by token comparison - full_input_ids = [np.array(f) for f in full_input_ids] - for full, concat in zip(full_input_ids, full_concat_input_ids, strict=True): - if len(full) != len(concat): - raise ValueError( - "The elements in 'full_input_ids' and 'full_concat_input_ids' must have the same pairwise length." - ) - - # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens - # can be merged together when tokenizing prompt+answer. This could result - # on the last token from the prompt being different when tokenized on its own - # vs when done as prompt+answer. - response_token_ids_start_idx = [len(p) for p in prompt_input_ids] - - # If tokenized prompt is different than both prompt+answer, then it means the - # last token has changed due to merging. - for idx, (p, f, r) in enumerate(zip(prompt_input_ids, full_input_ids, response_token_ids_start_idx, strict=True)): - if not np.array_equal(p, f[:r]): - response_token_ids_start_idx[idx] -= 1 - - prompt_input_ids = [f[:r] for f, r in zip(full_input_ids, response_token_ids_start_idx, strict=True)] - prompt_attention_mask = [f[:r] for f, r in zip(full_attention_mask, response_token_ids_start_idx, strict=True)] - - for p, m in zip(prompt_input_ids, prompt_attention_mask, strict=True): - if len(p) != len(m): - raise ValueError("Prompt input ids and attention mask should have the same length.") - - answer_input_ids = [f[r:] for f, r in zip(full_input_ids, response_token_ids_start_idx, strict=True)] - answer_attention_mask = [f[r:] for f, r in zip(full_attention_mask, response_token_ids_start_idx, strict=True)] - - output = dict( - prompt_input_ids=prompt_input_ids, - prompt_attention_mask=prompt_attention_mask, - answer_input_ids=answer_input_ids, - answer_attention_mask=answer_attention_mask, - ) - - if embedding_tokenizer is not None: - embedding_tokenized = embedding_tokenizer(batch["prompt"], add_special_tokens=False) - - output.update( - { - "embedding_input_ids": embedding_tokenized["input_ids"], - "embedding_attention_mask": embedding_tokenized["attention_mask"], - } - ) - - return output - - -def _process_tokens(example: dict[str, Any], model: "PreTrainedModel" = None, **kwargs) -> dict: - """Process tokens of a BCO specific dataset. - - At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation in case the prompt + - completion responses is/are too long. First we truncate the prompt; if we're still too long, we truncate the - completion. - - We also create the labels for the completion responses, which are of length equal to the sum of the length of the - prompt and the completion response, with `-100` for the prompt tokens. - """ - prompt = example["prompt"] - completion = example["completion"] - - batch = { - f"{kwargs['prefix']}prompt": prompt, - f"{kwargs['prefix']}completion": completion, - f"{kwargs['prefix']}label": example["label"], - } - - if not kwargs["is_encoder_decoder"]: - # Check issues below for more details - # 1. https://github.com/huggingface/trl/issues/907 - # 2. https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 - # 3. https://github.com/LianjiaTech/BELLE/issues/337 - - if not isinstance(prompt, str): - raise ValueError(f"prompt should be an str but got {type(prompt)}") - - if not isinstance(completion, str): - raise ValueError(f"completion should be an str but got {type(completion)}") - - # keys of format prompt_* refers to just the prompt and answer_* refers to just the answer - all_tokens = { - "prompt_input_ids": example["prompt_input_ids"], - "prompt_attention_mask": example["prompt_attention_mask"], - "answer_input_ids": example["answer_input_ids"], - "answer_attention_mask": example["answer_attention_mask"], - } - - # calculate max length by checking if BOS/EOS is already there - max_length = kwargs["max_length"] - bos_token_id = kwargs["tokenizer"].bos_token_id - eos_token_id = kwargs["tokenizer"].eos_token_id - if bos_token_id != all_tokens["prompt_input_ids"][0]: - max_length -= 1 - if eos_token_id != all_tokens["answer_input_ids"][-1]: - max_length -= 1 - - # if combined sequence is too long (> max_length - 1 for BOS token - 1 for EOS), truncate the response - if len(all_tokens["prompt_input_ids"]) + len(all_tokens["answer_input_ids"]) > max_length: - for k in ["answer_input_ids", "answer_attention_mask"]: - all_tokens[k] = all_tokens[k][: max_length - len(all_tokens["prompt_input_ids"])] - - # all input_ids and attention mask as is. We then check if we need to add BOS/EOS tokens - batch[f"{kwargs['prefix']}prompt_input_ids"] = all_tokens["prompt_input_ids"] - batch[f"{kwargs['prefix']}prompt_attention_mask"] = all_tokens["prompt_attention_mask"] - batch[f"{kwargs['prefix']}completion_input_ids"] = ( - all_tokens["prompt_input_ids"] + all_tokens["answer_input_ids"] - ) - batch[f"{kwargs['prefix']}completion_attention_mask"] = ( - all_tokens["prompt_attention_mask"] + all_tokens["answer_attention_mask"] - ) - - # add BOS, which affects both prompt and the full completion - if bos_token_id is not None: - if len(all_tokens["prompt_input_ids"]) == 0 or bos_token_id != all_tokens["prompt_input_ids"][0]: - batch[f"{kwargs['prefix']}prompt_input_ids"] = [bos_token_id] + batch[ - f"{kwargs['prefix']}prompt_input_ids" - ] - batch[f"{kwargs['prefix']}prompt_attention_mask"] = [1] + batch[ - f"{kwargs['prefix']}prompt_attention_mask" - ] - batch[f"{kwargs['prefix']}completion_input_ids"] = [bos_token_id] + batch[ - f"{kwargs['prefix']}completion_input_ids" - ] - batch[f"{kwargs['prefix']}completion_attention_mask"] = [1] + batch[ - f"{kwargs['prefix']}completion_attention_mask" - ] - # add EOS, which affects only the full completion - if len(all_tokens["answer_input_ids"]) == 0 or eos_token_id != all_tokens["answer_input_ids"][-1]: - batch[f"{kwargs['prefix']}completion_input_ids"] = batch[f"{kwargs['prefix']}completion_input_ids"] + [ - eos_token_id - ] - batch[f"{kwargs['prefix']}completion_attention_mask"] = batch[ - f"{kwargs['prefix']}completion_attention_mask" - ] + [1] - - batch[f"{kwargs['prefix']}completion_labels"] = batch[f"{kwargs['prefix']}completion_input_ids"][:] - batch[f"{kwargs['prefix']}completion_labels"][: len(batch[f"{kwargs['prefix']}prompt_input_ids"])] = [ - -100 - ] * len(batch[f"{kwargs['prefix']}prompt_input_ids"]) - else: - completion_tokens = kwargs["tokenizer"]( - completion, truncation=True, max_length=kwargs["max_completion_length"], add_special_tokens=True - ) - prompt_tokens = kwargs["tokenizer"](prompt, add_special_tokens=True) - - batch[f"{kwargs['prefix']}prompt_input_ids"] = prompt_tokens["input_ids"] - batch[f"{kwargs['prefix']}prompt_attention_mask"] = prompt_tokens["attention_mask"] - - batch[f"{kwargs['prefix']}completion_labels"] = completion_tokens["input_ids"] - batch[f"{kwargs['prefix']}completion_attention_mask"] = completion_tokens["attention_mask"] - if model is not None and hasattr(model, "prepare_decoder_input_ids_from_labels"): - batch[f"{kwargs['prefix']}completion_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( - labels=torch.tensor(batch["completion_labels"]) - ) - - return batch - - -class BCOTrainer(_BaseTrainer): - r""" - Initialize BCOTrainer from [BCO](https://huggingface.co/papers/2404.04656) paper. - - Args: - model ([`~transformers.PreTrainedModel`]): - The model to train, preferably an [`~transformers.AutoModelForSequenceClassification`]. - ref_model ([`~transformers.PreTrainedModel`]): - Hugging Face transformer model with a casual language modelling head. Used for implicit reward computation - and loss. If no reference model is provided, the trainer will create a reference model with the same - architecture as the model to be optimized. - args ([`experimental.bco.BCOConfig`]): - The arguments to use for training. - train_dataset ([`~datasets.Dataset`]): - The dataset to use for training. - eval_dataset ([`~datasets.Dataset`]): - The dataset to use for evaluation. - processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.BaseImageProcessor`], [`~transformers.FeatureExtractionMixin`] or [`~transformers.ProcessorMixin`], *optional*): - Processing class used to process the data. If provided, will be used to automatically process the inputs - for the model, and it will be saved along the model to make it easier to rerun an interrupted training or - reuse the fine-tuned model. - data_collator ([`~transformers.DataCollator`], *optional*): - The data collator to use for training. If None is specified, the default data collator - ([`experimental.utils.DPODataCollatorWithPadding`]) will be used which will pad the sequences to the - maximum length of the sequences in the batch, given a dataset of paired sequences. - model_init (`Callable[[], transformers.PreTrainedModel]`): - The model initializer to use for training. If None is specified, the default model initializer will be - used. - callbacks (`list[transformers.TrainerCallback]`): - The callbacks to use for training. - optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): - The optimizer and scheduler to use for training. - preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): - The function to use to preprocess the logits before computing the metrics. - peft_config ([`~peft.PeftConfig`], *optional*): - The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in - a PEFT model. - compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): - The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to - metric values. - model_adapter_name (`str`, defaults to `None`): - Name of the train target PEFT adapter, when using LoRA with multiple adapters. - ref_adapter_name (`str`, defaults to `None`): - Name of the reference PEFT adapter, when using LoRA with multiple adapters. - embedding_func (`Callable`, *optional*): - Function to compute prompt embeddings, used to train the underlying distribution matching (UDM) classifier - when the desirable and undesirable datasets have divergent prompt distributions. Requires the scikit-learn - and joblib libraries. - embedding_tokenizer ([`~transformers.PreTrainedTokenizerBase`], *optional*): - Tokenizer used to prepare prompts for `embedding_func`. - """ - - _tag_names = ["trl", "bco"] - _name = "BCO" - _paper = { - "title": "Binary Classifier Optimization for Large Language Model Alignment", - "id": "2404.04656", - # docstyle-ignore - "citation": textwrap.dedent("""\ - @article{jung2024binary, - title = {{Binary Classifier Optimization for Large Language Model Alignment}}, - author = {Seungjae Jung and Gunsoo Han and Daniel Wontae Nam and Kyoung{-}Woon On}, - year = 2024, - eprint = {arXiv:2404.04656} - }"""), - } - - def __init__( - self, - model: PreTrainedModel | nn.Module | str = None, - ref_model: PreTrainedModel | nn.Module | str | None = None, - args: BCOConfig = None, - train_dataset: Dataset | None = None, - eval_dataset: Dataset | dict[str, Dataset] | None = None, - processing_class: PreTrainedTokenizerBase - | BaseImageProcessor - | FeatureExtractionMixin - | ProcessorMixin - | None = None, - data_collator: DataCollator | None = None, - model_init: Callable[[], PreTrainedModel] | None = None, - callbacks: list[TrainerCallback] | None = None, - optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), - preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, - peft_config: "PeftConfig | None" = None, - compute_metrics: Callable[[EvalLoopOutput], dict] | None = None, - model_adapter_name: str | None = None, - ref_adapter_name: str | None = None, - embedding_func: Callable | None = None, - embedding_tokenizer: PreTrainedTokenizerBase | None = None, - ): - if embedding_func is not None and not (is_sklearn_available() and is_joblib_available()): - raise ImportError( - "BCOTrainer with UDM requires the scikit-learn and joblib libraries. Please install it with `pip install scikit-learn joblib`." - ) - - if type(args) is TrainingArguments: - raise ValueError("Please use `BCOConfig` instead `TrainingArguments`.") - - if train_dataset is None: - raise ValueError("`train_dataset` is required") - - if not isinstance(model, str) and model is not None and ref_model is model: - raise ValueError( - "`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the " - "same as `model`, you must mass a copy of it, or `None` if you use peft." - ) - - if args.model_init_kwargs is None: - model_init_kwargs = {} - elif not isinstance(model, str): - raise ValueError("You passed model_kwargs to the BCOTrainer. But your model is already instantiated.") - else: - model_init_kwargs = args.model_init_kwargs - dtype = model_init_kwargs.get("dtype", "auto") - if dtype is not None: - # Convert to `torch.dtype` if an str is passed - if isinstance(dtype, str) and dtype != "auto": - dtype = getattr(torch, dtype) - if dtype != "auto" and not isinstance(dtype, torch.dtype): - raise ValueError( - f"Invalid `dtype` passed to the BCOConfig. Expected a string with either `torch.dtype` or 'auto', but got {dtype}." - ) - model_init_kwargs["dtype"] = dtype - model_init_kwargs["device_map"] = model_init_kwargs.get("device_map", "auto") - - model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code) - model_revision = model_init_kwargs.get("revision") if isinstance(model, str) else None - - if isinstance(model, str): - model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs) - - if isinstance(ref_model, str): - ref_model = AutoModelForCausalLM.from_pretrained(ref_model, **model_init_kwargs) - - # PEFT - # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16` - # has been called in order to properly call autocast if needed. - self._peft_has_been_casted_to_bf16 = False - if peft_config is not None: - if not is_peft_available(): - raise ImportError( - "You passed `peft_config` but the `peft` library is not installed. " - "Install it with `pip install trl[peft]`." - ) - if not isinstance(peft_config, PeftConfig): - raise TypeError( - f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " - f"got {type(peft_config).__name__}." - ) - if is_peft_model(model): - raise ValueError( - "You passed a `PeftModel` instance together with a `peft_config` to the trainer. Please first " - "merge and unload the existing adapter, save the resulting base model, and then pass that base " - "model along with the new `peft_config` to the trainer." - ) - - if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False): - _support_gc_kwargs = hasattr( - args, "gradient_checkpointing_kwargs" - ) and "gradient_checkpointing_kwargs" in list( - inspect.signature(prepare_model_for_kbit_training).parameters - ) - - prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing} - - if _support_gc_kwargs: - prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs - - model = prepare_model_for_kbit_training(model, **prepare_model_kwargs) - elif args.gradient_checkpointing: - # For backward compatibility with older versions of transformers - if hasattr(model, "enable_input_require_grads"): - model.enable_input_require_grads() - else: - - def make_inputs_require_grad(module, input, output): - output.requires_grad_(True) - - model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) - - # get peft model with the given config - # ZeRO-3 + PEFT for non-quantized models: - # - PEFT's default autocast_adapter_dtype=True upcasts LoRA adapter params to fp32 even when the base model is bf16. - # - ZeRO-3's _allgather_params_coalesced allocates output buffers using the dtype of the first persistent parameter, - # so mixed-dtype persistent_parameters (bf16 base + fp32 LoRA) cause a TypeError on the first optimizer step. - # - Passing autocast_adapter_dtype=False keeps adapter params in the base model dtype (bf16), fixing the mismatch. - # - This is safe: the fp32 upcast is a QLoRA-specific concern (low-bit quantized base models), not needed for - # non-quantized bf16 training. - # - See: - # - TRL issue: https://github.com/huggingface/trl/issues/6089 - # - Upstream issue: https://github.com/deepspeedai/DeepSpeed/issues/8072 - _is_quantized_model = getattr(model, "is_loaded_in_4bit", False) or getattr( - model, "is_loaded_in_8bit", False - ) - get_peft_model_kwargs = {} - if args.deepspeed_plugin is not None and args.deepspeed_plugin.zero_stage == 3 and not _is_quantized_model: - get_peft_model_kwargs["autocast_adapter_dtype"] = False - model = get_peft_model(model, peft_config, **get_peft_model_kwargs) - if args.bf16 and getattr(model, "is_loaded_in_4bit", False): - peft_module_casting_to_bf16(model) - # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager - self._peft_has_been_casted_to_bf16 = True - - # For models that use gradient_checkpointing, we need to attach a hook that enables input - # to explicitly have `requires_grad=True`, otherwise training will either silently - # fail or completely fail. - elif args.gradient_checkpointing: - # For backward compatibility with older versions of transformers - if hasattr(model, "enable_input_require_grads"): - model.enable_input_require_grads() - else: - - def make_inputs_require_grad(module, input, output): - output.requires_grad_(True) - - model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) - - if args.generate_during_eval and not (is_wandb_available() or is_comet_available()): - raise ValueError( - "`generate_during_eval=True` requires Weights and Biases or Comet to be installed." - " Please install `wandb` or `comet-ml` to resolve." - ) - - if model is not None: - self.is_encoder_decoder = model.config.is_encoder_decoder - elif args.is_encoder_decoder is None: - raise ValueError("When no model is provided, you need to pass the parameter is_encoder_decoder.") - else: - self.is_encoder_decoder = args.is_encoder_decoder - - self.model_adapter_name = model_adapter_name - self.ref_adapter_name = ref_adapter_name - - if ref_model: - self.ref_model = ref_model - elif is_peft_model(model) or args.precompute_ref_log_probs: - # The `model` with adapters turned off will be used as the reference model - self.ref_model = None - else: - self.ref_model = create_reference_model(model) - - if processing_class is None: - processing_class = AutoTokenizer.from_pretrained( - get_config_model_id(model.config), revision=model_revision, trust_remote_code=args.trust_remote_code - ) - if args.max_length is None: - logger.warning( - "When using DPODataCollatorWithPadding, you should set `max_length` in the `BCOConfig`. " - "It will be set to `512` by default, but you should do it yourself in the future.", - ) - max_length = 512 - if args.max_length is not None: - max_length = args.max_length - - max_completion_length = None - if args.max_completion_length is None and self.is_encoder_decoder: - logger.warning( - "When using DPODataCollatorWithPadding with an encoder decoder architecture, you should set `max_completion_length` in the BCOTrainer's init" - " it will be set to `128` by default, but you should do it yourself in the future.", - ) - max_completion_length = 128 - if args.max_completion_length is not None and self.is_encoder_decoder: - max_completion_length = args.max_completion_length - - if data_collator is None: - data_collator = DPODataCollatorWithPadding( - pad_token_id=processing_class.pad_token_id, - is_encoder_decoder=self.is_encoder_decoder, - ) - - if args.remove_unused_columns: - args.remove_unused_columns = False - # warn users - logger.warning( - "When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your BCOConfig" - " we have set it for you, but you should do it yourself in the future.", - ) - - self.use_dpo_data_collator = True - else: - self.use_dpo_data_collator = False - - # Disable dropout in the model and reference model - if args.disable_dropout: - disable_dropout_in_model(model) - if self.ref_model is not None: - disable_dropout_in_model(self.ref_model) - - self.max_length = max_length - self.generate_during_eval = args.generate_during_eval - self.max_completion_length = max_completion_length - self.precompute_ref_log_probs = args.precompute_ref_log_probs - - # Since ref_logs are precomputed on the first call to get_train/eval_dataloader - # keep track of first called to avoid computation of future calls - self._precomputed_train_ref_log_probs = False - self._precomputed_eval_ref_log_probs = False - - # metric - self._stored_metrics = defaultdict(lambda: defaultdict(list)) - - # BCO parameter - self.beta = args.beta - self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) - self.aux_loss_coef = getattr(model.config, "router_aux_loss_coef", 0.0) - if self.aux_loss_enabled and self.aux_loss_coef == 0.0: - logger.warning( - "You set `output_router_logits` to `True` in the model config, but `router_aux_loss_coef` is set to " - "`0.0`, meaning the auxiliary loss will not be used. Either set `router_aux_loss_coef` to a value " - "greater than `0.0`, or set `output_router_logits` to `False` if you don't want to use the auxiliary " - "loss.", - ) - - # Underlying Distribution Matching argument - self.embedding_func = embedding_func - self.embedding_tokenizer = embedding_tokenizer - - with PartialState().main_process_first(): - # Extract the prompt if needed - train_dataset = train_dataset.map( - maybe_extract_prompt, num_proc=args.dataset_num_proc, desc="Extracting prompt from train dataset" - ) - # Unpair the dataset if needed - train_dataset = maybe_unpair_preference_dataset( - train_dataset, args.dataset_num_proc, desc="Unpairing train dataset" - ) - # Apply the chat template if needed - train_dataset = train_dataset.map( - maybe_apply_chat_template, - fn_kwargs={"processing_class": processing_class}, - num_proc=args.dataset_num_proc, - ) - if eval_dataset is not None: - # Extract the prompt if needed - eval_dataset = eval_dataset.map( - maybe_extract_prompt, num_proc=args.dataset_num_proc, desc="Extracting prompt from eval dataset" - ) - # Unpair the dataset if needed - eval_dataset = maybe_unpair_preference_dataset( - eval_dataset, args.dataset_num_proc, desc="Unpairing eval dataset" - ) - eval_dataset = eval_dataset.map( - maybe_apply_chat_template, - fn_kwargs={"processing_class": processing_class}, - num_proc=args.dataset_num_proc, - ) - - # Tokenize and prepare the training datasets - train_dataset = train_dataset.map( - _tokenize, - batched=True, - fn_kwargs={"tokenizer": processing_class, "embedding_tokenizer": self.embedding_tokenizer}, - num_proc=args.dataset_num_proc, - desc="Tokenizing train dataset", - ) - - # Prepare the datasets - fn_kwargs = { - "prefix": "", - "is_encoder_decoder": self.is_encoder_decoder, - "tokenizer": processing_class, - "max_length": self.max_length, - "max_completion_length": self.max_completion_length, - } - train_dataset = train_dataset.map( - _process_tokens, - fn_kwargs=fn_kwargs, - num_proc=args.dataset_num_proc, - desc="Processing tokenized train dataset", - ) - - if eval_dataset is not None: - # Tokenize - eval_dataset = eval_dataset.map( - _tokenize, - fn_kwargs={"tokenizer": processing_class, "embedding_tokenizer": self.embedding_tokenizer}, - batched=True, - num_proc=args.dataset_num_proc, - desc="Tokenizing eval dataset", - ) - - # Process - fn_kwargs = { - "prefix": "", - "is_encoder_decoder": self.is_encoder_decoder, - "tokenizer": processing_class, - "max_length": self.max_length, - "max_completion_length": self.max_completion_length, - } - eval_dataset = eval_dataset.map( - _process_tokens, - fn_kwargs=fn_kwargs, - num_proc=args.dataset_num_proc, - desc="Processing tokenized eval dataset", - ) - - desirable = train_dataset.filter( - lambda x: x["label"], num_proc=args.dataset_num_proc, desc="Filtering desirable examples" - ) - undesirable = train_dataset.filter( - lambda x: not x["label"], num_proc=args.dataset_num_proc, desc="Filtering undesirable examples" - ) - - # Transformers explicitly set use_reentrant=True in the past to silence a PyTorch warning, but the default was - # never updated once PyTorch switched to recommending use_reentrant=False. Until that change lands upstream - # (see https://github.com/huggingface/transformers/pull/43203) and is released (most likely in 5.0.0), we - # default to the recommended non-reentrant behavior here, while preserving any user-provided value. - if args.gradient_checkpointing and Version(transformers.__version__) < Version("5.0.0"): - args.gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} - args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False) - - super().__init__( - model=model, - args=args, - data_collator=data_collator, - train_dataset=train_dataset, - eval_dataset=eval_dataset, - processing_class=processing_class, - model_init=model_init, - compute_metrics=compute_metrics, - callbacks=callbacks, - optimizers=optimizers, - preprocess_logits_for_metrics=preprocess_logits_for_metrics, - ) - - # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the - # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set - # self.model_accepts_loss_kwargs to False to enable scaling. - self.model_accepts_loss_kwargs = False - - # Add tags for models that have been loaded with the correct transformers version - if hasattr(self.model, "add_model_tags"): - self.model.add_model_tags(self._tag_names) - - if not hasattr(self, "accelerator"): - raise AttributeError( - "Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`." - ) - - # Deepspeed Zero-3 does not support precompute_ref_log_probs - if self.is_deepspeed_enabled: - if self.accelerator.state.deepspeed_plugin.zero_stage == 3 and self.precompute_ref_log_probs: - raise ValueError( - "You cannot use `precompute_ref_log_probs=True` with Deepspeed ZeRO-3. Please set `precompute_ref_log_probs=False`." - ) - - if self.ref_model is None: - if not (is_peft_model(model) or self.precompute_ref_log_probs): - raise ValueError( - "No reference model and model is not a Peft model. Try setting `precompute_ref_log_probs=True`" - ) - else: - if self.is_deepspeed_enabled: - self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) - else: - self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) - - self.running = RunningMoments(accelerator=self.accelerator) - - if self.embedding_func is None or args.resume_from_checkpoint: - return - - chosen_embeddings = self._get_sample_prompt_embeddings(desirable, sample_size=self.args.prompt_sample_size) - rejected_embeddings = self._get_sample_prompt_embeddings(undesirable, sample_size=self.args.prompt_sample_size) - - embeddings = torch.cat((chosen_embeddings, rejected_embeddings), dim=0) - labels = torch.cat( - (torch.ones_like(chosen_embeddings[:, 0]), torch.zeros_like(rejected_embeddings[:, 0])), dim=0 - ) - - self.clf = LogisticRegression(class_weight="balanced").fit( - embeddings.cpu().float().numpy(), labels.cpu().numpy() - ) - chosen_mean = self.clf.score( - chosen_embeddings.cpu().float().numpy(), torch.ones_like(chosen_embeddings[:, 0]).cpu().numpy() - ) - rejected_mean = self.clf.score( - rejected_embeddings.cpu().float().numpy(), torch.zeros_like(rejected_embeddings[:, 0]).cpu().numpy() - ) - logger.info(f"UDM classifier training scores: chosen: {chosen_mean}, rejected: {rejected_mean}") - - @property - def match_underlying_distribution(self): - return self.embedding_func is not None and self.embedding_tokenizer is not None - - def _get_chosen_prob(self, prompt_embeddings: torch.FloatTensor) -> torch.FloatTensor: - """ - Calculates the probability if the given prompt embedding is from desirable dataset. This function calculates - the probability in the process and ensemble across processes. - """ - dtype = prompt_embeddings.dtype - device = prompt_embeddings.device - rank = self.accelerator.process_index - - padded_prompt_embeddings = self.accelerator.pad_across_processes( - prompt_embeddings, pad_index=self.embedding_tokenizer.pad_token_id - ) - sample_size = padded_prompt_embeddings.shape[0] - nonzero = padded_prompt_embeddings.mean(dim=1) != self.embedding_tokenizer.pad_token_id - prompt_embeddings = self.accelerator.gather(padded_prompt_embeddings) - - # cannot predict for all empty values - if prompt_embeddings.shape[0] == 0: - return torch.tensor([], device=device, dtype=dtype) - - prob = self.clf.predict_proba(prompt_embeddings.cpu().float().numpy())[:, 1] - prob = torch.as_tensor(prob, dtype=dtype, device=device) - prob = self.accelerator.reduce(prob, reduction="mean") - - prob = prob[sample_size * rank : sample_size * (rank + 1)] - prob = prob[nonzero] - - return prob - - def _vectorize_prompt(self, input_ids: torch.LongTensor, attention_mask: torch.LongTensor) -> torch.FloatTensor: - """ - Replaces processing_class.pad_token_id to embedding_tokenizer.pad_token_id and applies self.embedding_func - """ - input_ids = torch.where( - input_ids == self.processing_class.pad_token_id, - self.embedding_tokenizer.pad_token_id, - input_ids, - ) - - with torch.no_grad(): - embeddings = self.embedding_func( - input_ids=input_ids, - attention_mask=attention_mask, - ) - - return embeddings - - def _get_prompt_embeddings( - self, batch: dict[str, list | torch.LongTensor] - ) -> tuple[torch.FloatTensor, torch.FloatTensor]: - """Extract embeddings from frozen embedding model""" - - if not self.match_underlying_distribution: - return None, None - - embeddings = self._vectorize_prompt( - input_ids=batch["embedding_input_ids"], - attention_mask=batch["embedding_attention_mask"], - ) - - labels = torch.tensor(batch["label"], dtype=torch.bool, device=embeddings.device) - chosen_idx = torch.where(labels)[0] - rejected_idx = torch.where(~labels)[0] - - chosen_embeddings = embeddings[chosen_idx, ...] - rejected_embeddings = embeddings[rejected_idx, ...] - - return (chosen_embeddings, rejected_embeddings) - - def _get_sample_prompt_embeddings(self, dataset: Dataset, sample_size: int = 512) -> torch.FloatTensor: - """ - Sample instances from dataset and get prompt embeddings. Used for density ratio classifier training. - """ - n_samples = min(len(dataset), sample_size) - rand_indices = np.random.choice(len(dataset), size=(n_samples,)) - - embedding_dataset = dataset.select(rand_indices) - - dataloader_params = { - "batch_size": self.args.per_device_train_batch_size, - "collate_fn": self.data_collator, - "num_workers": self.args.dataloader_num_workers, - "pin_memory": self.args.dataloader_pin_memory, - "shuffle": False, - } - - # prepare dataloader - data_loader = self.accelerator.prepare(DataLoader(embedding_dataset, **dataloader_params)) - - with torch.no_grad(): - all_embeddings = torch.empty(0) - for padded_batch in tqdm( - iterable=data_loader, - desc="Building sample prompt embeddings", - disable=bool(os.environ.get("TQDM_DISABLE", "")), - ): - embeddings = self._vectorize_prompt( - input_ids=padded_batch["embedding_input_ids"], - attention_mask=padded_batch["embedding_attention_mask"], - ) - embeddings = self.accelerator.gather_for_metrics(embeddings) - all_embeddings = torch.cat((all_embeddings, embeddings.cpu())) - - return all_embeddings - - def _save_optimizer_and_scheduler(self, output_dir): - output_dir = output_dir if output_dir is not None else self.args.output_dir - super()._save_optimizer_and_scheduler(output_dir) - - if self.accelerator.is_main_process: - # When saving optimizer and scheduler to checkpoint, save also the running delta object. - self.running.save_to_json(os.path.join(output_dir, RUNNING_NAME)) - - if self.match_underlying_distribution: - joblib.dump(self.clf, os.path.join(output_dir, CLF_NAME), compress=True) - - def _load_optimizer_and_scheduler(self, checkpoint): - if checkpoint is None: - logger.warning_once(f"Missing Checkpoint {checkpoint}") - return - - super()._load_optimizer_and_scheduler(checkpoint) - - # when loading optimizer and scheduler from checkpoint, also load the running delta object. - running_file = os.path.join(checkpoint, RUNNING_NAME) - if os.path.isfile(running_file): - self.running = RunningMoments.load_from_json(self.accelerator, running_file) - - if self.match_underlying_distribution: - clf_file = os.path.join(checkpoint, CLF_NAME) - if os.path.isfile(clf_file): - self.clf = joblib.load(clf_file) - - @contextmanager - def null_ref_context(self): - """Context manager for handling null reference model (that is, peft adapter manipulation).""" - with ( - self.accelerator.unwrap_model(self.model).disable_adapter() - if is_peft_model(self.model) and not self.ref_adapter_name - else nullcontext() - ): - if self.ref_adapter_name: - self.model.set_adapter(self.ref_adapter_name) - yield - if self.ref_adapter_name: - self.model.set_adapter(self.model_adapter_name or "default") - - def get_train_dataloader(self) -> DataLoader: - """ - Returns the training [`~torch.utils.data.DataLoader`]. - - Subclass of transformers.src.transformers.trainer.get_train_dataloader to precompute `ref_log_probs`. - """ - - if self.precompute_ref_log_probs and not self._precomputed_train_ref_log_probs: - dataloader_params = { - "batch_size": self.args.per_device_train_batch_size, - "collate_fn": self.data_collator, - "num_workers": self.args.dataloader_num_workers, - "pin_memory": self.args.dataloader_pin_memory, - "shuffle": False, - } - - # prepare dataloader - data_loader = self.accelerator.prepare(DataLoader(self.train_dataset, **dataloader_params)) - reference_completion_logps = [] - - for padded_batch in tqdm( - iterable=data_loader, - desc="Train dataset reference log probs", - disable=bool(os.environ.get("TQDM_DISABLE", "")), - ): - reference_completion_logp = self.compute_reference_log_probs(padded_batch) - - reference_completion_logp = self.accelerator.gather_for_metrics(reference_completion_logp) - reference_completion_logps.append(reference_completion_logp.cpu()) - - self.train_dataset = self.train_dataset.add_column( - name="reference_logps", column=torch.cat(reference_completion_logps).float().numpy() - ) - - self._precomputed_train_ref_log_probs = True - - return super().get_train_dataloader() - - def get_eval_dataloader(self, eval_dataset: Dataset | None = None) -> DataLoader: - """ - Returns the evaluation [`~torch.utils.data.DataLoader`]. - - Subclass of transformers.src.transformers.trainer.get_eval_dataloader to precompute `ref_log_probs`. - - Args: - eval_dataset (`torch.utils.data.Dataset`, *optional*): - If provided, will override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns not accepted - by the `model.forward()` method are automatically removed. It must implement `__len__`. - """ - if eval_dataset is None and self.eval_dataset is None: - raise ValueError("Trainer: evaluation requires an eval_dataset.") - eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset - - if self.precompute_ref_log_probs and not self._precomputed_eval_ref_log_probs: - dataloader_params = { - "batch_size": self.args.per_device_eval_batch_size, - "collate_fn": self.data_collator, - "num_workers": self.args.dataloader_num_workers, - "pin_memory": self.args.dataloader_pin_memory, - "shuffle": False, - } - - # prepare dataloader - data_loader = self.accelerator.prepare(DataLoader(eval_dataset, **dataloader_params)) - - reference_completion_logps = [] - - for padded_batch in tqdm( - iterable=data_loader, - desc="Eval dataset reference log probs", - disable=bool(os.environ.get("TQDM_DISABLE", "")), - ): - reference_completion_logp = self.compute_reference_log_probs(padded_batch) - - reference_completion_logp = self.accelerator.gather_for_metrics(reference_completion_logp) - reference_completion_logps.append(reference_completion_logp.cpu()) - - eval_dataset = eval_dataset.add_column( - name="reference_logps", column=torch.cat(reference_completion_logps).float().numpy() - ) - - # Save calculated reference_chosen_logps and reference_rejected_logps to the eval_dataset for subsequent runs - if self.eval_dataset is not None: - self.eval_dataset = eval_dataset - self._precomputed_eval_ref_log_probs = True - - return super().get_eval_dataloader(eval_dataset=eval_dataset) - - def compute_reference_log_probs(self, padded_batch: dict) -> dict: - """Computes log probabilities of the reference model for a single padded batch of a BCO specific dataset.""" - with torch.no_grad(): - if self.ref_model is None: - with self.null_ref_context(): - if self.is_encoder_decoder: - completion_logits = self.model( - padded_batch["prompt_input_ids"], - attention_mask=padded_batch["prompt_attention_mask"], - decoder_input_ids=padded_batch.get("completion_decoder_input_ids"), - labels=padded_batch["completion_labels"], - ).logits - - else: - completion_logits = self.model( - padded_batch["completion_input_ids"], - attention_mask=padded_batch["completion_attention_mask"], - ).logits - - else: - if self.is_encoder_decoder: - completion_logits = self.ref_model( - padded_batch["prompt_input_ids"], - attention_mask=padded_batch["prompt_attention_mask"], - decoder_input_ids=padded_batch.get("completion_decoder_input_ids"), - labels=padded_batch["completion_labels"], - ).logits - - else: - completion_logits = self.ref_model( - padded_batch["completion_input_ids"], attention_mask=padded_batch["completion_attention_mask"] - ).logits - - completion_logps = self.get_batch_logps( - completion_logits, - padded_batch["completion_labels"], - average_log_prob=False, - is_encoder_decoder=self.is_encoder_decoder, - ) - - return completion_logps - - @staticmethod - def get_batch_logps( - logits: torch.FloatTensor, - labels: torch.LongTensor, - average_log_prob: bool = False, - is_encoder_decoder: bool = False, - ) -> torch.FloatTensor: - """Compute the log probabilities of the given labels under the given logits. - - Args: - logits: Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size) - labels: - Labels for which to compute the log probabilities. Label tokens with a value of `-100` are ignored. - Shape: (batch_size, sequence_length) - average_log_prob: - If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the - log probabilities of the (non-masked) tokens. - is_encoder_decoder: - Whether the model is an encoder-decoder model. If True, the labels are not shifted, and the logits are - assumed to already be aligned with the labels. If False, the labels are shifted to the right by one - position, and the logits are assumed to be aligned with the shifted labels. - - Returns: - A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the - given logits. - """ - if logits.shape[:-1] != labels.shape: - raise ValueError("Logits (batch and sequence length dim) and labels must have the same shape.") - - if not is_encoder_decoder: - labels = labels[:, 1:].clone() - logits = logits[:, :-1, :] - else: - # Fixes end-dec RuntimeError - labels = labels.clone() - - loss_mask = labels != -100 - - # dummy token; we'll ignore the losses on these tokens later - labels[labels == -100] = 0 - - per_token_logps = selective_log_softmax(logits, labels) - - if average_log_prob: - return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1) - else: - return (per_token_logps * loss_mask).sum(-1) - - def forward( - self, model: nn.Module, batch: dict[str, list | torch.LongTensor] - ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: - model_kwargs = ( - { - "labels": batch["completion_labels"], - "decoder_input_ids": batch.get("completion_decoder_input_ids"), - } - if self.is_encoder_decoder - else {} - ) - if self.aux_loss_enabled: - model_kwargs["output_router_logits"] = True - - outputs = model( - batch["completion_input_ids"], - attention_mask=batch["completion_attention_mask"], - **model_kwargs, - ) - completion_logits = outputs.logits - - completion_logps = self.get_batch_logps( - completion_logits, - batch["completion_labels"], - average_log_prob=False, - is_encoder_decoder=self.is_encoder_decoder, - ) - - if completion_logps.shape[0] != len(batch["label"]): - raise ValueError( - "There is a mismatch between the number of examples in this batch and the number of " - "examples for which an output sequence was predicted." - ) - - chosen_idx = [i for i in range(completion_logps.shape[0]) if batch["label"][i] is True] - rejected_idx = [i for i in range(completion_logps.shape[0]) if batch["label"][i] is False] - - chosen_logps = completion_logps[chosen_idx, ...] - rejected_logps = completion_logps[rejected_idx, ...] - - chosen_logits = completion_logits[chosen_idx, ...] - rejected_logits = completion_logits[rejected_idx, ...] - - if self.aux_loss_enabled: - return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, outputs.aux_loss) - else: - return (chosen_logps, rejected_logps, chosen_logits, rejected_logits) - - def _get_udm_weight(self, rejected_embeddings: torch.FloatTensor) -> torch.FloatTensor: - prob_desirable = self._get_chosen_prob(rejected_embeddings) - min_ratio = self.args.min_density_ratio - max_ratio = self.args.max_density_ratio - - weight = (prob_desirable / (1 - prob_desirable + 1e-8)).clamp(min=min_ratio, max=max_ratio) - - return weight - - def bco_loss( - self, - policy_chosen_logps: torch.FloatTensor, - policy_rejected_logps: torch.FloatTensor, - reference_chosen_logps: torch.FloatTensor, - reference_rejected_logps: torch.FloatTensor, - chosen_embeddings: torch.FloatTensor | None, - rejected_embeddings: torch.FloatTensor | None, - do_train: bool = True, - ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: - """Compute the BCO loss for a batch of policy and reference model log probabilities. - - Args: - policy_chosen_logps: - Log probabilities of the policy model for the chosen responses. Shape: (num(chosen) in batch_size,) - policy_rejected_logps: - Log probabilities of the policy model for the rejected responses. Shape: (num(rejected) in batch_size,) - reference_chosen_logps: - Log probabilities of the reference model for the chosen responses. Shape: (num(chosen) in batch_size,) - reference_rejected_logps: - Log probabilities of the reference model for the rejected responses. Shape: (num(rejected) in - batch_size,) - chosen_embeddings: embeddings of desirable prompts - rejected_embeddings: embeddings of undesirable prompts - do_train: whether to update the running delta value. Default is True. - - Returns: - A tuple of four tensors: (losses, chosen_rewards, rejected_rewards, delta). The losses tensor contains the - BCO loss for each example in the batch. The chosen_rewards and rejected_rewards tensors contain the rewards - for the chosen and rejected responses, respectively. The delta value contains the moving average of all - implicit rewards. - """ - - chosen_logratios = policy_chosen_logps - reference_chosen_logps - chosen_rewards = self.beta * chosen_logratios - - rejected_logratios = policy_rejected_logps - reference_rejected_logps - rejected_rewards = self.beta * rejected_logratios - - if do_train: - self.running.update(torch.cat((chosen_rewards, rejected_rewards), 0).detach()) - delta = torch.as_tensor(self.running.mean, device=chosen_rewards.device) - - chosen_losses = -F.logsigmoid(chosen_rewards - delta) - rejected_losses = -F.logsigmoid(-(rejected_rewards - delta)) - - if self.match_underlying_distribution: - chosen_weight = torch.ones_like(chosen_losses) - rejected_weight = self._get_udm_weight(rejected_embeddings) - - losses = torch.cat((chosen_weight * chosen_losses, rejected_weight * rejected_losses), dim=0) - else: - losses = torch.cat((chosen_losses, rejected_losses), dim=0) - - return losses, chosen_rewards, rejected_rewards, delta - - def get_batch_loss_metrics( - self, - model, - batch: dict[str, list | torch.LongTensor], - do_train: bool = True, - ): - """Compute the BCO loss and other metrics for the given batch of inputs for train or test.""" - metrics = {} - batch = {k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()} - - forward_output = self.forward(model, batch) - ( - policy_chosen_logps, - policy_rejected_logps, - policy_chosen_logits, - policy_rejected_logits, - ) = forward_output[:4] - if self.aux_loss_enabled: - aux_loss = forward_output[4] - - # if reference_logps in batch use them, otherwise use the reference model - if "reference_logps" in batch: - chosen_idx = [i for i in range(batch["reference_logps"].shape[0]) if batch["label"][i] is True] - rejected_idx = [i for i in range(batch["reference_logps"].shape[0]) if batch["label"][i] is False] - - reference_chosen_logps = batch["reference_logps"][chosen_idx, ...] - reference_rejected_logps = batch["reference_logps"][rejected_idx, ...] - else: - with torch.no_grad(): - if self.ref_model is None: - with self.null_ref_context(): - ( - reference_chosen_logps, - reference_rejected_logps, - _, - _, - ) = self.forward(self.model, batch)[:4] - else: - ( - reference_chosen_logps, - reference_rejected_logps, - _, - _, - ) = self.forward(self.ref_model, batch)[:4] - - chosen_embeddings, rejected_embeddings = self._get_prompt_embeddings(batch) - - losses, chosen_rewards, rejected_rewards, delta = self.bco_loss( - policy_chosen_logps, - policy_rejected_logps, - reference_chosen_logps, - reference_rejected_logps, - chosen_embeddings, - rejected_embeddings, - do_train=do_train, - ) - metrics["delta"] = self.accelerator.gather_for_metrics(delta).mean().item() - - num_chosen = torch.Tensor([len(chosen_rewards)]).to(self.accelerator.device) - num_rejected = torch.Tensor([len(rejected_rewards)]).to(self.accelerator.device) - - all_num_chosen = self.accelerator.gather_for_metrics(num_chosen).sum().item() - all_num_rejected = self.accelerator.gather_for_metrics(num_rejected).sum().item() - - if all_num_chosen > 0: - metrics["rewards/chosen_sum"] = ( - self.accelerator.gather_for_metrics(chosen_rewards.nansum()).nansum().item() - ) - metrics["logps/chosen_sum"] = ( - self.accelerator.gather_for_metrics(policy_chosen_logps.nansum()).nansum().item() - ) - metrics["logits/chosen_sum"] = ( - self.accelerator.gather_for_metrics(policy_chosen_logits.nansum()).nansum().item() - ) - metrics["count/chosen"] = all_num_chosen - - if all_num_rejected > 0: - metrics["rewards/rejected_sum"] = ( - self.accelerator.gather_for_metrics(rejected_rewards.nansum()).nansum().item() - ) - metrics["logps/rejected_sum"] = ( - self.accelerator.gather_for_metrics(policy_rejected_logps.nansum()).nansum().item() - ) - metrics["logits/rejected_sum"] = ( - self.accelerator.gather_for_metrics(policy_rejected_logits.nansum()).nansum().item() - ) - metrics["count/rejected"] = all_num_rejected - - loss = losses.nanmean() - if self.aux_loss_enabled: - loss += self.aux_loss_coef * aux_loss - - return loss, metrics - - def compute_loss( - self, - model: PreTrainedModel | nn.Module, - inputs: dict[str, torch.Tensor | Any], - return_outputs=False, - num_items_in_batch=None, - ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]: - compute_loss_context_manager = ( - autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() - ) - - with compute_loss_context_manager: - loss, metrics = self.get_batch_loss_metrics(model, inputs) - - # Make sure to move the loss to the device the original accumulating loss is at back in the `Trainer` class: - loss = loss.to(self.args.device) - # force log the metrics - if self.accelerator.is_main_process: - self.store_metrics(metrics, train_eval="train") - - if return_outputs: - return (loss, metrics) - return loss - - def store_metrics(self, metrics: dict[str, float], train_eval: Literal["train", "eval"] = "train") -> None: - for key, value in metrics.items(): - self._stored_metrics[train_eval][key].append(value) - - def _get_train_sampler(self, dataset: Dataset | None = None) -> torch.utils.data.Sampler | None: - if dataset is None: - dataset = self.train_dataset - if dataset is None or not has_length(dataset): - return None - return SequentialSampler(dataset) - - def generate_from_model_and_ref(self, model, batch: dict[str, torch.LongTensor]) -> tuple[str, str]: - """Generate samples from the model and reference model for the given batch of inputs.""" - - # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with - # the torch amp context manager as some hidden states are silently casted to full precision. - generate_context_manager = ( - autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() - ) - with generate_context_manager: - policy_output = model.generate( - input_ids=batch["prompt_input_ids"], - attention_mask=batch["prompt_attention_mask"], - max_length=self.max_length, - do_sample=True, - pad_token_id=self.processing_class.pad_token_id, - ) - - # if reference_output in batch use that otherwise use the reference model - if "reference_output" in batch: - reference_output = batch["reference_output"] - else: - if self.ref_model is None: - with self.null_ref_context(): - reference_output = self.model.generate( - input_ids=batch["prompt_input_ids"], - attention_mask=batch["prompt_attention_mask"], - max_length=self.max_length, - do_sample=True, - pad_token_id=self.processing_class.pad_token_id, - ) - else: - reference_output = self.ref_model.generate( - input_ids=batch["prompt_input_ids"], - attention_mask=batch["prompt_attention_mask"], - max_length=self.max_length, - do_sample=True, - pad_token_id=self.processing_class.pad_token_id, - ) - - policy_output = pad_to_length(policy_output, self.max_length, self.processing_class.pad_token_id) - policy_output_decoded = self.processing_class.batch_decode(policy_output, skip_special_tokens=True) - - reference_output = pad_to_length(reference_output, self.max_length, self.processing_class.pad_token_id) - reference_output_decoded = self.processing_class.batch_decode(reference_output, skip_special_tokens=True) - - return policy_output_decoded, reference_output_decoded - - def prediction_step( - self, - model: PreTrainedModel | nn.Module, - inputs: dict[str, torch.Tensor | Any], - prediction_loss_only: bool, - ignore_keys: list[str] | None = None, - ): - if ignore_keys is None: - if hasattr(model, "config"): - ignore_keys = getattr(model.config, "keys_to_ignore_at_inference", []) - else: - ignore_keys = [] - - prediction_context_manager = ( - autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() - ) - with torch.no_grad(), prediction_context_manager: - loss, metrics = self.get_batch_loss_metrics(model, inputs, do_train=False) - - # force log the metrics - if self.accelerator.is_main_process: - self.store_metrics(metrics, train_eval="eval") - - if prediction_loss_only: - return (loss.detach(), None, None) - - # logits for the chosen and rejected samples from model - logits_dict = {} - if "logits/chosen_sum" in metrics: - logits_dict["eval_logits/chosen"] = metrics["logits/chosen_sum"] - if "logits/rejected_sum" in metrics: - logits_dict["eval_logits/rejected"] = metrics["logits/rejected_sum"] - logits = [v for k, v in logits_dict.items() if k not in ignore_keys] - logits = torch.tensor(logits, device=self.accelerator.device) - labels = torch.zeros(logits.shape[0], device=self.accelerator.device) - - return (loss.detach(), logits, labels) - - def evaluation_loop( - self, - dataloader: DataLoader, - description: str, - prediction_loss_only: bool | None = None, - ignore_keys: list[str] | None = None, - metric_key_prefix: str = "eval", - ) -> EvalLoopOutput: - """ - Overriding built-in evaluation loop to store metrics for each batch. Prediction/evaluation loop, shared by - `Trainer.evaluate()` and `Trainer.predict()`. - - Works both with or without labels. - """ - - # Sample and save to game log if requested (for one batch to save time) - if self.generate_during_eval: - # Generate random indices within the range of the total number of samples - num_samples = len(dataloader.dataset) - random_indices = random.sample(range(num_samples), k=self.args.eval_batch_size) - - # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader - random_batch_dataset = dataloader.dataset.select(random_indices) - random_batch = self.data_collator(random_batch_dataset) - random_batch = self._prepare_inputs(random_batch) - - target_labels = torch.tensor(random_batch["label"], dtype=torch.bool, device=self.accelerator.device) - target_indices = torch.where(~target_labels)[0] - target_batch = { - "prompt_input_ids": random_batch["prompt_input_ids"][target_indices], - "prompt_attention_mask": random_batch["prompt_attention_mask"][target_indices], - "prompt": itemgetter(*target_indices)(random_batch["prompt"]), - } - policy_output_decoded, ref_output_decoded = self.generate_from_model_and_ref(self.model, target_batch) - - table = pd.DataFrame( - columns=["Prompt", "Policy", "Ref Model"], - data=[ - [prompt, pol[len(prompt) :], ref[len(prompt) :]] - for prompt, pol, ref in zip( - target_batch["prompt"], policy_output_decoded, ref_output_decoded, strict=True - ) - ], - ) - if "wandb" in self.args.report_to: - wandb.log({"game_log": wandb.Table(data=table)}) - - if "comet_ml" in self.args.report_to: - log_table_to_comet_experiment( - name="game_log.csv", - table=table, - ) - - # Base evaluation - initial_output = super().evaluation_loop( - dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix - ) - - return initial_output - - def log(self, logs: dict[str, float], start_time: float | None = None) -> None: - """ - Log `logs` on the various objects watching training, including stored metrics. - - Args: - logs (`dict[str, float]`): - The values to log. - start_time (`float`, *optional*): - Start time of the training. - """ - # logs either has 'loss' or 'eval_loss' - train_eval = "train" if "loss" in logs else "eval" - # train metrics should have no prefix, eval should have 'eval_' - prefix = "eval_" if train_eval == "eval" else "" - # accumulate average metrics from sums and lengths - for split in ["chosen", "rejected"]: - if f"count/{split}" in self._stored_metrics[train_eval]: - count_sum = torch.Tensor(self._stored_metrics[train_eval][f"count/{split}"]).sum().item() - for metric in ["rewards", "logps", "logits"]: - logs[f"{prefix}{metric}/{split}"] = ( - torch.Tensor(self._stored_metrics[train_eval][f"{metric}/{split}_sum"]).sum().item() - / count_sum - ) - # delete obsolete metric - del self._stored_metrics[train_eval][f"{metric}/{split}_sum"] - del self._stored_metrics[train_eval][f"count/{split}"] - # calculate reward margin - if f"{prefix}rewards/chosen" in logs and f"{prefix}rewards/rejected" in logs: - logs[f"{prefix}rewards/margins"] = logs[f"{prefix}rewards/chosen"] - logs[f"{prefix}rewards/rejected"] - # Add averaged stored metrics to logs - for key, metrics in self._stored_metrics[train_eval].items(): - logs[f"{prefix}{key}"] = torch.Tensor(metrics).mean().item() - del self._stored_metrics[train_eval] - return super().log(logs, start_time) - - # Ensure the model card is saved along with the checkpoint - def _save_checkpoint(self, model, trial): - if self.args.hub_model_id is None: - model_name = Path(self.args.output_dir).name - else: - model_name = self.args.hub_model_id.split("/")[-1] - self.create_model_card(model_name=model_name) - super()._save_checkpoint(model, trial) diff --git a/trl/import_utils.py b/trl/import_utils.py index 325ef505516..fc080a9bea1 100644 --- a/trl/import_utils.py +++ b/trl/import_utils.py @@ -69,10 +69,6 @@ def is_jmespath_available() -> bool: return _is_package_available("jmespath") -def is_joblib_available() -> bool: - return _is_package_available("joblib") - - def is_liger_kernel_available(min_version: str = LIGER_KERNEL_MIN_VERSION) -> bool: _liger_kernel_available, _liger_kernel_version = _is_package_available("liger_kernel", return_version=True) return _liger_kernel_available and Version(_liger_kernel_version) >= Version(min_version) diff --git a/trl/trainer/base_trainer.py b/trl/trainer/base_trainer.py index 22a57af155e..4775b2e607e 100644 --- a/trl/trainer/base_trainer.py +++ b/trl/trainer/base_trainer.py @@ -45,7 +45,6 @@ "A2POTrainer", "AsyncDistillationTrainer", "AsyncGRPOTrainer", - "BCOTrainer", "CPOTrainer", "GKDTrainer", "GMPOTrainer",