WIP for YuisekinAI

Author: yuiseki

Makefile

@@ -18,10 +18,10 @@ eval-all: $(targets)
 	python3 src/eval.py recipes/RTX_3060_12GB/sql-coder.yaml
 
 output/tinyllama-color-coder-v1/checkpoint-200/README.md:
-	python3 src/train.py recipes/RTX_3060_12GB/color-coder.yaml
+	accelerate launch src/train.py recipes/RTX_3060_12GB/color-coder.yaml
 
 output/tinyllama-sql-coder-v1/checkpoint-200/README.md:
-	python3 src/train.py recipes/RTX_3060_12GB/sql-coder.yaml
+	accelerate launch src/train.py recipes/RTX_3060_12GB/sql-coder.yaml
 
 docker:
 	docker build --no-cache -t yuiseki/infinite-tinyllama:latest .

README.md

@@ -40,6 +40,10 @@ VRAM 24GB で日本語データセットでファインチューニングしよ
 conda create -n peft
 ```
 
+```bash
+conda install -c nvidia cuda-toolkit=12.1
+```
+
 ```bash
 conda activate peft
 ```

recipes/A5000_24GB_x8/YuisekinAI-mistral-1.1B-aya.yaml

@@ -0,0 +1,30 @@
+target_task: tasks/i18n/ja.md
+base_model_id: yuiseki/YuisekinAI-mistral-1.1B
+model_name: YuisekinAI-mistral-1.1B-aya
+output_base_dir: /data/output
+dataset_id: CohereForAI/aya_dataset
+dataset_input_field_name: inputs
+dataset_output_field_name: targets
+dataset_filter_field_name: language_code
+dataset_filter_field_value: jpn
+dataset_train_split_seed: 42
+dataset_train_split_test_size: 0.2
+lora_r: 8
+lora_alpha: 16
+lora_dropout: 0.05
+train_claim_gpu_num: 1
+train_per_device_train_batch_size: 4
+train_gradient_accumulation_steps: 32
+train_num_train_epochs: 4
+train_fp16: True
+inference_max_new_tokens: 32
+evaluations:
+  -
+    prompt: "火縄銃の威力が全国に知られる事となった、１５７５年に織田・徳川連合軍が鉄砲隊を用いて武田勝頼率いる騎馬隊を破った戦いを何というでしょう？"
+    expected_output: "長篠の戦いです。"
+  -
+    prompt: "ベトナム戦争終結や米ソ戦略兵器削減交渉などを進めたものの、１９７４年にウォーターゲート事件の責任をとって辞任したアメリカの第３７代大統領は誰でしょう？"
+    expected_output: "リチャード・ニクソンです。"
+  -
+    prompt: "格闘家ボブ・サップの出身国はどこでしょう？"
+    expected_output: "アメリカです。"

recipes/A5000_24GB_x8/i18n-ja-wikibook-hs.yaml

@@ -0,0 +1,28 @@
+target_task: tasks/i18n/ja.md
+base_model_id: TinyLlama/TinyLlama-1.1B-intermediate-step-715k-1.5T
+model_name: tinyllama-ja-wikibook-hs-v0.1
+output_base_dir: /data/output
+dataset_id: DataPilot/wikibook_High_School_textbooks
+dataset_input_field_name: text
+dataset_train_split_seed: 42
+dataset_train_split_test_size: 0.2
+lora_r: 8
+lora_alpha: 16
+lora_dropout: 0.05
+train_claim_gpu_num: 4
+train_per_device_train_batch_size: 8
+train_gradient_accumulation_steps: 4
+train_num_train_epochs: 4
+train_max_steps: 2000
+train_fp16: True
+inference_max_new_tokens: 32
+evaluations:
+  -
+    prompt: "２つ以上の文字を組み合わせて図案化したもののことで、特にルイ・ヴィトンのものが知られるのは何でしょう？"
+    expected_output: "モノグラム"
+  -
+    prompt: "幾つかの布切れを縫いあわせ、飾りや模様を作る手芸方法を何というでしょう？"
+    expected_output: "パッチワーク"
+  -
+    prompt: "格闘家ボブ・サップの出身国はどこでしょう？"
+    expected_output: "アメリカ"

recipes/A5000_24GB_x8/i18n-ja-wikibook-jhs.yaml

@@ -0,0 +1,28 @@
+target_task: tasks/i18n/ja.md
+base_model_id: TinyLlama/TinyLlama-1.1B-intermediate-step-715k-1.5T
+model_name: tinyllama-ja-wikibook-jhs-v0.1
+output_base_dir: /data/output
+dataset_id: DataPilot/wikibook_Junior_High_School_textbooks_ja
+dataset_input_field_name: text
+dataset_train_split_seed: 42
+dataset_train_split_test_size: 0.2
+lora_r: 8
+lora_alpha: 16
+lora_dropout: 0.05
+train_claim_gpu_num: 4
+train_per_device_train_batch_size: 8
+train_gradient_accumulation_steps: 4
+train_num_train_epochs: 4
+train_max_steps: 2000
+train_fp16: True
+inference_max_new_tokens: 32
+evaluations:
+  -
+    prompt: "２つ以上の文字を組み合わせて図案化したもののことで、特にルイ・ヴィトンのものが知られるのは何でしょう？"
+    expected_output: "モノグラム"
+  -
+    prompt: "幾つかの布切れを縫いあわせ、飾りや模様を作る手芸方法を何というでしょう？"
+    expected_output: "パッチワーク"
+  -
+    prompt: "格闘家ボブ・サップの出身国はどこでしょう？"
+    expected_output: "アメリカ"

requirements.txt

@@ -4,5 +4,6 @@ bitsandbytes
 transformers
 trl
 wandb
+flash_attn
 
 types-PyYAML

src/train.py

@@ -1,8 +1,9 @@
 import os
 import sys
 
-import wandb
+import torch
 import yaml
+from accelerate import PartialState
 from datasets.arrow_dataset import Dataset
 from datasets.load import load_dataset
 from peft import LoraConfig
@@ -14,6 +15,8 @@
 )
 from trl import SFTTrainer
 
+import wandb
+
 os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
 
 
@@ -176,24 +179,35 @@ def load_model_and_tokenizer(model_id):
     # tokenizer.add_special_tokens({'pad_token': '[PAD]'})
     # NOTE: tokenizer.add_special_tokensやるならこれは不要
     tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.padding_side = "right"
 
     # Define the quantization configuration for memory-efficient training.
     bnb_config = BitsAndBytesConfig(
         # Load the model weights in 4-bit quantized format.
         load_in_4bit=True,
+        # Specify whether to use double quantization for 4-bit quantization.
+        bnb_4bit_use_double_quant=True,
         # Specify the quantization type to use for 4-bit quantization.
         bnb_4bit_quant_type="nf4",
         # Specify the data type to use for computations during training.
-        bnb_4bit_compute_dtype="float16",
-        # Specify whether to use double quantization for 4-bit quantization.
-        bnb_4bit_use_double_quant=True,
+        bnb_4bit_compute_dtype=torch.float16,
     )
     # Load the model from the specified model ID and apply the quantization configuration.
+
     model = AutoModelForCausalLM.from_pretrained(
+        # Base model id
         model_id,
+        # BitsAndBytes configuration
         quantization_config=bnb_config,
+        # Set torch dtype
+        torch_dtype=torch.float16,
+        # Trust remote code
         trust_remote_code=True,
-        device_map="auto",
+        # Set device map to auto
+        # device_map="auto",
+        device_map={"": PartialState().process_index},
+        # Set the attention impl
+        attn_implementation="flash_attention_2",
     )
     # Disable cache to improve training speed.
     model.config.use_cache = False
@@ -222,12 +236,6 @@ def load_model_and_tokenizer(model_id):
 os.environ["WANDB_PROJECT"] = "infinite-tinyllama"
 os.environ["WANDB_LOG_MODEL"] = "false"
 os.environ["WANDB_WATCH"] = "all"
-wandb.init(
-    project="infinite-tinyllama",
-    name=train_config["model_name"],
-    group=train_config["model_name"],
-    config=train_config,
-)
 
 #
 # Define LoRA and PEFT config
@@ -249,12 +257,11 @@ def load_model_and_tokenizer(model_id):
     optim="paged_adamw_32bit",
     learning_rate=2e-4,
     lr_scheduler_type="cosine",
-    save_strategy="steps",
-    save_steps=100,
+    save_strategy="epoch",
     logging_steps=10,
     num_train_epochs=int(train_config["train_num_train_epochs"]),
-    max_steps=int(train_config["train_max_steps"]),
     fp16=True,
+    run_name=train_config["model_name"],
 )
 
 trainer = SFTTrainer(