We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
sft_llama2 error
Hello! 

I made no changes to the sample code or the dataset. I just simply wanted to get it to run the stacked llama2 example. However, I get this error:

UserWarning: The passed formatting_func has more than one argument. Usually that function should have a single argument `example` which corresponds to the dictionary returned by each element of the dataset. Make sure you know what you are doing.

 in __init__
    raise ValueError("the `--group_by_length` option is only available for `Dataset`, not `IterableDataset")
ValueError: the `--group_by_length` option is only available for `Dataset`, not `IterableDataset

(Sorry, I must omit certain parts of the error message since I am working on an institution's server). 

Thank you!

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/research_projects/stack_llama_2/scripts/sft_llama2.py`
Content:
```
1 # Fine-Tune Llama2-7b on SE paired dataset
2 import os
3 from dataclasses import dataclass, field
4 from typing import Optional
5 
6 import torch
7 from datasets import load_dataset
8 from peft import AutoPeftModelForCausalLM, LoraConfig
9 from tqdm import tqdm
10 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, HfArgumentParser, TrainingArguments
11 
12 from trl import SFTTrainer
13 from trl.trainer import ConstantLengthDataset
14 
15 
16 @dataclass
17 class ScriptArguments:
18     model_name: Optional[str] = field(default="meta-llama/Llama-2-7b-hf", metadata={"help": "the model name"})
19     log_with: Optional[str] = field(default="wandb", metadata={"help": "use 'wandb' to log with wandb"})
20 
21     dataset_name: Optional[str] = field(default="lvwerra/stack-exchange-paired", metadata={"help": "the dataset name"})
22     subset: Optional[str] = field(default="data/finetune", metadata={"help": "the subset to use"})
23     split: Optional[str] = field(default="train", metadata={"help": "the split to use"})
24     size_valid_set: Optional[int] = field(default=4000, metadata={"help": "the size of the validation set"})
25     streaming: Optional[bool] = field(default=True, metadata={"help": "whether to stream the dataset"})
26     shuffle_buffer: Optional[int] = field(default=5000, metadata={"help": "the shuffle buffer size"})
27     seq_length: Optional[int] = field(default=1024, metadata={"help": "the sequence length"})
28 
29     max_steps: Optional[int] = field(default=500, metadata={"help": "the maximum number of sgd steps"})
30     logging_steps: Optional[int] = field(default=10, metadata={"help": "the logging frequency"})
31     save_steps: Optional[int] = field(default=10, metadata={"help": "the saving frequency"})
32     per_device_train_batch_size: Optional[int] = field(default=4, metadata={"help": "the per device train batch size"})
33     per_device_eval_batch_size: Optional[int] = field(default=1, metadata={"help": "the per device eval batch size"})
34     gradient_accumulation_steps: Optional[int] = field(default=2, metadata={"help": "the gradient accumulation steps"})
35     gradient_checkpointing: Optional[bool] = field(
36         default=True, metadata={"help": "whether to use gradient checkpointing"}
37     )
38     group_by_length: Optional[bool] = field(default=True, metadata={"help": "whether to group by length"})
39 
40     lora_alpha: Optional[float] = field(default=16, metadata={"help": "the lora alpha parameter"})
41     lora_dropout: Optional[float] = field(default=0.05, metadata={"help": "the lora dropout parameter"})
42     lora_r: Optional[int] = field(default=8, metadata={"help": "the lora r parameter"})
43 
44     learning_rate: Optional[float] = field(default=1e-4, metadata={"help": "the learning rate"})
45     lr_scheduler_type: Optional[str] = field(default="cosine", metadata={"help": "the lr scheduler type"})
46     num_warmup_steps: Optional[int] = field(default=100, metadata={"help": "the number of warmup steps"})
47     weight_decay: Optional[float] = field(default=0.05, metadata={"help": "the weight decay"})
48     optimizer_type: Optional[str] = field(default="paged_adamw_32bit", metadata={"help": "the optimizer type"})
49 
50     output_dir: Optional[str] = field(default="./results", metadata={"help": "the output directory"})
51     log_freq: Optional[int] = field(default=1, metadata={"help": "the logging frequency"})
52 
53 
54 parser = HfArgumentParser(ScriptArguments)
55 script_args = parser.parse_args_into_dataclasses()[0]
56 
57 
58 def chars_token_ratio(dataset, tokenizer, nb_examples=400):
59     """
60     Estimate the average number of characters per token in the dataset.
61     """
62     total_characters, total_tokens = 0, 0
63     for _, example in tqdm(zip(range(nb_examples), iter(dataset)), total=nb_examples):
64         text = prepare_sample_text(example)
65         total_characters += len(text)
66         if tokenizer.is_fast:
67             total_tokens += len(tokenizer(text).tokens())
68         else:
69             total_tokens += len(tokenizer.tokenize(text))
70 
71     return total_characters / total_tokens
72 
73 
74 def print_trainable_parameters(model):
75     """
76     Prints the number of trainable parameters in the model.
77     """
78     trainable_params = 0
79     all_param = 0
80     for _, param in model.named_parameters():
81         all_param += param.numel()
82         if param.requires_grad:
83             trainable_params += param.numel()
84     print(
85         f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
86     )
87 
88 
89 def prepare_sample_text(example):
90     """Prepare the text from a sample of the dataset."""
91     text = f"Question: {example['question']}\n\nAnswer: {example['response_j']}"
92     return text
93 
94 
95 def create_datasets(tokenizer, args):
96     dataset = load_dataset(
97         args.dataset_name,
98         data_dir=args.subset,
99         split=args.split,
100         use_auth_token=True,
101         num_proc=args.num_workers if not args.streaming else None,
102         streaming=args.streaming,
103     )
104     if args.streaming:
105         print("Loading the dataset in streaming mode")
106         valid_data = dataset.take(args.size_valid_set)
107         train_data = dataset.skip(args.size_valid_set)
108         train_data = train_data.shuffle(buffer_size=args.shuffle_buffer, seed=None)
109     else:
110         dataset = dataset.train_test_split(test_size=0.005, seed=None)
111         train_data = dataset["train"]
112         valid_data = dataset["test"]
113         print(f"Size of the train set: {len(train_data)}. Size of the validation set: {len(valid_data)}")
114 
115     chars_per_token = chars_token_ratio(train_data, tokenizer)
116     print(f"The character to token ratio of the dataset is: {chars_per_token:.2f}")
117 
118     train_dataset = ConstantLengthDataset(
119         tokenizer,
120         train_data,
121         formatting_func=prepare_sample_text,
122         infinite=True,
123         seq_length=args.seq_length,
124         chars_per_token=chars_per_token,
125     )
126     valid_dataset = ConstantLengthDataset(
127         tokenizer,
128         valid_data,
129         formatting_func=prepare_sample_text,
130         infinite=False,
131         seq_length=args.seq_length,
132         chars_per_token=chars_per_token,
133     )
134     return train_dataset, valid_dataset
135 
136 
137 bnb_config = BitsAndBytesConfig(
138     load_in_4bit=True,
139     bnb_4bit_quant_type="nf4",
140     bnb_4bit_compute_dtype=torch.bfloat16,
141 )
142 
143 base_model = AutoModelForCausalLM.from_pretrained(
144     script_args.model_name,
145     quantization_config=bnb_config,
146     device_map={"": 0},
147     trust_remote_code=True,
148     use_auth_token=True,
149 )
150 base_model.config.use_cache = False
151 
152 peft_config = LoraConfig(
153     r=script_args.lora_r,
154     lora_alpha=script_args.lora_alpha,
155     lora_dropout=script_args.lora_dropout,
156     target_modules=["q_proj", "v_proj"],
157     bias="none",
158     task_type="CAUSAL_LM",
159 )
160 
161 tokenizer = AutoTokenizer.from_pretrained(script_args.model_name, trust_remote_code=True)
162 tokenizer.pad_token = tokenizer.eos_token
163 tokenizer.padding_side = "right"  # Fix weird overflow issue with fp16 training
164 
165 
166 training_args = TrainingArguments(
167     output_dir=script_args.output_dir,
168     per_device_train_batch_size=script_args.per_device_train_batch_size,
169     gradient_accumulation_steps=script_args.gradient_accumulation_steps,
170     per_device_eval_batch_size=script_args.per_device_eval_batch_size,
171     learning_rate=script_args.learning_rate,
172     logging_steps=script_args.logging_steps,
173     max_steps=script_args.max_steps,
174     report_to=script_args.log_with,
175     save_steps=script_args.save_steps,
176     group_by_length=script_args.group_by_length,
177     lr_scheduler_type=script_args.lr_scheduler_type,
178     warmup_steps=script_args.num_warmup_steps,
179     optim=script_args.optimizer_type,
180     bf16=True,
181     remove_unused_columns=False,
182     run_name="sft_llama2",
183 )
184 
185 train_dataset, eval_dataset = create_datasets(tokenizer, script_args)
186 
187 trainer = SFTTrainer(
188     model=base_model,
189     train_dataset=train_dataset,
190     eval_dataset=eval_dataset,
191     peft_config=peft_config,
192     packing=True,
193     max_seq_length=None,
194     tokenizer=tokenizer,
195     args=training_args,
196 )
197 trainer.train()
198 trainer.save_model(script_args.output_dir)
199 
200 output_dir = os.path.join(script_args.output_dir, "final_checkpoint")
201 trainer.model.save_pretrained(output_dir)
202 
203 # Free memory for merging weights
204 del base_model
205 torch.cuda.empty_cache()
206 
207 model = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map="auto", torch_dtype=torch.bfloat16)
208 model = model.merge_and_unload()
209 
210 output_merged_dir = os.path.join(script_args.output_dir, "final_merged_checkpoint")
211 model.save_pretrained(output_merged_dir, safe_serialization=True)
212
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/examples/research_projects/stack_llama_2/scripts/sft_llama2.py b/examples/research_projects/stack_llama_2/scripts/sft_llama2.py
--- a/examples/research_projects/stack_llama_2/scripts/sft_llama2.py
+++ b/examples/research_projects/stack_llama_2/scripts/sft_llama2.py
@@ -35,7 +35,8 @@
     gradient_checkpointing: Optional[bool] = field(
         default=True, metadata={"help": "whether to use gradient checkpointing"}
     )
-    group_by_length: Optional[bool] = field(default=True, metadata={"help": "whether to group by length"})
+    group_by_length: Optional[bool] = field(default=False, metadata={"help": "whether to group by length"})
+    packing: Optional[bool] = field(default=True, metadata={"help": "whether to use packing for SFTTrainer"})
 
     lora_alpha: Optional[float] = field(default=16, metadata={"help": "the lora alpha parameter"})
     lora_dropout: Optional[float] = field(default=0.05, metadata={"help": "the lora dropout parameter"})
@@ -54,6 +55,9 @@
 parser = HfArgumentParser(ScriptArguments)
 script_args = parser.parse_args_into_dataclasses()[0]
 
+if script_args.group_by_length and script_args.packing:
+    raise ValueError("Cannot use both packing and group by length")
+
 
 def chars_token_ratio(dataset, tokenizer, nb_examples=400):
     """
@@ -189,7 +193,7 @@
     train_dataset=train_dataset,
     eval_dataset=eval_dataset,
     peft_config=peft_config,
-    packing=True,
+    packing=script_args.packing,
     max_seq_length=None,
     tokenizer=tokenizer,
     args=training_args,

{"golden_diff": "diff --git a/examples/research_projects/stack_llama_2/scripts/sft_llama2.py b/examples/research_projects/stack_llama_2/scripts/sft_llama2.py\n--- a/examples/research_projects/stack_llama_2/scripts/sft_llama2.py\n+++ b/examples/research_projects/stack_llama_2/scripts/sft_llama2.py\n@@ -35,7 +35,8 @@\n     gradient_checkpointing: Optional[bool] = field(\n         default=True, metadata={\"help\": \"whether to use gradient checkpointing\"}\n     )\n-    group_by_length: Optional[bool] = field(default=True, metadata={\"help\": \"whether to group by length\"})\n+    group_by_length: Optional[bool] = field(default=False, metadata={\"help\": \"whether to group by length\"})\n+    packing: Optional[bool] = field(default=True, metadata={\"help\": \"whether to use packing for SFTTrainer\"})\n \n     lora_alpha: Optional[float] = field(default=16, metadata={\"help\": \"the lora alpha parameter\"})\n     lora_dropout: Optional[float] = field(default=0.05, metadata={\"help\": \"the lora dropout parameter\"})\n@@ -54,6 +55,9 @@\n parser = HfArgumentParser(ScriptArguments)\n script_args = parser.parse_args_into_dataclasses()[0]\n \n+if script_args.group_by_length and script_args.packing:\n+    raise ValueError(\"Cannot use both packing and group by length\")\n+\n \n def chars_token_ratio(dataset, tokenizer, nb_examples=400):\n     \"\"\"\n@@ -189,7 +193,7 @@\n     train_dataset=train_dataset,\n     eval_dataset=eval_dataset,\n     peft_config=peft_config,\n-    packing=True,\n+    packing=script_args.packing,\n     max_seq_length=None,\n     tokenizer=tokenizer,\n     args=training_args,\n", "issue": "sft_llama2 error\nHello! \r\n\r\nI made no changes to the sample code or the dataset. I just simply wanted to get it to run the stacked llama2 example. However, I get this error:\r\n\r\nUserWarning: The passed formatting_func has more than one argument. Usually that function should have a single argument `example` which corresponds to the dictionary returned by each element of the dataset. Make sure you know what you are doing.\r\n\r\n in __init__\r\n    raise ValueError(\"the `--group_by_length` option is only available for `Dataset`, not `IterableDataset\")\r\nValueError: the `--group_by_length` option is only available for `Dataset`, not `IterableDataset\r\n\r\n(Sorry, I must omit certain parts of the error message since I am working on an institution's server). \r\n\r\nThank you!\n", "before_files": [{"content": "# Fine-Tune Llama2-7b on SE paired dataset\nimport os\nfrom dataclasses import dataclass, field\nfrom typing import Optional\n\nimport torch\nfrom datasets import load_dataset\nfrom peft import AutoPeftModelForCausalLM, LoraConfig\nfrom tqdm import tqdm\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, HfArgumentParser, TrainingArguments\n\nfrom trl import SFTTrainer\nfrom trl.trainer import ConstantLengthDataset\n\n\n@dataclass\nclass ScriptArguments:\n    model_name: Optional[str] = field(default=\"meta-llama/Llama-2-7b-hf\", metadata={\"help\": \"the model name\"})\n    log_with: Optional[str] = field(default=\"wandb\", metadata={\"help\": \"use 'wandb' to log with wandb\"})\n\n    dataset_name: Optional[str] = field(default=\"lvwerra/stack-exchange-paired\", metadata={\"help\": \"the dataset name\"})\n    subset: Optional[str] = field(default=\"data/finetune\", metadata={\"help\": \"the subset to use\"})\n    split: Optional[str] = field(default=\"train\", metadata={\"help\": \"the split to use\"})\n    size_valid_set: Optional[int] = field(default=4000, metadata={\"help\": \"the size of the validation set\"})\n    streaming: Optional[bool] = field(default=True, metadata={\"help\": \"whether to stream the dataset\"})\n    shuffle_buffer: Optional[int] = field(default=5000, metadata={\"help\": \"the shuffle buffer size\"})\n    seq_length: Optional[int] = field(default=1024, metadata={\"help\": \"the sequence length\"})\n\n    max_steps: Optional[int] = field(default=500, metadata={\"help\": \"the maximum number of sgd steps\"})\n    logging_steps: Optional[int] = field(default=10, metadata={\"help\": \"the logging frequency\"})\n    save_steps: Optional[int] = field(default=10, metadata={\"help\": \"the saving frequency\"})\n    per_device_train_batch_size: Optional[int] = field(default=4, metadata={\"help\": \"the per device train batch size\"})\n    per_device_eval_batch_size: Optional[int] = field(default=1, metadata={\"help\": \"the per device eval batch size\"})\n    gradient_accumulation_steps: Optional[int] = field(default=2, metadata={\"help\": \"the gradient accumulation steps\"})\n    gradient_checkpointing: Optional[bool] = field(\n        default=True, metadata={\"help\": \"whether to use gradient checkpointing\"}\n    )\n    group_by_length: Optional[bool] = field(default=True, metadata={\"help\": \"whether to group by length\"})\n\n    lora_alpha: Optional[float] = field(default=16, metadata={\"help\": \"the lora alpha parameter\"})\n    lora_dropout: Optional[float] = field(default=0.05, metadata={\"help\": \"the lora dropout parameter\"})\n    lora_r: Optional[int] = field(default=8, metadata={\"help\": \"the lora r parameter\"})\n\n    learning_rate: Optional[float] = field(default=1e-4, metadata={\"help\": \"the learning rate\"})\n    lr_scheduler_type: Optional[str] = field(default=\"cosine\", metadata={\"help\": \"the lr scheduler type\"})\n    num_warmup_steps: Optional[int] = field(default=100, metadata={\"help\": \"the number of warmup steps\"})\n    weight_decay: Optional[float] = field(default=0.05, metadata={\"help\": \"the weight decay\"})\n    optimizer_type: Optional[str] = field(default=\"paged_adamw_32bit\", metadata={\"help\": \"the optimizer type\"})\n\n    output_dir: Optional[str] = field(default=\"./results\", metadata={\"help\": \"the output directory\"})\n    log_freq: Optional[int] = field(default=1, metadata={\"help\": \"the logging frequency\"})\n\n\nparser = HfArgumentParser(ScriptArguments)\nscript_args = parser.parse_args_into_dataclasses()[0]\n\n\ndef chars_token_ratio(dataset, tokenizer, nb_examples=400):\n    \"\"\"\n    Estimate the average number of characters per token in the dataset.\n    \"\"\"\n    total_characters, total_tokens = 0, 0\n    for _, example in tqdm(zip(range(nb_examples), iter(dataset)), total=nb_examples):\n        text = prepare_sample_text(example)\n        total_characters += len(text)\n        if tokenizer.is_fast:\n            total_tokens += len(tokenizer(text).tokens())\n        else:\n            total_tokens += len(tokenizer.tokenize(text))\n\n    return total_characters / total_tokens\n\n\ndef print_trainable_parameters(model):\n    \"\"\"\n    Prints the number of trainable parameters in the model.\n    \"\"\"\n    trainable_params = 0\n    all_param = 0\n    for _, param in model.named_parameters():\n        all_param += param.numel()\n        if param.requires_grad:\n            trainable_params += param.numel()\n    print(\n        f\"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}\"\n    )\n\n\ndef prepare_sample_text(example):\n    \"\"\"Prepare the text from a sample of the dataset.\"\"\"\n    text = f\"Question: {example['question']}\\n\\nAnswer: {example['response_j']}\"\n    return text\n\n\ndef create_datasets(tokenizer, args):\n    dataset = load_dataset(\n        args.dataset_name,\n        data_dir=args.subset,\n        split=args.split,\n        use_auth_token=True,\n        num_proc=args.num_workers if not args.streaming else None,\n        streaming=args.streaming,\n    )\n    if args.streaming:\n        print(\"Loading the dataset in streaming mode\")\n        valid_data = dataset.take(args.size_valid_set)\n        train_data = dataset.skip(args.size_valid_set)\n        train_data = train_data.shuffle(buffer_size=args.shuffle_buffer, seed=None)\n    else:\n        dataset = dataset.train_test_split(test_size=0.005, seed=None)\n        train_data = dataset[\"train\"]\n        valid_data = dataset[\"test\"]\n        print(f\"Size of the train set: {len(train_data)}. Size of the validation set: {len(valid_data)}\")\n\n    chars_per_token = chars_token_ratio(train_data, tokenizer)\n    print(f\"The character to token ratio of the dataset is: {chars_per_token:.2f}\")\n\n    train_dataset = ConstantLengthDataset(\n        tokenizer,\n        train_data,\n        formatting_func=prepare_sample_text,\n        infinite=True,\n        seq_length=args.seq_length,\n        chars_per_token=chars_per_token,\n    )\n    valid_dataset = ConstantLengthDataset(\n        tokenizer,\n        valid_data,\n        formatting_func=prepare_sample_text,\n        infinite=False,\n        seq_length=args.seq_length,\n        chars_per_token=chars_per_token,\n    )\n    return train_dataset, valid_dataset\n\n\nbnb_config = BitsAndBytesConfig(\n    load_in_4bit=True,\n    bnb_4bit_quant_type=\"nf4\",\n    bnb_4bit_compute_dtype=torch.bfloat16,\n)\n\nbase_model = AutoModelForCausalLM.from_pretrained(\n    script_args.model_name,\n    quantization_config=bnb_config,\n    device_map={\"\": 0},\n    trust_remote_code=True,\n    use_auth_token=True,\n)\nbase_model.config.use_cache = False\n\npeft_config = LoraConfig(\n    r=script_args.lora_r,\n    lora_alpha=script_args.lora_alpha,\n    lora_dropout=script_args.lora_dropout,\n    target_modules=[\"q_proj\", \"v_proj\"],\n    bias=\"none\",\n    task_type=\"CAUSAL_LM\",\n)\n\ntokenizer = AutoTokenizer.from_pretrained(script_args.model_name, trust_remote_code=True)\ntokenizer.pad_token = tokenizer.eos_token\ntokenizer.padding_side = \"right\"  # Fix weird overflow issue with fp16 training\n\n\ntraining_args = TrainingArguments(\n    output_dir=script_args.output_dir,\n    per_device_train_batch_size=script_args.per_device_train_batch_size,\n    gradient_accumulation_steps=script_args.gradient_accumulation_steps,\n    per_device_eval_batch_size=script_args.per_device_eval_batch_size,\n    learning_rate=script_args.learning_rate,\n    logging_steps=script_args.logging_steps,\n    max_steps=script_args.max_steps,\n    report_to=script_args.log_with,\n    save_steps=script_args.save_steps,\n    group_by_length=script_args.group_by_length,\n    lr_scheduler_type=script_args.lr_scheduler_type,\n    warmup_steps=script_args.num_warmup_steps,\n    optim=script_args.optimizer_type,\n    bf16=True,\n    remove_unused_columns=False,\n    run_name=\"sft_llama2\",\n)\n\ntrain_dataset, eval_dataset = create_datasets(tokenizer, script_args)\n\ntrainer = SFTTrainer(\n    model=base_model,\n    train_dataset=train_dataset,\n    eval_dataset=eval_dataset,\n    peft_config=peft_config,\n    packing=True,\n    max_seq_length=None,\n    tokenizer=tokenizer,\n    args=training_args,\n)\ntrainer.train()\ntrainer.save_model(script_args.output_dir)\n\noutput_dir = os.path.join(script_args.output_dir, \"final_checkpoint\")\ntrainer.model.save_pretrained(output_dir)\n\n# Free memory for merging weights\ndel base_model\ntorch.cuda.empty_cache()\n\nmodel = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map=\"auto\", torch_dtype=torch.bfloat16)\nmodel = model.merge_and_unload()\n\noutput_merged_dir = os.path.join(script_args.output_dir, \"final_merged_checkpoint\")\nmodel.save_pretrained(output_merged_dir, safe_serialization=True)\n", "path": "examples/research_projects/stack_llama_2/scripts/sft_llama2.py"}], "after_files": [{"content": "# Fine-Tune Llama2-7b on SE paired dataset\nimport os\nfrom dataclasses import dataclass, field\nfrom typing import Optional\n\nimport torch\nfrom datasets import load_dataset\nfrom peft import AutoPeftModelForCausalLM, LoraConfig\nfrom tqdm import tqdm\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, HfArgumentParser, TrainingArguments\n\nfrom trl import SFTTrainer\nfrom trl.trainer import ConstantLengthDataset\n\n\n@dataclass\nclass ScriptArguments:\n    model_name: Optional[str] = field(default=\"meta-llama/Llama-2-7b-hf\", metadata={\"help\": \"the model name\"})\n    log_with: Optional[str] = field(default=\"wandb\", metadata={\"help\": \"use 'wandb' to log with wandb\"})\n\n    dataset_name: Optional[str] = field(default=\"lvwerra/stack-exchange-paired\", metadata={\"help\": \"the dataset name\"})\n    subset: Optional[str] = field(default=\"data/finetune\", metadata={\"help\": \"the subset to use\"})\n    split: Optional[str] = field(default=\"train\", metadata={\"help\": \"the split to use\"})\n    size_valid_set: Optional[int] = field(default=4000, metadata={\"help\": \"the size of the validation set\"})\n    streaming: Optional[bool] = field(default=True, metadata={\"help\": \"whether to stream the dataset\"})\n    shuffle_buffer: Optional[int] = field(default=5000, metadata={\"help\": \"the shuffle buffer size\"})\n    seq_length: Optional[int] = field(default=1024, metadata={\"help\": \"the sequence length\"})\n\n    max_steps: Optional[int] = field(default=500, metadata={\"help\": \"the maximum number of sgd steps\"})\n    logging_steps: Optional[int] = field(default=10, metadata={\"help\": \"the logging frequency\"})\n    save_steps: Optional[int] = field(default=10, metadata={\"help\": \"the saving frequency\"})\n    per_device_train_batch_size: Optional[int] = field(default=4, metadata={\"help\": \"the per device train batch size\"})\n    per_device_eval_batch_size: Optional[int] = field(default=1, metadata={\"help\": \"the per device eval batch size\"})\n    gradient_accumulation_steps: Optional[int] = field(default=2, metadata={\"help\": \"the gradient accumulation steps\"})\n    gradient_checkpointing: Optional[bool] = field(\n        default=True, metadata={\"help\": \"whether to use gradient checkpointing\"}\n    )\n    group_by_length: Optional[bool] = field(default=False, metadata={\"help\": \"whether to group by length\"})\n    packing: Optional[bool] = field(default=True, metadata={\"help\": \"whether to use packing for SFTTrainer\"})\n\n    lora_alpha: Optional[float] = field(default=16, metadata={\"help\": \"the lora alpha parameter\"})\n    lora_dropout: Optional[float] = field(default=0.05, metadata={\"help\": \"the lora dropout parameter\"})\n    lora_r: Optional[int] = field(default=8, metadata={\"help\": \"the lora r parameter\"})\n\n    learning_rate: Optional[float] = field(default=1e-4, metadata={\"help\": \"the learning rate\"})\n    lr_scheduler_type: Optional[str] = field(default=\"cosine\", metadata={\"help\": \"the lr scheduler type\"})\n    num_warmup_steps: Optional[int] = field(default=100, metadata={\"help\": \"the number of warmup steps\"})\n    weight_decay: Optional[float] = field(default=0.05, metadata={\"help\": \"the weight decay\"})\n    optimizer_type: Optional[str] = field(default=\"paged_adamw_32bit\", metadata={\"help\": \"the optimizer type\"})\n\n    output_dir: Optional[str] = field(default=\"./results\", metadata={\"help\": \"the output directory\"})\n    log_freq: Optional[int] = field(default=1, metadata={\"help\": \"the logging frequency\"})\n\n\nparser = HfArgumentParser(ScriptArguments)\nscript_args = parser.parse_args_into_dataclasses()[0]\n\nif script_args.group_by_length and script_args.packing:\n    raise ValueError(\"Cannot use both packing and group by length\")\n\n\ndef chars_token_ratio(dataset, tokenizer, nb_examples=400):\n    \"\"\"\n    Estimate the average number of characters per token in the dataset.\n    \"\"\"\n    total_characters, total_tokens = 0, 0\n    for _, example in tqdm(zip(range(nb_examples), iter(dataset)), total=nb_examples):\n        text = prepare_sample_text(example)\n        total_characters += len(text)\n        if tokenizer.is_fast:\n            total_tokens += len(tokenizer(text).tokens())\n        else:\n            total_tokens += len(tokenizer.tokenize(text))\n\n    return total_characters / total_tokens\n\n\ndef print_trainable_parameters(model):\n    \"\"\"\n    Prints the number of trainable parameters in the model.\n    \"\"\"\n    trainable_params = 0\n    all_param = 0\n    for _, param in model.named_parameters():\n        all_param += param.numel()\n        if param.requires_grad:\n            trainable_params += param.numel()\n    print(\n        f\"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}\"\n    )\n\n\ndef prepare_sample_text(example):\n    \"\"\"Prepare the text from a sample of the dataset.\"\"\"\n    text = f\"Question: {example['question']}\\n\\nAnswer: {example['response_j']}\"\n    return text\n\n\ndef create_datasets(tokenizer, args):\n    dataset = load_dataset(\n        args.dataset_name,\n        data_dir=args.subset,\n        split=args.split,\n        use_auth_token=True,\n        num_proc=args.num_workers if not args.streaming else None,\n        streaming=args.streaming,\n    )\n    if args.streaming:\n        print(\"Loading the dataset in streaming mode\")\n        valid_data = dataset.take(args.size_valid_set)\n        train_data = dataset.skip(args.size_valid_set)\n        train_data = train_data.shuffle(buffer_size=args.shuffle_buffer, seed=None)\n    else:\n        dataset = dataset.train_test_split(test_size=0.005, seed=None)\n        train_data = dataset[\"train\"]\n        valid_data = dataset[\"test\"]\n        print(f\"Size of the train set: {len(train_data)}. Size of the validation set: {len(valid_data)}\")\n\n    chars_per_token = chars_token_ratio(train_data, tokenizer)\n    print(f\"The character to token ratio of the dataset is: {chars_per_token:.2f}\")\n\n    train_dataset = ConstantLengthDataset(\n        tokenizer,\n        train_data,\n        formatting_func=prepare_sample_text,\n        infinite=True,\n        seq_length=args.seq_length,\n        chars_per_token=chars_per_token,\n    )\n    valid_dataset = ConstantLengthDataset(\n        tokenizer,\n        valid_data,\n        formatting_func=prepare_sample_text,\n        infinite=False,\n        seq_length=args.seq_length,\n        chars_per_token=chars_per_token,\n    )\n    return train_dataset, valid_dataset\n\n\nbnb_config = BitsAndBytesConfig(\n    load_in_4bit=True,\n    bnb_4bit_quant_type=\"nf4\",\n    bnb_4bit_compute_dtype=torch.bfloat16,\n)\n\nbase_model = AutoModelForCausalLM.from_pretrained(\n    script_args.model_name,\n    quantization_config=bnb_config,\n    device_map={\"\": 0},\n    trust_remote_code=True,\n    use_auth_token=True,\n)\nbase_model.config.use_cache = False\n\npeft_config = LoraConfig(\n    r=script_args.lora_r,\n    lora_alpha=script_args.lora_alpha,\n    lora_dropout=script_args.lora_dropout,\n    target_modules=[\"q_proj\", \"v_proj\"],\n    bias=\"none\",\n    task_type=\"CAUSAL_LM\",\n)\n\ntokenizer = AutoTokenizer.from_pretrained(script_args.model_name, trust_remote_code=True)\ntokenizer.pad_token = tokenizer.eos_token\ntokenizer.padding_side = \"right\"  # Fix weird overflow issue with fp16 training\n\n\ntraining_args = TrainingArguments(\n    output_dir=script_args.output_dir,\n    per_device_train_batch_size=script_args.per_device_train_batch_size,\n    gradient_accumulation_steps=script_args.gradient_accumulation_steps,\n    per_device_eval_batch_size=script_args.per_device_eval_batch_size,\n    learning_rate=script_args.learning_rate,\n    logging_steps=script_args.logging_steps,\n    max_steps=script_args.max_steps,\n    report_to=script_args.log_with,\n    save_steps=script_args.save_steps,\n    group_by_length=script_args.group_by_length,\n    lr_scheduler_type=script_args.lr_scheduler_type,\n    warmup_steps=script_args.num_warmup_steps,\n    optim=script_args.optimizer_type,\n    bf16=True,\n    remove_unused_columns=False,\n    run_name=\"sft_llama2\",\n)\n\ntrain_dataset, eval_dataset = create_datasets(tokenizer, script_args)\n\ntrainer = SFTTrainer(\n    model=base_model,\n    train_dataset=train_dataset,\n    eval_dataset=eval_dataset,\n    peft_config=peft_config,\n    packing=script_args.packing,\n    max_seq_length=None,\n    tokenizer=tokenizer,\n    args=training_args,\n)\ntrainer.train()\ntrainer.save_model(script_args.output_dir)\n\noutput_dir = os.path.join(script_args.output_dir, \"final_checkpoint\")\ntrainer.model.save_pretrained(output_dir)\n\n# Free memory for merging weights\ndel base_model\ntorch.cuda.empty_cache()\n\nmodel = AutoPeftModelForCausalLM.from_pretrained(output_dir, device_map=\"auto\", torch_dtype=torch.bfloat16)\nmodel = model.merge_and_unload()\n\noutput_merged_dir = os.path.join(script_args.output_dir, \"final_merged_checkpoint\")\nmodel.save_pretrained(output_merged_dir, safe_serialization=True)\n", "path": "examples/research_projects/stack_llama_2/scripts/sft_llama2.py"}]}