Fine-tune Llama 3.1 8B на RTX 4090

TL;DR

Fine-tuning Llama 3.1 8B с LoRA на арендованном RTX 4090 стоит $0,61/час и использует 21,4 ГБ из 24 ГБ VRAM карты (89,3% использования). Датасет из 10k samples завершается примерно за 2 часа. Общие затраты на 3-epoch прохода составят около $1,22-$1,55, включая передачу данных.

Почему LoRA вместо полного fine-tuning?

Hu et al. (2021) представили LoRA как способ внедрения обучаемых матриц низкого ранга в attention layers при замораживании всего остального. Для Llama 3.1 8B это снижает количество обучаемых параметров с 8,03 миллиардов до примерно 8 миллионов - 0,1% от оригинала.

Практический результат - уменьшение VRAM в 10-15 раз. Полный fine-tuning на 24 ГБ карте требует gradient checkpointing только чтобы избежать OOM, что снижает пропускную способность. LoRA подходит без этого.

Для domain adaptation стандартная точка старта - r=8 и alpha=16. Больший ранг (r=16 или r=32) добавляет выразительную мощность и требует больше VRAM. Для большинства задач r=8 достаточно, и catastrophic forgetting не является реальной проблемой в этом масштабе.

Требования VRAM и размер batch

Базовая модель Llama 3.1 8B загруженная в 8-bit занимает примерно 16 ГБ согласно карте модели Meta Llama 3.1. Добавьте состояния LoRA adapter плюс optimizer states и вы получите около 21,4 ГБ в сумме. На 24 ГБ RTX 4090 это 89,3% использования. Плотно, но стабильно.

Начните с per_device_train_batch_size=4. Это безопасный стандарт на 24 ГБ карте с 8-bit загрузкой. Переход на 4-bit через bitsandbytes позволит вам довести до 8 без проблем.

RTX 4090 является практическим минимумом для этой модели. A100 40 GB дает больше простора но стоит в 3-5 раз дороже в час на большинстве платформ.

Настройка LoRA config

Установите стек:

pip install torch transformers peft datasets accelerate bitsandbytes

Проверьте CUDA:

python -c "import torch; print(torch.cuda.is_available())"

Загрузите модель в 8-bit, затем оберните её с LoRA config. Согласно документации Hugging Face PEFT, LoRA adapters остаются в float32 даже когда базовая модель квантизирована.

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, TaskType

model_id = "meta-llama/Llama-3.1-8B"

bnb_config = BitsAndBytesConfig(load_in_8bit=True)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto"
)

lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type=TaskType.CAUSAL_LM
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# trainable params: 8,388,608 || all params: 8,038,957,056 || trainable%: 0.10

target_modules - это самое важное решение здесь. q_proj и v_proj - это стандартный выбор. Добавление k_proj и o_proj поднимает обучаемые параметры до примерно 16 M и помогает на задачах требующих больше структурных изменений, ценой немного больше VRAM.

Размер датасета vs. время обучения

Датасет из 10k samples с 5-10 M tokens обучается примерно 2 часа при batch size 4 на RTX 4090.

Отложите 5-10% как validation set. Проверяйте eval loss каждые 500 шагов - на датасете из 10k это примерно один раз за epoch, чего достаточно чтобы рано поймать overfitting.

from datasets import load_dataset
dataset = load_dataset("tatsu-lab/alpaca")

Данные custom domain работают быстрее если хорошо отформатированы. JSONL файл с одним полем "text" за строку - это самая простая структура которую PEFT и Transformers обрабатывают без предварительной обработки.

Стоимость аренды RTX 4090

On-demand RTX 4090 инстанции работают $0,50-$0,70/час (GPUniq pricing). Spot pricing на AWS и GCP составляет $0,30-$0,40/час, но spot инстанции могут быть прерваны во время выполнения.

2-часовой прохода при $0,61/час стоит $1,22. Добавьте примерно $0,10 за загрузку 5-10 ГБ модели с Hugging Face. Итого: примерно $1,32.

Покупка RTX 4090 обойдется в $1600-$2000. При $0,61/час вам нужно более 2600 часов GPU времени прежде чем покупка окупится. Для случайных fine-tuning экспериментов аренда - очевидный выбор.

Быстрая проверка convergence

Следите за training loss каждые 100 шагов. К концу epoch 1 вы должны видеть падение на 30-50% от начального значения. Если нет - сначала проверьте learning rate.

Используйте perplexity на validation set вместо запуска полного inference. Это примерно в 10 раз быстрее и дает тот же сигнал для decisions о convergence.

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    output_dir="./lora_output",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    learning_rate=2e-4,
    save_steps=500,
    eval_steps=500,
    evaluation_strategy="steps",
    logging_steps=100,
    fp16=True,
    report_to="tensorboard"
)

Установите eval_steps=500 и вы получите значение каждые пол-epoch на датасете из 10k. Остановите обучение если validation loss останавливается на 2 последовательных epoch. На меньших custom датасетах это может произойти быстро.

Частые точки отказа

CUDA OOM. Порядок исправления: сначала уменьшите batch size до 2, затем переключитесь с 8-bit на 4-bit загрузку, затем включите gradient checkpointing как последний ресурс (это стоит пропускной способности).

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16
)

Медленное обучение обычно означает что GPU использование низко. Запустите nvidia-smi в отдельном терминале. Ниже 80% и ваш data pipeline - узкое место - увеличьте num_workers в DataLoader или предварительно обработайте и кэшируйте датасет перед началом обучения.

NaN loss почти всегда означает что learning rate слишком высок. Начните с 2e-4 для LoRA. NaN в первые 50 шагов - снизьте до 1e-4. Полный fine-tuning обычно использует 1e-5 до 5e-5; LoRA терпит более высокие rates потому что обновляется только 0,1% весов.

Еще одно: примите лицензию модели на Hugging Face перед загрузкой. Пропущенное согласие дает запутанную ошибку 403, а не четкое сообщение.

Какие open датасеты обучаются быстрее?

Данные custom domain побеждают по скорости. 5-10k хорошо отформатированных samples может завершиться за час и превзойдет generic датасет на вашей конкретной задаче.

Без domain данных Alpaca - практический default. 52k QA пары, domain-agnostic, 1-2 часа на RTX 4090. OpenAssistant имеет лучшее разнообразие и более длинные разговоры, но ожидайте 4-6 часов на полный прохода.

Отформатируйте custom данные как JSONL с полем "text":

{"text": "### Instruction:\nSummarize the following document.\n\n### Input:\n...\n\n### Response:\n..."}

Держите формат последовательным во всем датасете. Несоответствие заставляет модель тратить обучающую мощность на изучение шаблона вместо содержания.

Стоимость за epoch

cost = hourly_rate * wall_clock_hours_per_epoch

При $0,61/час и 30 минутах за epoch на датасете из 10k, каждый epoch стоит $0,31.

Чтобы получить фактическую пропускную способность, логируйте tokens per second во время обучения. Умножьте на 3600 чтобы получить tokens per hour, затем разделите количество tokens вашего датасета на это число чтобы получить часов за epoch. Калькулятор стоимости GPUniq принимает измеренную пропускную способность напрямую.

Прекратите инстанцию сразу после сохранения adapter. Забыть отключить - это как $1,22 эксперимент превращается в $15.

Полный training прохода, шаг за шагом

1. Арендуйте RTX 4090 инстанцию

Зарегистрируйтесь на GPUniq, Lambda Labs или Vast.ai. Запустите on-demand RTX 4090 инстанцию с Ubuntu 22.04 и CUDA 12.1. SSH и проверьте:

nvidia-smi
# Should show: NVIDIA GeForce RTX 4090, 24576MiB

2. Установите dependencies

pip install torch transformers peft datasets accelerate bitsandbytes
python -c "import torch; print(torch.cuda.is_available())"
# True

3. Загрузите Llama 3.1 8B и датасет

Примите лицензию на huggingface.co/meta-llama/Llama-3.1-8B сначала.

huggingface-cli login

from datasets import load_dataset
dataset = load_dataset("tatsu-lab/alpaca")

4. Настройте LoRA и training

Используйте config блоки из разделов выше. Ключевые значения: r=8, lora_alpha=16, target_modules=["q_proj", "v_proj"], learning_rate=2e-4, per_device_train_batch_size=4.

5. Запустите обучение

python train.py
# In a second terminal:
watch -n 2 nvidia-smi

Ожидайте падения loss на 30-50% к epoch 1. Общее wall clock на 10k samples: примерно 2 часа.

6. Сохраните и перезагрузите adapter

model.save_pretrained("./lora_output")

from peft import PeftModel
base_model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
model = PeftModel.from_pretrained(base_model, "./lora_output")

7. Рассчитайте стоимость и отключите

Умножьте wall-clock часы на вашу ставку аренды. Добавьте примерно $0,10 за 5 ГБ egress. Затем прекратите инстанцию. Биллинг останавливается когда инстанция удалена, не когда вы вышли.