Add end-of-training model, README, tokenizer

README.md

@@ -1,3 +1,17 @@
 ---
-license: mit
+library_name: peft
+base_model: google/gemma-2b
+widget:
+- text: "Salut, ce zi minunata pentru"
 ---
+
+## Model details
+
+A (decent) try at fine-tuning a Gemma 2B model on about ~1.6GB of high-quality
+Romanian.
+
+All the scripts used + data are available in this repo.
+
+### Framework versions
+
+- PEFT 0.9.0

adapter_config.json

@@ -0,0 +1,33 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "google/gemma-2b",
+  "bias": "none",
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 32,
+  "lora_dropout": 0.1,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 8,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "down_proj",
+    "up_proj",
+    "o_proj",
+    "v_proj",
+    "gate_proj",
+    "k_proj",
+    "q_proj"
+  ],
+  "task_type": "CAUSAL_LM",
+  "use_dora": false,
+  "use_rslora": false
+}

adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:55ffb0e9dd4622929f14a38e560e535b02ffdd2da430c5bd6597af450619e38a
+size 39256456

fine-tuning.py

@@ -0,0 +1,127 @@
+# start with torchrun --nproc-per-node <n-gpu's> fine-tuning.py
+import os
+
+import torch
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    DataCollatorForLanguageModeling,
+    TrainingArguments,
+    Trainer,
+    BitsAndBytesConfig,
+    TrainerCallback,
+)
+from datasets import load_from_disk
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from peft.tuners.lora import LoraLayer
+from accelerate import Accelerator
+
+
+batch_size = 2
+
+checkpoint = "google/gemma-2b"
+data_dir = "dataset_ro_small_v1/"
+save_dir = "gemma-2b-romanian-1.6gb-finetuned-qlora"
+log_dir = "training_logs/"
+
+# load dataset
+tokenized_datasets = load_from_disk(f'tokenized_{data_dir}')
+
+tokenized_datasets = tokenized_datasets.shuffle(seed=42)
+
+print(tokenized_datasets)
+
+# load quantized model
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_quant_dtype=torch.float16,
+    bnb_4bit_compute_dtype=torch.float16,
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    checkpoint,
+    load_in_8bit=False,
+    quantization_config=bnb_config,
+    device_map={ "": Accelerator().process_index }, # see https://github.com/huggingface/trl/issues/1348
+    torch_dtype=torch.float16,
+    trust_remote_code=True,
+    attn_implementation='sdpa',#'flash_attention_2',
+    use_cache=False,
+)
+model = prepare_model_for_kbit_training(model)
+
+# load qlora config
+lora_config = LoraConfig(
+    lora_alpha=32,
+    lora_dropout=0.1,
+    r=8,
+    bias="none",
+    task_type="CAUSAL_LM",
+    target_modules=["q_proj", "o_proj", "k_proj", "v_proj", "gate_proj", "up_proj", "down_proj"],
+)
+model = get_peft_model(model, lora_config)
+
+model.print_trainable_parameters()
+
+# load tokenizer from checkpoint
+tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+
+tokenizer.pad_token = tokenizer.eos_token
+data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+# training args
+args = TrainingArguments(
+    output_dir='training_checkpoints/',
+    logging_dir=log_dir,
+    per_device_train_batch_size=batch_size,
+    per_device_eval_batch_size=batch_size,
+    evaluation_strategy='no',
+    logging_steps=100,
+    save_strategy='steps',
+    save_steps=100,
+    save_total_limit=10,
+    gradient_accumulation_steps=4,
+    gradient_checkpointing=True,
+    gradient_checkpointing_kwargs={ "use_reentrant": False },
+    num_train_epochs=1,
+    warmup_steps=1_000,
+    weight_decay=0.001,
+    lr_scheduler_type='cosine',
+    learning_rate=1e-4,
+    max_grad_norm=0.3,
+    fp16=True,
+    ddp_find_unused_parameters=False,
+)
+
+# stop the training loop after 1000 updates
+class StopCallback(TrainerCallback):
+    def on_step_end(self, args, state, control, **kwargs):
+        if state.global_step != 0 and state.global_step % 1000 == 0:
+            # stop training
+            control.should_training_stop = True
+
+# train as usual
+trainer = Trainer(
+    model=model,
+    args=args,
+    data_collator=data_collator,
+    train_dataset=tokenized_datasets['train'],
+    eval_dataset=tokenized_datasets['test'],
+    tokenizer=tokenizer,
+)
+trainer.add_callback(StopCallback)
+
+print("Starting training...")
+
+train_checkpoint = os.getenv("TRAIN_CHECKPOINT")
+if train_checkpoint is not None:
+    trainer.train(train_checkpoint) # resume training from checkpoint dir
+else:
+    trainer.train()
+
+# save trainer state at end
+torch.save(trainer.state.log_history, "trainer_log_history.pth")
+
+model.save_pretrained(save_dir)
+tokenizer.save_pretrained(save_dir)

inference.py

@@ -0,0 +1,31 @@
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+)
+from peft import PeftModel, PeftConfig
+import torch
+
+orig_checkpoint = 'google/gemma-2b'
+checkpoint = 'checkpoint-4000'
+HF_TOKEN = ''
+PROMPT = 'Salut, ca sa imi schimb buletinul pot sa'
+
+seq_len = 2048
+
+# load original model first
+tokenizer = AutoTokenizer.from_pretrained(orig_checkpoint, token=HF_TOKEN)
+
+config = PeftConfig.from_pretrained(checkpoint)
+model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, token=HF_TOKEN)
+
+# then merge trained QLoRA weights
+model = PeftModel.from_pretrained(model, checkpoint)
+model.merge_and_unload()
+
+model = model.cuda()
+
+# generate normally
+inputs = tokenizer.encode(PROMPT, return_tensors="pt").cuda()
+outputs = model.generate(inputs, max_new_tokens=seq_len)
+
+print(tokenizer.decode(outputs[0]))

preprocessing.py

@@ -0,0 +1,38 @@
+from transformers import AutoTokenizer
+from datasets import load_dataset
+
+
+checkpoint = "google/gemma-2b"
+data_dir = "dataset_ro_small_v1/"
+
+seq_len = 2048
+
+raw_datasets = load_dataset("json", data_dir=data_dir, split='train')
+raw_datasets = raw_datasets.remove_columns(['url', 'date_download', 'digest',
+                                            'length', 'nlines', 'source_domain',
+                                            'title', 'cc_segment',
+                                            'original_nlines',
+                                            'original_length', 'line_ids',
+                                            'language', 'language_score'])
+raw_datasets = raw_datasets.rename_column('raw_content', 'text')
+raw_datasets = raw_datasets.train_test_split(test_size=0.1)
+
+print(raw_datasets)
+
+# load tokenizer from checkpoint
+tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+
+def tokenize_fn(examples):
+    return tokenizer(examples['text'],
+                     max_length=seq_len,
+                     return_overflowing_tokens=True,
+                     truncation=True)
+
+tokenizer.pad_token = tokenizer.eos_token
+tokenized_datasets = raw_datasets.map(
+    tokenize_fn,
+    batched=True,
+    remove_columns=raw_datasets['train'].column_names
+)
+
+tokenized_datasets.save_to_disk(f'tokenized_{data_dir}')

special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<bos>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<eos>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<eos>",
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,49 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<eos>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<bos>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<bos>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<eos>",
+  "legacy": null,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<eos>",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "GemmaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}