""" Production QLoRA fine-tuning pipeline for WorldDisasterLM. Hardware requirements --------------------- Minimum : 1× NVIDIA A100 40 GB (or 2× RTX 3090/4090 24 GB with gradient checkpointing) Preferred: 2–8× A100/H100 80 GB for full batch sizes Consumer : 1× RTX 4090 24 GB → set per_device_train_batch_size=1, gradient_accumulation_steps=16 Prerequisites ------------- 1. Accept Meta Llama 3.1 license on Hugging Face: https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct 2. Set HF_TOKEN in your .env file: huggingface-cli login (or export HF_TOKEN=hf_...) Run --- python scripts/train_production.py --dataset data/processed/instruction_dataset.jsonl """ from __future__ import annotations import json import logging import os from dataclasses import dataclass, field from pathlib import Path from worlddisasterlm.training.chat_format import apply_template logger = logging.getLogger(__name__) @dataclass class QLoRAConfig: # ── Model ───────────────────────────────────────────────────────────────── base_model: str = "meta-llama/Llama-3.1-8B-Instruct" output_dir: str = "checkpoints/worlddisasterlm-qlora" # ── Data ────────────────────────────────────────────────────────────────── dataset_path: str = "data/processed/instruction_dataset.jsonl" max_seq_length: int = 2048 # ── QLoRA / PEFT ────────────────────────────────────────────────────────── use_4bit: bool = True bnb_4bit_quant_type: str = "nf4" # "nf4" or "fp4" use_double_quant: bool = True compute_dtype: str = "bfloat16" # "bfloat16" or "float16" lora_r: int = 16 lora_alpha: int = 32 lora_dropout: float = 0.05 lora_target_modules: list[str] = field( default_factory=lambda: [ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ] ) # ── Optimisation ────────────────────────────────────────────────────────── epochs: int = 3 learning_rate: float = 2e-4 per_device_train_batch_size: int = 2 gradient_accumulation_steps: int = 8 warmup_ratio: float = 0.03 max_grad_norm: float = 0.3 lr_scheduler_type: str = "cosine" weight_decay: float = 0.0 bf16: bool = True fp16: bool = False gradient_checkpointing: bool = True # ── Logging & saving ────────────────────────────────────────────────────── logging_steps: int = 10 save_steps: int = 200 save_total_limit: int = 3 report_to: str = "none" # "mlflow" | "wandb" | "none" # ── Misc ────────────────────────────────────────────────────────────────── seed: int = 42 def load_jsonl(path: str) -> list[dict]: records = [] with open(path, encoding="utf-8") as handle: for line in handle: line = line.strip() if line: records.append(json.loads(line)) return records def build_dataset(path: str, tokenizer): from datasets import Dataset raw = load_jsonl(path) logger.info("Loaded %d instruction samples from %s", len(raw), path) texts = [] for sample in raw: try: text = apply_template( tokenizer, instruction=sample.get("instruction", ""), context=sample.get("input", ""), output=sample.get("output", ""), ) texts.append(text) except Exception as exc: logger.warning("Failed to format sample: %s", exc) logger.info("Formatted %d samples with chat template", len(texts)) return Dataset.from_dict({"text": texts}) def setup_model_tokenizer(config: QLoRAConfig): import torch from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training # ── Tokenizer ───────────────────────────────────────────────────────────── tokenizer = AutoTokenizer.from_pretrained( config.base_model, trust_remote_code=True, padding_side="right", ) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.pad_token_id = tokenizer.eos_token_id # ── Quantization ────────────────────────────────────────────────────────── torch_dtype = torch.bfloat16 if config.bf16 else torch.float16 bnb_config = None if config.use_4bit: bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type=config.bnb_4bit_quant_type, bnb_4bit_compute_dtype=torch_dtype, bnb_4bit_use_double_quant=config.use_double_quant, ) # ── Base model ──────────────────────────────────────────────────────────── model = AutoModelForCausalLM.from_pretrained( config.base_model, quantization_config=bnb_config, device_map="auto", torch_dtype=torch_dtype, trust_remote_code=True, attn_implementation="flash_attention_2" if _flash_attn_available() else "eager", ) if config.use_4bit: model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=config.gradient_checkpointing) elif config.gradient_checkpointing: model.gradient_checkpointing_enable() model.config.use_cache = False # ── LoRA ────────────────────────────────────────────────────────────────── peft_config = LoraConfig( task_type=TaskType.CAUSAL_LM, r=config.lora_r, lora_alpha=config.lora_alpha, target_modules=config.lora_target_modules, lora_dropout=config.lora_dropout, bias="none", inference_mode=False, ) model = get_peft_model(model, peft_config) model.print_trainable_parameters() return model, tokenizer def _flash_attn_available() -> bool: try: import flash_attn # noqa: F401 return True except ImportError: return False def train(config: QLoRAConfig) -> None: from trl import SFTConfig, SFTTrainer # Seed from worlddisasterlm.utils.seed import seed_everything seed_everything(config.seed) Path(config.output_dir).mkdir(parents=True, exist_ok=True) # MLflow / W&B config if config.report_to == "mlflow" and os.getenv("MLFLOW_TRACKING_URI"): import mlflow mlflow.set_tracking_uri(os.environ["MLFLOW_TRACKING_URI"]) mlflow.set_experiment("worlddisasterlm-training") model, tokenizer = setup_model_tokenizer(config) train_dataset = build_dataset(config.dataset_path, tokenizer) report_to_list = [config.report_to] if config.report_to != "none" else [] sft_config = SFTConfig( output_dir=config.output_dir, num_train_epochs=config.epochs, per_device_train_batch_size=config.per_device_train_batch_size, gradient_accumulation_steps=config.gradient_accumulation_steps, gradient_checkpointing=config.gradient_checkpointing, learning_rate=config.learning_rate, max_grad_norm=config.max_grad_norm, warmup_ratio=config.warmup_ratio, lr_scheduler_type=config.lr_scheduler_type, weight_decay=config.weight_decay, bf16=config.bf16, fp16=config.fp16, logging_steps=config.logging_steps, save_steps=config.save_steps, save_total_limit=config.save_total_limit, load_best_model_at_end=False, max_seq_length=config.max_seq_length, dataset_text_field="text", packing=False, report_to=report_to_list, seed=config.seed, ) trainer = SFTTrainer( model=model, tokenizer=tokenizer, train_dataset=train_dataset, args=sft_config, ) logger.info("Starting QLoRA fine-tuning: %d samples, %d epochs", len(train_dataset), config.epochs) trainer.train() logger.info("Saving adapter to %s", config.output_dir) trainer.save_model(config.output_dir) tokenizer.save_pretrained(config.output_dir) logger.info("Training complete.")