Spaces:
Sleeping
Sleeping
File size: 2,843 Bytes
dbacab4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | from __future__ import annotations
import math
from typing import Any
import torch
from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
from transformers import BitsAndBytesConfig, get_cosine_schedule_with_warmup
def str_to_torch_dtype(name: str | None) -> torch.dtype | None:
if name is None:
return None
value = str(name).lower()
if value in {"bfloat16", "bf16"}:
return torch.bfloat16
if value in {"float16", "fp16"}:
return torch.float16
if value in {"float32", "fp32"}:
return torch.float32
raise ValueError(f"Unsupported torch dtype: {name}")
def quantization_config_from_config(config: dict[str, Any]) -> BitsAndBytesConfig | None:
if not config.get("load_in_4bit", False):
return None
return BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type=config.get("bnb_4bit_quant_type", "nf4"),
bnb_4bit_compute_dtype=str_to_torch_dtype(config.get("bnb_4bit_compute_dtype", "bfloat16")),
bnb_4bit_use_double_quant=bool(config.get("bnb_4bit_use_double_quant", True)),
)
def apply_lora(model: torch.nn.Module, config: dict[str, Any]) -> torch.nn.Module:
if config.get("load_in_4bit", False):
model = prepare_model_for_kbit_training(
model,
use_gradient_checkpointing=bool(config.get("gradient_checkpointing", True)),
)
lora_config = LoraConfig(
r=int(config.get("lora_r", 16)),
lora_alpha=int(config.get("lora_alpha", 32)),
lora_dropout=float(config.get("lora_dropout", 0.05)),
target_modules=list(config.get("target_modules", [])),
bias="none",
task_type=TaskType.CAUSAL_LM,
)
return get_peft_model(model, lora_config)
def build_optimizer(model: torch.nn.Module, config: dict[str, Any]) -> torch.optim.Optimizer:
trainable = [p for p in model.parameters() if p.requires_grad]
if not trainable:
raise ValueError("No trainable parameters found. Check LoRA target_modules.")
return torch.optim.AdamW(
trainable,
lr=float(config.get("learning_rate", 1e-4)),
weight_decay=float(config.get("weight_decay", 0.0)),
)
def build_scheduler(optimizer: torch.optim.Optimizer, config: dict[str, Any], steps_per_epoch: int):
epochs = int(config.get("num_train_epochs", 1))
total_steps = max(1, steps_per_epoch * epochs)
warmup_steps = math.ceil(total_steps * float(config.get("warmup_ratio", 0.03)))
return get_cosine_schedule_with_warmup(optimizer, warmup_steps, total_steps)
def oom_help() -> str:
return (
"CUDA out of memory. Try reducing max_frames, train_max_samples, or "
"per_device_train_batch_size; use a smaller model_name; keep load_in_4bit true; "
"or run on a GPU with more VRAM."
)
|