from __future__ import annotations import json import os from dataclasses import dataclass from pathlib import Path from typing import Any import torch from datasets import Dataset from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training from transformers import ( AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, Trainer, TrainerCallback, TrainingArguments, set_seed, ) from transformers.trainer_utils import get_last_checkpoint from .common import ARTIFACT_ROOT, CHECKPOINT_ROOT, DATASET_ROOT, atomic_json, ensure_dirs, log_event, update_status BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen3-1.7B") MAX_LENGTH = int(os.environ.get("MAX_LENGTH", "2048")) OUTPUT_NAME = os.environ.get("OUTPUT_NAME", "adapter") TRAIN_LIMIT = int(os.environ.get("TRAIN_LIMIT", "0")) MAX_STEPS = int(os.environ.get("MAX_STEPS", "-1")) SEED = int(os.environ.get("TRAIN_SEED", "42")) def load_rows(path: Path, limit: int = 0) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] with path.open("r", encoding="utf-8") as handle: for line in handle: rows.append(json.loads(line)) if limit and len(rows) >= limit: break return rows def common_prefix_length(left: list[int], right: list[int]) -> int: length = 0 for a, b in zip(left, right): if a != b: break length += 1 return length def as_token_ids(value: Any) -> list[int]: if isinstance(value, dict) or hasattr(value, "keys"): value = value["input_ids"] if isinstance(value, torch.Tensor): value = value.tolist() if value and isinstance(value[0], list): value = value[0] return list(value) def tokenized_dataset(tokenizer: Any, rows: list[dict[str, Any]]) -> Dataset: dataset = Dataset.from_list([{"messages": row["messages"], "id": row["id"]} for row in rows]) def encode(example: dict[str, Any]) -> dict[str, Any]: messages = example["messages"] prefix = messages[:-1] prefix_ids = as_token_ids( tokenizer.apply_chat_template( prefix, tokenize=True, add_generation_prompt=True, enable_thinking=False, ) ) full_ids = as_token_ids( tokenizer.apply_chat_template( messages, tokenize=True, add_generation_prompt=False, enable_thinking=False, ) ) if tokenizer.eos_token_id is not None and (not full_ids or full_ids[-1] != tokenizer.eos_token_id): full_ids.append(tokenizer.eos_token_id) mask_to = common_prefix_length(prefix_ids, full_ids) input_ids = full_ids[:MAX_LENGTH] labels = [-100] * min(mask_to, len(input_ids)) + input_ids[min(mask_to, len(input_ids)) :] if not any(label != -100 for label in labels): raise ValueError(f"Assistant response was fully truncated for example {example['id']}") return {"input_ids": input_ids, "attention_mask": [1] * len(input_ids), "labels": labels} encoded = dataset.map(encode, remove_columns=dataset.column_names, desc="Tokenizing conversations") packed: list[dict[str, list[int]]] = [] current = {"input_ids": [], "attention_mask": [], "labels": []} for item in encoded: item_length = len(item["input_ids"]) if item_length > MAX_LENGTH: raise ValueError("Tokenized example exceeded MAX_LENGTH after truncation") if current["input_ids"] and len(current["input_ids"]) + item_length > MAX_LENGTH: packed.append(current) current = {"input_ids": [], "attention_mask": [], "labels": []} current["input_ids"].extend(item["input_ids"]) current["attention_mask"].extend(item["attention_mask"]) current["labels"].extend(item["labels"]) if current["input_ids"]: packed.append(current) return Dataset.from_list(packed) @dataclass class AssistantOnlyCollator: pad_token_id: int def __call__(self, features: list[dict[str, list[int]]]) -> dict[str, torch.Tensor]: max_len = max(len(feature["input_ids"]) for feature in features) input_ids = [] attention_mask = [] labels = [] for feature in features: padding = max_len - len(feature["input_ids"]) input_ids.append(feature["input_ids"] + [self.pad_token_id] * padding) attention_mask.append(feature["attention_mask"] + [0] * padding) labels.append(feature["labels"] + [-100] * padding) return { "input_ids": torch.tensor(input_ids, dtype=torch.long), "attention_mask": torch.tensor(attention_mask, dtype=torch.long), "labels": torch.tensor(labels, dtype=torch.long), } class StatusCallback(TrainerCallback): def on_log(self, args: TrainingArguments, state: Any, control: Any, logs: dict[str, Any] | None = None, **_: Any) -> None: if not logs: return numeric = {key: round(float(value), 6) for key, value in logs.items() if isinstance(value, (int, float))} update_status( "smoke_train" if OUTPUT_NAME == "smoke_adapter" else "full_train", f"Training step {state.global_step}", global_step=state.global_step, max_steps=state.max_steps, metrics=numeric, ) def main() -> None: ensure_dirs() if not torch.cuda.is_available(): raise RuntimeError("QLoRA training requires the requested CUDA GPU") set_seed(SEED) output_dir = CHECKPOINT_ROOT / OUTPUT_NAME adapter_dir = ARTIFACT_ROOT / OUTPUT_NAME output_dir.mkdir(parents=True, exist_ok=True) adapter_dir.mkdir(parents=True, exist_ok=True) tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True, trust_remote_code=False) if tokenizer.pad_token_id is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "right" train_rows = load_rows(DATASET_ROOT / "train.jsonl", TRAIN_LIMIT) eval_limit = min(200, max(30, len(train_rows) // 10)) if TRAIN_LIMIT else 500 eval_rows = load_rows(DATASET_ROOT / "validation.jsonl", eval_limit) train_dataset = tokenized_dataset(tokenizer, train_rows) eval_dataset = tokenized_dataset(tokenizer, eval_rows) quantization = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16, ) model = AutoModelForCausalLM.from_pretrained( BASE_MODEL, quantization_config=quantization, torch_dtype=torch.bfloat16, device_map={"": 0}, trust_remote_code=False, attn_implementation="sdpa", ) model.config.use_cache = False model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True) lora = LoraConfig( r=int(os.environ.get("LORA_R", "16")), lora_alpha=int(os.environ.get("LORA_ALPHA", "32")), lora_dropout=float(os.environ.get("LORA_DROPOUT", "0.05")), bias="none", task_type="CAUSAL_LM", target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], ) model = get_peft_model(model, lora) model.print_trainable_parameters() save_strategy = os.environ.get("SAVE_STRATEGY", "steps") smoke = OUTPUT_NAME == "smoke_adapter" args = TrainingArguments( output_dir=str(output_dir), overwrite_output_dir=False, num_train_epochs=float(os.environ.get("EPOCHS", "2")), max_steps=MAX_STEPS, per_device_train_batch_size=int(os.environ.get("BATCH_SIZE", "4")), per_device_eval_batch_size=int(os.environ.get("EVAL_BATCH_SIZE", "4")), gradient_accumulation_steps=int(os.environ.get("GRAD_ACCUM", "4")), learning_rate=float(os.environ.get("LEARNING_RATE", "1e-4")), lr_scheduler_type="cosine", warmup_ratio=0.03, weight_decay=0.0, max_grad_norm=1.0, optim="paged_adamw_8bit", bf16=True, fp16=False, tf32=True, logging_steps=1 if smoke else 10, eval_strategy="no" if smoke else "steps", eval_steps=100, save_strategy=save_strategy, save_steps=100, save_total_limit=3, load_best_model_at_end=not smoke, metric_for_best_model="eval_loss", greater_is_better=False, report_to="none", remove_unused_columns=False, dataloader_num_workers=2, dataloader_pin_memory=True, gradient_checkpointing=True, seed=SEED, data_seed=SEED, ) trainer = Trainer( model=model, args=args, train_dataset=train_dataset, eval_dataset=None if smoke else eval_dataset, data_collator=AssistantOnlyCollator(tokenizer.pad_token_id), callbacks=[StatusCallback()], ) last_checkpoint = get_last_checkpoint(str(output_dir)) if output_dir.exists() and not smoke else None update_status( "smoke_train" if smoke else "full_train", "Model loaded; optimizer starting", base_model=BASE_MODEL, train_examples=len(train_dataset), validation_examples=len(eval_dataset), resume_from=last_checkpoint, gpu=torch.cuda.get_device_name(0), ) result = trainer.train(resume_from_checkpoint=last_checkpoint) trainer.save_metrics("train", result.metrics) model.save_pretrained(adapter_dir, safe_serialization=True) tokenizer.save_pretrained(adapter_dir) metrics = {key: float(value) for key, value in result.metrics.items() if isinstance(value, (int, float))} atomic_json(adapter_dir / "run_metrics.json", metrics) atomic_json( adapter_dir / "training_config.json", { "base_model": BASE_MODEL, "method": "QLoRA NF4 double-quantization with BF16 compute", "assistant_only_loss": True, "packing": True, "thinking": False, "max_length": MAX_LENGTH, "seed": SEED, "source_train_examples": len(train_rows), "source_validation_examples": len(eval_rows), "packed_train_sequences": len(train_dataset), "packed_validation_sequences": len(eval_dataset), "lora_r": lora.r, "lora_alpha": lora.lora_alpha, }, ) log_event("training_complete", output=OUTPUT_NAME, metrics=metrics) update_status( "smoke_complete" if smoke else "training_complete", f"{OUTPUT_NAME} training completed", metrics=metrics, adapter=str(adapter_dir), ) if __name__ == "__main__": main()