Update train.py

train.py

@@ -1,30 +1,28 @@
-from transformers import MT5Tokenizer, MT5ForConditionalGeneration, Trainer, TrainingArguments
-from transformers import ByT5Tokenizer, T5ForConditionalGeneration
+from transformers import T5ForConditionalGeneration, ByT5Tokenizer, Trainer, TrainingArguments
 from datasets import load_dataset
 import os
 import wandb
+import torch
 
-# 🔧 Название запуска (используется и как run_name, и как output_dir)
-run_name = "byt5-finetuning-run"
+# 🔧 Название модели и путь
+model_name = "google/byt5-small"
+run_id = "byt5-autobatch"
+output_dir = f"./{run_id}"
+start_batch_size = 300
+step_batch_size = 5
 
-# 🧠 Название модели для фантюнинга
-model_id = "google/byt5-small"
+# 📦 Загружаем модель и токенизатор
+model = T5ForConditionalGeneration.from_pretrained(model_name)
+tokenizer = ByT5Tokenizer.from_pretrained(model_name)
 
-# 📂 Куда сохранять результат обучения
-output_dir = f"./{run_name}"
-
-# Загружаем модель и токенизатор
-model = T5ForConditionalGeneration.from_pretrained(model_id)
-tokenizer = ByT5Tokenizer.from_pretrained(model_id)
-
-# Загружаем датасет
+# 📂 Загружаем датасет
 data_files = {
     "train": "mt5_training_data-1.jsonl",
     "validation": "mt5_validation_data-1.jsonl"
 }
 dataset = load_dataset("json", data_files=data_files)
 
-# Токенизация
+# 🔠 Токенизация
 def tokenize_function(examples):
     model_inputs = tokenizer(examples["text"], max_length=256, truncation=True, padding="max_length")
     labels = tokenizer(examples["target"], max_length=64, truncation=True, padding="max_length")
@@ -33,45 +31,60 @@ def tokenize_function(examples):
 
 tokenized_datasets = dataset.map(tokenize_function, batched=True)
 
-# Авторизация в Weights & Biases
+# 🔑 Авторизация W&B
 wandb.login(key="5f028bc0142fb7fa45bdacdde3c00dbbaf8bf98e")
 
-# Аргументы обучения
-training_args = TrainingArguments(
-    output_dir=output_dir,
-    evaluation_strategy="steps",
-    eval_steps=100,
-    learning_rate=5e-5,
-    per_device_train_batch_size=200,
-    per_device_eval_batch_size=200,
-    fp16=True,
-    num_train_epochs=3,
-    logging_steps=100,
-    warmup_ratio=0.06,
-    logging_first_step=True,
-    weight_decay=0.01,
-    logging_dir="./logs",
-    save_total_limit=2,
-    save_strategy="epoch",
-    report_to="wandb",
-    run_name=run_name,
-    disable_tqdm=False,
-    max_grad_norm=1.0
-)
+# 🚀 Функция автоподбора batch size
+def try_training_with_batch_size(batch_size_start):
+    batch_size = batch_size_start
+    while batch_size > 0:
+        try:
+            print(f"\n🚀 Пробуем batch_size = {batch_size}")
+            training_args = TrainingArguments(
+                output_dir=output_dir,
+                evaluation_strategy="steps",
+                eval_steps=100,
+                learning_rate=5e-5,
+                per_device_train_batch_size=batch_size,
+                per_device_eval_batch_size=batch_size,
+                fp16=True,
+                num_train_epochs=3,
+                logging_steps=100,
+                warmup_ratio=0.06,
+                logging_first_step=True,
+                weight_decay=0.01,
+                logging_dir="./logs",
+                save_total_limit=2,
+                save_strategy="epoch",
+                report_to="wandb",
+                run_name=run_id,
+                disable_tqdm=False,
+                max_grad_norm=1.0
+            )
+
+            trainer = Trainer(
+                model=model,
+                args=training_args,
+                train_dataset=tokenized_datasets["train"],
+                eval_dataset=tokenized_datasets["validation"]
+            )
 
-# Инициализируем Trainer
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=tokenized_datasets["train"],
-    eval_dataset=tokenized_datasets["validation"]
-)
+            trainer.train()
+            return batch_size
+        except RuntimeError as e:
+            if "CUDA out of memory" in str(e):
+                print(f"❌ OOM на batch_size = {batch_size}, уменьшаем...")
+                torch.cuda.empty_cache()
+                batch_size -= step_batch_size
+            else:
+                raise e
+    raise RuntimeError("Не удалось подобрать подходящий batch size 😢")
 
-# Обучение
-trainer.train()
-# trainer.train(resume_from_checkpoint=True)
+# 🏁 Запуск с автоподбором
+optimal_batch_size = try_training_with_batch_size(start_batch_size)
+print(f"\n✅ Успешно обучено с batch_size = {optimal_batch_size}")
 
-# Сохраняем модель
+# 💾 Сохраняем модель
 model.save_pretrained(output_dir)
 tokenizer.save_pretrained(output_dir)
-print(f"✅ Модель сохранена локально в {output_dir}")
+print(f"✅ Модель сохранена в {output_dir}")