Method comparison evaluation suite (#2395)

Introduction of a method evaluation suite.

We generally face the problem that there is little knowledge on what PEFT methods perform best. To this end we decided to build an evaluation suite that has defined tasks, shared hyper-parameters and can be extended with new tasks and new method configurations over time.

For the sake of comparison we've not decided to incorporate user-submitted results but we encourage users to inspect the results, suggest new experiments and improve the configuration of methods if they're deemed unfavorable.

As of now there's only one task based on the MetaMathQA dataset which has the benefit of being complex while still fitting on a consumer GPU.

Notable changes in this squash:

* Add default training params

The experiment specific training params use the default training params
but can override any parameter from it if needed. However, this way it's
easier to make a change to all experiments (say, I want to change the
base model, I don't need to change each individual
training_parameters.json).

* Add possibility to change attn implementation

However, both flash attention 2 and flex attention are slower on my
system. Thus, stay with default None (-> SDPA).

* Refactor to use GenerationConfig

Allows to more easily use, say, static cache, which is the new default,
as it's faster (apart from the first pass)

* Better parsing of answers

E.g. 1/2 == 0.5

* Keep adapter file by default after train run

But add --clean to delete it.

Keeping the adapter can be useful if the user wants to run further tests
with the trained model.

---------

Co-authored-by: Benjamin Bossan <[email protected]>

Author: githubnemo

.gitignore

@@ -139,3 +139,7 @@ dmypy.json
 
 # More test things
 wandb
+
+# method_comparison logs
+method_comparison/MetaMathQA/cancelled_results/
+method_comparison/MetaMathQA/temporary_results/

method_comparison/Makefile

@@ -0,0 +1,213 @@
+# PEFT method comparison on the MetaMathQA and GSM8K datasets
+
+## Goal
+
+This goal is to provide a benchmarking framework for the different PEFT methods that are implemented. It is important that evaluating different PEFT methods is reproducible, idempotent, and version-controlled. Results for more PEFT methods can be added over time.
+
+## Dataset
+
+This task trains on the [MetaMathQA]((https://huggingface.co/datasets/meta-math/MetaMathQA)) dataset and validates/tests on the [GSM8K](https://huggingface.co/datasets/openai/gsm8k) dataset ("main").
+
+For the model to attain good accuracy, it needs to learn to adhere to the output format and it must express basic chain of thought reasoning capabilities to get to the correct result in the first place. The task is challenging for models in the sub 7B parameter range.
+
+The train set uses the whole of MetaMathQA. The validation set is a random sample from the train set of GSM8K. The test set is the whole of the GSM8K test set.
+
+## Running
+
+Create an experiment in the `experiment/<peft-method>` folder of your choice and give it a name (the name itself does not matter but helps identify the experiment). An example would be `experiments/lora/llama-3.2-3B-rank32/`. Inside that directory, create 2 files:
+
+- `adapter_config.json`
+- Optional: `training_parameters.json`
+
+### `adapter_config.json`
+
+This must be a valid PEFT configuration. It is easiest to create it programmatically, e.g.:
+
+```python
+from peft import LoraConfig
+
+config = LoraConfig(...)
+config.save_pretrained(<path-to-experiment>)
+```
+
+### `training_parameters.json`
+
+There is a default file for the non-PEFT parameters: `default_training_params.json`. This contains all the other parameters that are relevant for training, e.g. the base model id, number of steps, batch size, learning rate, etc. If parameters that differ from the defaults are needed for a specific experiment, place a `training_parameters.json` into the experiment directory and adjust the parameters that need changing. The other parametes are taken from the aforementioned default config.
+
+For an overview of all possible arguments, you can also check the `TrainConfig` `dataclass` in `utils.py`.
+
+### Runtime performance
+
+Several factors should be considered to achieve a fast runtime performance. Besides the obvious factors like `max_steps` or the base model size, we found the following factors to have a significant impact:
+
+#### Eval batch size
+
+Regarding the `batch_size_eval` parameter, it is quite critical since evaluation takes up a significant portion of the training time and batching helps with reducing that. It should be possible to choose a value that is multiple times higher than the batch size used for training (`batch_size`). You should also pay attention to the size of the validation set -- e.g. if it's 50, don't choose a `batch_size_eval` of 40, as that results in a large batch of 30 and a small batch of 10. 25 might be a better choice. Also, ensure via a quick train run that the batch size does not lead to out of memory errors -- getting this error at the very end on evaluating the test set would be quite a loss of time.
+
+#### Generation length
+
+During testing, we discovered that the validation time is greatly inflated by just a few very long generations. Those can inflate the validation time by a factor of 3 or more. At the same time, we discovered that these long generations do not help with accuracy -- in fact, if they exceed the maximum configured length, they're just cut off mid sentence and would thus produce an accuracy of 0 anyway.
+
+To remedy this, we now set both `max_length` and `max_new_tokens` for the generation kwargs in the default training parameters. Normally, this is not possible when using transformers, as the latter argument overrides the former. However, we have added special logic inside of `get_generation_config` which takes both and chooses the smaller of the two. This way, we can get rid of these excessively long generations, thus considerably reducing eval times, while still guaranteeing a maximum total generation length to guard against OOM errors. Testing showed that this does not hamper test accuracy. It is therefore recommended not to change these settings.
+
+#### Bucketing
+
+The length of the sequences in the training data can vary a lot. Therefore, if samples are taken randomly from the training dataset, we will end up with batches containing very short and very long sequences. This is bad because the batch will be padded to the longest sequence, slowing down training. The obvious solution would be to sort the whole dataset by sequence length, but this is also bad because it introduces an order bias (e.g. first training on only short and then on only long answers).
+
+The solution is to find a trade off between the two factors. This is achieved by the `BucketIterator`. It first creates buckets that contain multiple batches, e.g. 20x the batch size. The bucket is then sorted by sequence length and then batches are yielded from the bucket. Therefore, we have a small order bias within a bucket but not between buckets, stricking a good balance between training speed and training loss.
+
+From practical experiments, for a batch size of 4, a bucket size of 80 provides a good balance with only slightly lower training loss but cutting training time by 25%. For eval, we don't use the iterator since there, the batch size is relatively big and thus there is little upside.
+
+### Start a run
+
+Once everything is set up properly, start a run by using the `run.py` script. Pass `-v` for verbose output to the console (recommended if observing the progress is desired). As an example, for `experiments/lora/llama-3.2-3B-rank32/` the invocation would be:
+
+```sh
+python run.py -v experiments/lora/llama-3.2-3B-rank32/
+```
+
+By default, the adapter will be saved in a temporary file for further inspection if needed. The prevent this, add the `--clean` flag to the call.
+
+### Run status
+
+The run can be categorized 3 different states:
+
+1. Main run: You are on the `main` branch and the run ended successfully. The results are stored in the `results` folder and are used for further analysis.
+2. Test run: You are not on the `main` branch and the run ended successfully. The results are stored in the `temporary_results` folder and are not used for further analysis.
+3. The run was cancelled (`ctrl + c`). The results are stored in the `cancelled_results` folder and are not used for further analysis.
+
+## Outputs
+
+Results are stored in one of the result directories. An example output could look like so:
+
+```js
+{
+  "run_info": {
+    "created_at": "2025-03-05T13:50:05+00:00",
+    "total_time": 2711.0915009640157,
+    "experiment_name": "ia3/lr_0.001",
+    "peft_branch": "ben-method-comparison",
+    "train_config": {
+      "model_id": "meta-llama/Llama-3.2-3B",
+      "dtype": "bfloat16",
+      "max_seq_length": 768,
+      "batch_size": 4,
+      "batch_size_eval": 51,
+      "max_steps": 5000,
+      "eval_steps": 250,
+      "compile": false,
+      "query_template": "Question: {query} Think step by step.\nAnswer:",
+      "seed": 0,
+      "grad_norm_clip": 1.0,
+      "optimizer_kwargs": {
+        "lr": 0.001
+      },
+      "lr_scheduler": "cosine",
+      "use_amp": false,
+      "generation_kwargs": {
+        "max_length": 800
+      },
+      "attn_implementation": null
+    },
+    "peft_config": {
+      "task_type": null,
+      "peft_type": "IA3",
+      "auto_mapping": null,
+      "base_model_name_or_path": "meta-llama/Llama-3.2-3B",
+      "revision": null,
+      "inference_mode": false,
+      "target_modules": [
+        "v_proj",
+        "k_proj",
+        "down_proj"
+      ],
+      "exclude_modules": null,
+      "feedforward_modules": [
+        "down_proj"
+      ],
+      "fan_in_fan_out": false,
+      "modules_to_save": null,
+      "init_ia3_weights": true
+    }
+  },
+  "train_info": {
+    "cuda_memory_reserved_avg": 14229219940,
+    "cuda_memory_max": 24847056896,
+    "cuda_memory_reserved_99th": 19115624366,
+    "train_time": 2238.65277833899,
+    "file_size": 1157064,
+    "status": "success",
+    "metrics": [
+      {
+        "step": 250,
+        "valid accuracy": 0.0784313725490196,
+        "train loss": 1.1336498007774354,
+        "train samples": 1000
+      },
+      [...]
+      {
+        "step": 5000,
+        "valid accuracy": 0.21568627450980393,
+        "train loss": 0.6345920492410659,
+        "train samples": 20000
+      },
+      {
+        "step": 5000,
+        "test accuracy": 0.35129740518962077,
+        "train loss": 0.6345920492410659,
+        "train samples": 20000,
+        "train total tokens": 4197579
+      }
+    ]
+  },
+  "meta_info": {
+    "model_sha": "13afe5124825b4f3751f836b40dafda64c1ed062",
+    "model_created_at": "2024-09-18T15:23:48+00:00",
+    "dataset_sha": "aa4f34d3d2d3231299b5b03d9b3e5a20da45aa18",
+    "dataset_created_at": "2023-09-21T17:22:46+00:00",
+    "package_info": {
+      "transformers-version": "4.50.0.dev0",
+      "transformers-commit-hash": "752ef3fd4e70869626ec70657a770a85c0ad9219",
+      "peft-version": "0.14.1.dev0",
+      "peft-commit-hash": "a447a4e5ecd87b7d57733f4df9616a328cf130f4",
+      "datasets-version": "3.3.2",
+      "datasets-commit-hash": null,
+      "bitsandbytes-version": "0.45.2",
+      "bitsandbytes-commit-hash": null,
+      "torch-version": "2.6.0+cu124",
+      "torch-commit-hash": null
+    },
+    "system_info": {
+      "system": "Linux",
+      "release": "6.11.0-17-generic",
+      "version": "#17~24.04.2-Ubuntu SMP PREEMPT_DYNAMIC Mon Jan 20 22:48:29 UTC 2",
+      "machine": "x86_64",
+      "processor": "x86_64",
+      "gpu": "NVIDIA GeForce RTX 4090"
+    },
+    "pytorch_info": "PyTorch built with: [...]"
+  }
+}
+```
+
+## Dependencies
+
+Apart from the normal PEFT dependencies, ensure that the packages in the `requirements.txt` are installed, e.g. via:
+
+```sh
+python -m pip install -r requirements.txt
+```
+
+Python 3.12+ is required.
+
+## Open tasks
+
+- consider using `DataLoader`
+- consider adding https://github.com/huggingface/Math-Verify
+- consider adding `weight` argument to cross entropy calculation to downweight the EOS token, but it would require calculating the loss manually instead of relying on transformers (see https://github.com/huggingface/transformers/blob/6a876462c308bd7cd7d3ca8e93abaa7d5b02e90e/src/transformers/loss/loss_utils.py#L24-L48)
+- do a sanity check against/comparison with transformers Trainer
+- consider using vLLM to potentially speed up generations, at least for the test set
+- using `torch.compile` leads to a huge slowdown, investigate (maybe recompiles), although it does save memory
+- AMP does not appear to help, investigate
+- packing of sequences (but this probably requires adjusting the attention matrix)
+- clean up what gets printed and where (stdout, stderr)

method_comparison/MetaMathQA/README.md

@@ -0,0 +1,109 @@
+# Copyright 2025-present the HuggingFace Inc. team.
+#
+# 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.
+
+"""
+All utilities related to data handling.
+"""
+
+from functools import partial
+from typing import Callable
+
+import datasets
+import numpy as np
+from datasets import Dataset, load_dataset
+
+
+# with a token limit of 768 for query + response, we have to exclude all texts with length > 1304; this leaves 93.8% of
+# the dataset
+CHAR_LIMIT = 1300
+# train/valid/test split -- note that evaluation takes quite long, so don't choose too large sizes for the valid set,
+# since it's run multiple times during training; test is only run once at the end and thus can be larger
+VALID_SIZE = 50
+
+
+def get_filtered_dataset(*, ds: datasets.Dataset, print_fn: Callable[..., None]) -> Dataset:
+    """Return the filtered dataset, with long queries removed.
+
+    We determined that 99% of queries have 529 or fewer characters. Characters roughly correspond to tokens, so this is
+    a good proxy. We cannot use tokens directly, as that depends on the tokenizer, which can be different for each
+    model, but we want the same filter for each model.
+
+    """
+    char_lengths = [len(f"{q} {r}") for q, r in zip(ds["query"], ds["response"])]
+    idx_filtered = [i for i, length in enumerate(char_lengths) if length <= CHAR_LIMIT]
+    print_fn(f"Filtered dataset: {100 * len(idx_filtered) / len(ds):.1f}% of the original dataset")
+    return ds.select(idx_filtered)
+
+
+def get_train_valid_test_datasets(
+    *, tokenizer, query_template: str, print_fn: Callable[..., None]
+) -> tuple[Dataset, Dataset, Dataset]:
+    """
+    Return the indices of the train, valid, and test splits of the dataset.
+
+    We cannot use ds.train_test_split(..., stratify_by_column="type") as it gives:
+
+    > ValueError: Stratifying by column is only supported for ClassLabel column, and column type is Value.
+
+    even after calling ds_filtered.class_encode_column("type"). Thus, using sklearn's StratifiedKFold instead.
+    """
+    metamath = load_dataset("meta-math/MetaMathQA")["train"]
+    metamath = get_filtered_dataset(ds=metamath, print_fn=print_fn)
+
+    # gsmk8k does not need to be filtered as query and response are short enough
+    gsm8k = load_dataset("openai/gsm8k", "main")
+    gsm8k = gsm8k.rename_columns({"question": "query", "answer": "response"})
+    gsm8k_train = gsm8k["train"]
+    gsm8k_test = gsm8k["test"]
+
+    np.random.seed(0)
+    indices = np.arange(len(gsm8k_train))
+    np.random.shuffle(indices)
+    idx_valid = indices[:VALID_SIZE]
+
+    ds_train = metamath
+    ds_valid = gsm8k_train.select(idx_valid)
+    ds_test = gsm8k_test
+
+    print_fn(f"Train size: {len(ds_train)}")
+    print_fn(f"Valid size: {len(ds_valid)}")
+    print_fn(f"Test size: {len(ds_test)}")
+
+    tokenize_with_answer_ = partial(tokenize_with_answer, tokenizer=tokenizer, template=query_template)
+    tokenize_wo_answer_ = partial(tokenize_wo_answer, tokenizer=tokenizer, template=query_template)
+    ds_train = ds_train.map(tokenize_with_answer_, batched=True).remove_columns(["type", "query", "original_question"])
+    ds_valid = ds_valid.map(tokenize_wo_answer_, batched=True).remove_columns(["query"])
+    ds_test = ds_test.map(tokenize_wo_answer_, batched=True).remove_columns(["query"])
+
+    return ds_train, ds_valid, ds_test
+
+
+def tokenize_with_answer(samples, tokenizer, template):
+    queries = [template.format(query=sample) + answer for sample, answer in zip(samples["query"], samples["response"])]
+    tokenized = tokenizer(queries)
+    tokenized["input_ids"] = [input_ids[: tokenizer.model_max_length] for input_ids in tokenized["input_ids"]]
+    tokenized["attention_mask"] = [
+        input_ids[: tokenizer.model_max_length] for input_ids in tokenized["attention_mask"]
+    ]
+    return tokenized
+
+
+def tokenize_wo_answer(samples, tokenizer, template):
+    queries = [template.format(query=sample) for sample in samples["query"]]
+    tokenized = tokenizer(queries)
+    tokenized["input_ids"] = [input_ids[: tokenizer.model_max_length] for input_ids in tokenized["input_ids"]]
+    tokenized["attention_mask"] = [
+        input_ids[: tokenizer.model_max_length] for input_ids in tokenized["attention_mask"]
+    ]
+    return tokenized

method_comparison/MetaMathQA/cancelled_results/.gitkeep

@@ -0,0 +1,25 @@
+{
+  "model_id": "meta-llama/Llama-3.2-3B",
+  "dtype": "bfloat16",
+  "max_seq_length": 768,
+  "batch_size": 4,
+  "batch_size_eval": 50,
+  "max_steps": 5000,
+  "eval_steps": 250,
+  "compile": false,
+  "seed": 0,
+  "grad_norm_clip": 1.0,
+  "optimizer_kwargs": {
+    "lr": 1e-4,
+    "weight_decay": 0.1
+  },
+  "lr_scheduler": "cosine",
+  "use_amp": false,
+  "autocast_adapter_dtype": true,
+  "attn_implementation": null,
+  "generation_kwargs": {
+    "max_length": 800,
+    "max_new_tokens": 300
+  },
+  "query_template": "Question: {query} Think step by step.\nAnswer:"
+}

method_comparison/MetaMathQA/data.py

@@ -0,0 +1,39 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": null,
+  "beta1": 0.85,
+  "beta2": 0.85,
+  "bias": "none",
+  "corda_config": null,
+  "deltaT": 1,
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": false,
+  "init_lora_weights": true,
+  "init_r": 64,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 8,
+  "lora_bias": false,
+  "lora_dropout": 0.0,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "orth_reg_weight": 0.5,
+  "peft_type": "ADALORA",
+  "r": 8,
+  "rank_pattern": null,
+  "revision": null,
+  "target_modules": null,
+  "target_r": 32,
+  "task_type": null,
+  "tfinal": 500,
+  "tinit": 200,
+  "total_step": 5000,
+  "use_dora": false,
+  "use_rslora": false
+}