Update 1_finetune_mixtral.py

1_finetune_mixtral.py

@@ -3,27 +3,24 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from peft import LoraConfig
 from trl import SFTTrainer
 from datasets import load_dataset
-import config as cfg # Импортируем наш конфиг
+import config as cfg
 
 def main():
     # --- 1. Загрузка датасета ---
-    # Мы сохранили данные в jsonl, так что можем загрузить их так:
     try:
-        dataset = load_dataset("json", data_files={"train": cfg.OUTPUT_TRAIN_FILE}, split="train")
-        # Если у вас есть и валидационный набор:
-        # dataset = load_dataset("json", data_files={"train": "data/train.jsonl", "validation": "data/validation.jsonl"})
+        dataset = load_dataset("json", data_files={"train": cfg.TRAINING_DATA_JSONL}, split="train")
     except Exception as e:
-        print(f"Ошибка загрузки датасета: {e}")
-        print("Убедитесь, что файл data/train_dataset_llm.jsonl существует и не пуст, и что скрипт 0_prepare_data.py успешно отработал.")
+        print(f"Ошибка загрузки датасета из {cfg.TRAINING_DATA_JSONL}: {e}")
+        print("Убедитесь, что файл существует и не пуст, и что скрипт 0_prepare_data.py успешно отработал.")
         return
         
-    if not dataset:
+    if not dataset or len(dataset) == 0:
         print("Датасет не загружен или пуст. Прерывание.")
         return
     
     print(f"Загружен датасет с {len(dataset)} примерами.")
-    print("Пример первого элемента датасета:", dataset[0])
-
+    if len(dataset) > 0:
+        print("Пример первого элемента датасета:", dataset[0])
 
     # --- 2. Конфигурация квантизации (BitsAndBytes) ---
     bnb_config = BitsAndBytesConfig(
@@ -38,25 +35,23 @@ def main():
         cfg.BASE_MODEL_NAME,
         quantization_config=bnb_config,
         torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
-        device_map="auto", # Распределит модель по доступным GPU
+        device_map="auto",
         trust_remote_code=True,
-        # attn_implementation="flash_attention_2" # Если установлено и поддерживается
+        # attn_implementation="flash_attention_2" # Если flash-attn установлен
     )
-    model.config.use_cache = False # Важно для PEFT
+    model.config.use_cache = False
 
     tokenizer = AutoTokenizer.from_pretrained(cfg.BASE_MODEL_NAME, trust_remote_code=True)
     tokenizer.pad_token = tokenizer.eos_token
     tokenizer.padding_side = "right"
 
     # --- 4. Конфигурация LoRA ---
-    # target_modules для Mixtral могут включать 'q_proj', 'k_proj', 'v_proj', 'o_proj',
-    # 'gate_proj', 'up_proj', 'down_proj'. Начните с основных для проекций внимания.
-    lora_target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"]
+    lora_target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"] #, "gate_proj", "up_proj", "down_proj"]
 
     peft_config = LoraConfig(
         lora_alpha=16,
-        lora_dropout=0.05,
-        r=8,
+        lora_dropout=0.05, # Было 0.1, уменьшил для возможной борьбы с переобучением на малых данных
+        r=8, # или 16
         bias="none",
         task_type="CAUSAL_LM",
         target_modules=lora_target_modules
@@ -67,20 +62,19 @@ def main():
         output_dir=cfg.OUTPUT_DIR,
         per_device_train_batch_size=cfg.TRAIN_BATCH_SIZE,
         gradient_accumulation_steps=cfg.GRAD_ACCUMULATION_STEPS,
-        optim="paged_adamw_32bit", # Экономит память
+        optim="paged_adamw_32bit",
         learning_rate=cfg.LEARNING_RATE,
         num_train_epochs=cfg.NUM_EPOCHS,
-        # max_steps=100, # Для быстрого теста, потом установите -1 для использования num_train_epochs
         lr_scheduler_type="cosine",
         warmup_ratio=0.03,
         logging_steps=10,
         save_strategy="epoch",
-        # evaluation_strategy="epoch", # Если есть eval_dataset
-        # load_best_model_at_end=True, # Если есть eval_dataset
-        fp16=not torch.cuda.is_bf16_supported(), # Используйте fp16 если bfloat16 не доступен
+        fp16=not torch.cuda.is_bf16_supported(),
         bf16=torch.cuda.is_bf16_supported(),
         gradient_checkpointing=True,
-        report_to="tensorboard", # или "wandb"
+        report_to="tensorboard",
+        # evaluation_strategy="epoch", # Если есть eval_dataset
+        # load_best_model_at_end=True, # Если есть eval_dataset
     )
 
     # --- 6. Инициализация SFTTrainer ---
@@ -88,12 +82,11 @@ def main():
         model=model,
         tokenizer=tokenizer,
         args=training_args,
-        train_dataset=dataset, # dataset.select(range(100)) для теста на малом подмножестве
-        # eval_dataset=dataset["validation"], # Если есть
+        train_dataset=dataset,
         peft_config=peft_config,
-        dataset_text_field="text", # Название колонки с текстом в вашем датасете
+        dataset_text_field="text",
         max_seq_length=cfg.MAX_SEQ_LENGTH,
-        packing=False, # Упаковка может быть полезна для коротких последовательностей, но начните без нее
+        packing=False, # С вашими данными, вероятно, лучше False
     )
 
     # --- 7. Запуск обучения ---
@@ -102,40 +95,18 @@ def main():
         trainer.train()
     except Exception as e:
         print(f"Ошибка во время обучения: {e}")
-        print("Возможные причины: нехватка VRAM (уменьшите batch_size, max_seq_length, LoRA r), проблемы с данными.")
         return
 
     # --- 8. Сохранение адаптера LoRA ---
-    trainer.save_model(cfg.FINETUNED_ADAPTER_PATH) # Сохраняем адаптер
-    # tokenizer.save_pretrained(cfg.FINETUNED_ADAPTER_PATH) # Токенизатор тоже можно сохранить рядом
+    trainer.save_model(cfg.FINETUNED_ADAPTER_PATH)
     print(f"Обучение завершено. Адаптер LoRA сохранен в: {cfg.FINETUNED_ADAPTER_PATH}")
 
-    # (Опционально) Слияние и сохранение полной модели
-    # Это потребует больше RAM/VRAM
-    # print("Слияние модели...")
-    # merged_model = model.merge_and_unload() # Если использовали get_peft_model
-    # Если SFTTrainer сам создал PeftModel, то нужно сначала получить базовую модель и PeftModel
-    # base_model_for_merge = AutoModelForCausalLM.from_pretrained(
-    #    cfg.BASE_MODEL_NAME,
-    #    torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
-    #    device_map="cpu", # Сливаем на CPU, если VRAM мало
-    #    trust_remote_code=True
-    # )
-    # merged_model = PeftModel.from_pretrained(base_model_for_merge, cfg.FINETUNED_ADAPTER_PATH)
-    # merged_model = merged_model.merge_and_unload()
-    # merged_model.save_pretrained(cfg.MERGED_MODEL_PATH)
-    # tokenizer.save_pretrained(cfg.MERGED_MODEL_PATH)
-    # print(f"Смерженная модель сохранена в: {cfg.MERGED_MODEL_PATH}")
-
-
 if __name__ == "__main__":
     if not torch.cuda.is_available():
-        print("CUDA недоступна. Обучение на CPU будет невозможным или крайне медленным для Mixtral.")
+        print("CUDA недоступна.")
     else:
         print(f"Доступно CUDA устройств: {torch.cuda.device_count()}")
         print(f"Текущее устройство CUDA: {torch.cuda.current_device()} ({torch.cuda.get_device_name(torch.cuda.current_device())})")
-        if torch.cuda.is_bf16_supported():
-            print("BF16 поддерживается.")
-        else:
-            print("BF16 НЕ поддерживается. Будет использоваться FP16 (если включено) или FP32.")
+        if torch.cuda.is_bf16_supported(): print("BF16 поддерживается.")
+        else: print("BF16 НЕ поддерживается.")
         main()