diff --git a/fine_tuning/alignment/README.md b/fine_tuning/alignment/README.md new file mode 100644 index 0000000..255576e --- /dev/null +++ b/fine_tuning/alignment/README.md @@ -0,0 +1,69 @@ +# Preference Alignment Techniques + +This directory contains implementations of various preference alignment techniques for language model fine-tuning. + +## Available Methods + +### DPO (Direct Preference Optimization) +- **Location**: `dpo/` +- **Description**: Direct preference optimization without requiring separate reward models +- **Use case**: Stable, efficient preference alignment with human feedback data +- **Key advantages**: + - Simpler than PPO-based RLHF + - More computationally efficient + - Better stability during training + +## Directory Structure + +``` +alignment/ +├── README.md # This file +└── dpo/ # Direct Preference Optimization implementation + ├── dpo.py # Main DPO training script + ├── sft.py # Supervised fine-tuning for DPO initialization + ├── compare_base_vs_dpo.py # Model comparison utilities + ├── dpo_pipeline_sanity.sh # End-to-end pipeline + ├── Dockerfile # Container setup + ├── Makefile # Build system + └── README.md # DPO-specific documentation +``` + +## Comparison with RLHF + +Traditional RLHF (in `../rlhf/`) uses reinforcement learning with separate reward models, while alignment techniques here focus on direct optimization approaches that are often more efficient and stable. + +| Aspect | RLHF (PPO) | Alignment (DPO) | +|--------|------------|-----------------| +| Reward Model | Required | Not needed | +| Training Stability | Can be unstable | More stable | +| Computational Cost | Higher | Lower | +| Implementation Complexity | Complex | Simpler | +| Hyperparameter Sensitivity | High | Lower | + +## Getting Started + +For DPO training, see the detailed documentation in `dpo/README.md`. + +Example quick start (interactive shell + manual pipeline): +```bash +cd dpo/ +make build +make run HF_TOKEN=<> MODEL_NAME=google/gemma-3-270m + +# Inside container (if not auto-run) +./dpo_pipeline_sanity.sh google/gemma-3-270m +``` + +Automated pipeline targets (recommended): +```bash +# Standard speed +make run-pipeline HF_TOKEN=<> MODEL_NAME=google/gemma-3-270m + +# Fast demo mode (<~20 min) with reduced steps & samples +make run-pipeline-fast HF_TOKEN=<> MODEL_NAME=google/gemma-3-270m +``` + +Fast mode details: +- Enables internal `--demo_fast` overrides in `sft.py`, `dpo.py`, `compare_base_vs_dpo.py`. +- Shrinks sequence lengths, steps, and sample counts; disables eval loops for speed. +- Suitable for stakeholder demos, not for final quality training. \ No newline at end of file diff --git a/fine_tuning/alignment/dpo/Dockerfile b/fine_tuning/alignment/dpo/Dockerfile new file mode 100644 index 0000000..20f1e4f --- /dev/null +++ b/fine_tuning/alignment/dpo/Dockerfile @@ -0,0 +1,38 @@ +FROM vault.habana.ai/gaudi-docker/1.21.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest + +# Set working directory +WORKDIR /workspace + +# Upgrade pip and install core Python dependencies +RUN pip install --upgrade pip && \ + pip install \ + optimum-habana==1.18.1 \ + transformers==4.51.0 \ + datasets==2.19.2 \ + pyarrow \ + rouge_score \ + nltk \ + evaluate \ + peft==0.12.0 \ + trl==0.9.6 \ + accelerate>=0.33 \ + sentencepiece!=0.1.92 \ + scipy \ + scikit-learn==1.5.2 \ + protobuf \ + torch>=1.3 \ + tyro +RUN pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.21.0 + +# Create directory structure +RUN mkdir -p /workspace/alignment/dpo + +# Copy all files from current directory (ppo) to container +COPY . /workspace/alignment/dpo/ + +# Set working directory to the copied files +WORKDIR /workspace/alignment/dpo + +# Default command +CMD ["/bin/bash"] + diff --git a/fine_tuning/alignment/dpo/Makefile b/fine_tuning/alignment/dpo/Makefile new file mode 100644 index 0000000..15ae1a6 --- /dev/null +++ b/fine_tuning/alignment/dpo/Makefile @@ -0,0 +1,20 @@ + +# Makefile for Fine-Tuning Workflow + +.PHONY: build run + +DOCKER_IMAGE=rlhf-ppo-llama-3-8b-habana +ENV_FLAGS=\ + -e HF_TOKEN=$(HF_TOKEN) + +DOCKER_RUN=docker run -it --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host $(ENV_FLAGS) $(DOCKER_IMAGE) + +build: + docker build -t $(DOCKER_IMAGE) . + +run: +ifeq ($(HF_TOKEN),) + @echo "ERROR: HF_TOKEN is not provided. Please provide your Hugging Face token: make run HF_TOKEN=your_hf_token_here" && exit 1 +else + $(DOCKER_RUN) +endif diff --git a/fine_tuning/alignment/dpo/README.md b/fine_tuning/alignment/dpo/README.md new file mode 100644 index 0000000..7dcf5e6 --- /dev/null +++ b/fine_tuning/alignment/dpo/README.md @@ -0,0 +1,124 @@ +# DPO Fine-Tuning on Gaudi (Gaudi2 / Gaudi3) — Gemma 3 & Llama Support + +This directory contains scripts and utilities for fine-tuning and reinforcement learning with language models using HuggingFace TRL, adapted for Habana Gaudi3 hardware. + +## ⚠️ Important: Model Arguments Required + +All scripts require explicit model specification. No default model is assumed. + +- Single scripts: pass `--model_name_or_path ` +- Pipeline script: `MODEL_NAME= ./dpo_pipeline_sanity.sh` +- If arguments are missing the script aborts with a clear message. + +## Contents + +- `sft.py` — Supervised fine-tuning (optional pre-step before DPO; can start from a base model). +- `dpo.py` — Direct Preference Optimization training (pairwise preference data). +- `compare_base_vs_dpo.py` — Compare base and DPO-finetuned models via logprob margins & qualitative outputs. +- `merge_peft_adapter.py` — Merge PEFT adapters into base model for deployment. +- `dpo_pipeline_sanity.sh` — One-shot SFT (optional) + DPO + comparison. +- `Dockerfile` — Container setup for Gaudi training. +- `Makefile` — Build and run targets. + +## End-to-End DPO Pipeline Sanity Check + +## Running the DPO Pipeline + +The `dpo_pipeline_sanity.sh` script demonstrates an ultra-fast workflow (<20 min target) on Gaudi: optional SFT warmup, DPO training on a tiny subset or synthetic pairs, and a qualitative comparison. All steps require explicit model specification. + +**To execute:** + +```bash +make build +make run HF_TOKEN=<> +chmod +x dpo_pipeline_sanity.sh +MODEL_NAME=google/gemma-3-270m ./dpo_pipeline_sanity.sh +``` + +### Steps Overview + +To achieve robust and production-quality fine-tuning of Llama models, the pipeline is split into key stages. Each step is essential for building a high-performing RLHF system: + +All steps run sequentially, logging to `dpo_pipeline_sanity.log`. + +1. **Supervised Fine-Tuning (SFT)** + - Trains the base Llama model on curated human demonstration data, teaching it to follow instructions and generate useful responses. LoRA adapters and Habana optimizations accelerate and scale this process. + - Example command (**model_name_or_path is required**): + ```bash + python sft.py \ + --model_name_or_path google/gemma-3-270m \ + --model_name_or_path google/gemma-3-270m \ + --dataset_name lvwerra/stack-exchange-paired \ + --output_dir ./sft_sanity \ + --do_train \ + --max_steps 50 \ + --per_device_train_batch_size 2 \ + --gradient_accumulation_steps 2 \ + --learning_rate 3e-5 \ + --lora_target_modules "q_proj" "k_proj" "v_proj" "o_proj" \ + --bf16 \ + --use_habana \ + --use_lazy_mode + ``` + +2. **Merge SFT Adapters** + - Integrates the learned LoRA weights into the base model, producing a single checkpoint for further training and evaluation. This simplifies deployment and downstream usage. + - Example command (**all arguments are required**): + ```bash + python merge_peft_adapter.py \ + --base_model_name "google/gemma-3-270m" \ + --base_model_name "google/gemma-3-270m" \ + --adapter_model_name "./sft_sanity" \ + --output_name "./sft_sanity_merged" + ``` + +3. **DPO Training** + - Directly optimizes the policy using preference pairs (chosen vs rejected) without a reward model, simplifying RLHF. + - Example command: + ```bash + PT_HPU_LAZY_MODE=1 python dpo.py \ + --model_name_or_path ./sft_sanity_merged \ + --tokenizer_name_or_path google/gemma-3-270m \ + --output_dir ./dpo_sanity \ + --beta 0.1 \ + --max_steps 40 \ + --per_device_train_batch_size 2 \ + --gradient_accumulation_steps 2 \ + --logging_steps 10 \ + --save_steps 999999 \ + --eval_steps 999999 \ + --lora_target_modules "q_proj" "k_proj" "v_proj" "o_proj" \ + --seed 42 + ``` + +4. **Compare Base vs DPO** + - Simple qualitative + logprob margin comparison between the base (or SFT) model and the DPO-finetuned model. + - Example command: + ```bash + PT_HPU_LAZY_MODE=1 python compare_base_vs_dpo.py \ + --base_model google/gemma-3-270m \ + --finetuned_model ./dpo_sanity \ + --seed 123 + ``` + +### Logging + +All output is logged to `ppo_pipeline_sanity.log` for review and debugging. + + + +## Comparing Base and DPO Models + +Use `compare_base_vs_dpo.py` for quick margins & qualitative outputs. Reward model arguments are unnecessary; DPO directly leverages preference pairs. + +```bash +PT_HPU_LAZY_MODE=1 python compare_base_vs_dpo.py \ + --base_model google/gemma-3-270m \ + --finetuned_model ./dpo_sanity \ + --seed 123 +``` + +This outputs logprob margins (chosen minus rejected) and selected generations to confirm improvement. + +--- + diff --git a/fine_tuning/alignment/dpo/compare_base_vs_dpo.py b/fine_tuning/alignment/dpo/compare_base_vs_dpo.py new file mode 100644 index 0000000..ba5a897 --- /dev/null +++ b/fine_tuning/alignment/dpo/compare_base_vs_dpo.py @@ -0,0 +1,418 @@ +#!/usr/bin/env python +# compare_base_vs_dpo.py (adapted from PPO comparison; reward model optional for qualitative check) + +import argparse +import warnings +import json +import logging +import os +import random +import time +from pathlib import Path +from typing import Dict + +import numpy as np +import torch +from transformers import ( + AutoTokenizer, + AutoModelForCausalLM, +) +try: + from peft import PeftConfig, PeftModel + _PEFT_AVAILABLE = True +except Exception: + _PEFT_AVAILABLE = False + +from logging_utils import setup_logging + + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument("--base_model", required=True, help="Base model path (REQUIRED)") + p.add_argument("--finetuned_model", "--ppo_model", dest="ppo_model", required=True, help="Fine-tuned/DPO model path (REQUIRED)") + p.add_argument("--max_new_tokens", type=int, default=96) + p.add_argument("--seed", type=int, default=None, help="Seed for reproducibility. If unset -> non-deterministic sampling") + p.add_argument("--do_sample", type=lambda v: str(v).lower() in {"1","true","yes"}, default=True) + p.add_argument("--top_p", type=float, default=0.95) + p.add_argument("--top_k", type=int, default=0) + p.add_argument("--temperature", type=float, default=1.05) + p.add_argument("--repetition_penalty", type=float, default=1.12, help="Penalty for repeating tokens (improves diversity)") + p.add_argument("--greedy", action="store_true", help="Override: force greedy generation (do_sample False)") + p.add_argument("--prompts_file", type=str, help="Optional path to a text file with one prompt per line") + p.add_argument( + "--output_json", type=str, default="compare_results.json", help="Where to write structured results" + ) + p.add_argument("--show_text", action="store_true", help="Print full generated texts for both models") + p.add_argument("--heuristics", action="store_true", help="Compute simple heuristic metrics (length, sentence count, keyword hits, diversity)") + p.add_argument("--min_new_tokens", type=int, default=0, help="Force at least this many new tokens by regenerating if shorter (simple safeguard against early EOS)") + p.add_argument("--stackexchange_style", action="store_true", help="Wrap prompts as 'Question: ...\n\nAnswer:' to mirror training format before generation") + return p.parse_args() + + +args = parse_args() +setup_logging() +warnings.filterwarnings( + "ignore", + message="Trainer.tokenizer is now deprecated. You should use Trainer.processing_class instead.", + category=UserWarning, +) +logger = logging.getLogger(__name__) + + +def set_all_seeds(seed: int): + if seed is None: + logger.info("[Seed] No seed provided -> results will vary run-to-run (sampling stochastic).") + return + logger.info(f"[Seed] Setting seed={seed}") + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + # Habana specific (if available) + if hasattr(torch, "hpu"): + try: + torch.hpu.manual_seed(seed) # type: ignore[attr-defined] + except Exception: + pass + os.environ["PYTHONHASHSEED"] = str(seed) + + +set_all_seeds(args.seed) + +base_model = args.base_model +ppo_model = args.ppo_model + +# Validate required arguments +if not base_model or base_model.strip() == "": + raise ValueError("--base_model is required and cannot be empty. Please specify the base model path.") +if not ppo_model or ppo_model.strip() == "": + raise ValueError("--finetuned_model is required and cannot be empty. Please specify the fine-tuned model path.") +logger.info( + f"[Init] base_model={base_model} finetuned_model={ppo_model} seed={args.seed}" +) + +def _hpu_available(): + # Robust check for HPU presence + try: + if not hasattr(torch, "hpu"): + import importlib + importlib.import_module("habana_frameworks.torch.core") # side-effect registers torch.hpu + return hasattr(torch, "hpu") and torch.hpu.is_available() + except Exception as e: + logger.warning(f"[HPU][Diagnose] Import error: {e}") + return False + +# Early diagnostics +logger.info(f"[HPU][Env] HPU env vars summary: " + ", ".join([f"{k}={v}" for k,v in os.environ.items() if k.startswith("HABANA") or k.startswith("PT_HPU")])) +logger.info(f"[HPU][Torch] hasattr(torch,'hpu')={hasattr(torch,'hpu')}") + +if _hpu_available(): + device = "hpu" + dtype = torch.bfloat16 + try: + import habana_frameworks.torch.core.hpu as hpu_mod # noqa: F401 + logger.info("[HPU][Import] habana_frameworks.torch.core.hpu imported successfully") + except Exception as e: + logger.warning(f"[HPU][Import] Failed optional import: {e}") + logger.info("[Device] Using Habana HPU") +elif torch.cuda.is_available(): + device = "cuda" + # Prefer bfloat16 if supported else float16 for small model speed + if torch.cuda.is_bf16_supported(): + dtype = torch.bfloat16 + else: + dtype = torch.float16 + logger.info(f"[Device] Using CUDA ({dtype})") +else: + device = "cpu" + dtype = torch.float32 + logger.info("[Device] Using CPU fallback (dtype=float32). This will be slower.") + +# Prompts +if args.prompts_file and Path(args.prompts_file).is_file(): + prompts = [ln.strip() for ln in Path(args.prompts_file).read_text().splitlines() if ln.strip()] + logger.info(f"[Prompts] Loaded {len(prompts)} prompts from {args.prompts_file}") +else: + prompts = [ + "Why do programmers prefer Python over Java for machine learning?", + "What are the advantages of Docker for deploying applications?", + ] + logger.info(f"[Prompts] Using default {len(prompts)} hard-coded prompts") + +# Tokenizer (shared) +tokenizer = AutoTokenizer.from_pretrained(base_model) +if getattr(tokenizer, "pad_token", None) is None: + tokenizer.pad_token = tokenizer.eos_token +tokenizer.padding_side = "left" # safer for causal LM generation when padding +logger.info("[Tokenizer] Loaded tokenizer & set pad_token -> eos_token") + +logger.info("[RewardModel] Removed reward model logic for pure generation comparison.") + +# Load both causal models once +def load_causal(path: str, is_adapter: bool = False, base_model: str = None): + if is_adapter and not _PEFT_AVAILABLE: + logger.warning("[ModelLoad][PEFT] Adapter indicated but PEFT not available; attempting direct load.") + if is_adapter and _PEFT_AVAILABLE: + if not base_model: + raise ValueError("Base model path required to load PEFT adapter") + logger.info(f"[ModelLoad][PEFT] Loading base model then applying adapter: base={base_model} adapter={path}") + base = AutoModelForCausalLM.from_pretrained( + base_model, + torch_dtype=dtype, + low_cpu_mem_usage=True, + ) + base.to(device) + base.eval() + wrapped = PeftModel.from_pretrained(base, path) + try: + logger.info("[ModelLoad][PEFT] Merging LoRA adapter for faster generation") + merged = wrapped.merge_and_unload() + merged.to(device) + merged.eval() + return merged + except Exception as e: + logger.warning(f"[ModelLoad][PEFT] merge_and_unload failed: {e}; using wrapped model.") + wrapped.to(device) + wrapped.eval() + return wrapped + else: + logger.info(f"[ModelLoad] Loading causal LM: {path}") + m = AutoModelForCausalLM.from_pretrained( + path, + torch_dtype=dtype, + low_cpu_mem_usage=True, + ) + m.to(device) + m.eval() + return m + + +def _is_peft_adapter_dir(path: str) -> bool: + if not os.path.isdir(path): + return False + # Common adapter marker files + markers = ["adapter_config.json", "adapter_model.bin", "adapter_model.safetensors"] + for m in markers: + if os.path.exists(os.path.join(path, m)): + return True + return False + +is_adapter = _is_peft_adapter_dir(ppo_model) +if is_adapter: + logger.info(f"[Detect][PEFT] Finetuned path appears to be a PEFT adapter directory: {ppo_model}") +else: + logger.info(f"[Detect] Finetuned path treated as full model directory: {ppo_model}") + +models: Dict[str, torch.nn.Module] = { + "Base": load_causal(base_model), + "DPO": load_causal(ppo_model, is_adapter=is_adapter, base_model=base_model if is_adapter else None), +} + +if args.greedy: + args.do_sample = False +elif args.stackexchange_style and args.min_new_tokens == 0: + # Auto-raise safeguard when using Q/A wrapper + args.min_new_tokens = 24 + logger.info(f"[GenConfig][Auto] stackexchange_style active -> setting min_new_tokens={args.min_new_tokens}") +generation_kwargs = { + "max_new_tokens": args.max_new_tokens, + "do_sample": args.do_sample, + "top_p": args.top_p, + "top_k": args.top_k, + "temperature": args.temperature, + "repetition_penalty": args.repetition_penalty, + "pad_token_id": tokenizer.pad_token_id, + "eos_token_id": tokenizer.eos_token_id, +} +logger.info(f"[GenConfig] {generation_kwargs}") + + +def generate_response(model, prompt: str) -> str: + gen_prompt = prompt + if args.stackexchange_style: + gen_prompt = f"Question: {prompt}\n\nAnswer:" + inputs = tokenizer(gen_prompt, return_tensors="pt").to(device) + with torch.no_grad(): + outputs = model.generate(**inputs, **generation_kwargs) + text = tokenizer.decode(outputs[0], skip_special_tokens=True) + # Simple safeguard for min_new_tokens: re-run once if below threshold + if args.min_new_tokens > 0: + new_part = text[len(gen_prompt):] + token_len = len(tokenizer(new_part, add_special_tokens=False)["input_ids"]) + if token_len < args.min_new_tokens: + logger.info(f"[Gen][Retry] Token length {token_len} < min_new_tokens {args.min_new_tokens}; retrying once.") + with torch.no_grad(): + outputs = model.generate(**inputs, **generation_kwargs) + text = tokenizer.decode(outputs[0], skip_special_tokens=True) + return text + +def expand_if_too_short(model, original_prompt: str, response: str, min_tokens: int = 40) -> str: + """Iteratively expand very short responses with domain-specific hints until min_tokens reached.""" + gen_prefix = f"Question: {original_prompt}\n\nAnswer:" + content = response[len(gen_prefix):] if response.startswith(gen_prefix) else response + token_len = len(tokenizer(content, add_special_tokens=False)["input_ids"]) + if token_len >= min_tokens: + return response + logger.info(f"[Gen][Expand] Response tokens {token_len} < {min_tokens}; starting guided expansion loop.") + hints = [ + " Please provide a concise but informative explanation including at least three concrete advantages:", + " List key benefits such as portability, isolation, reproducibility, efficient resource usage, and simplified CI/CD:", + " Add a short practical example of how an engineering team benefits from these features:", + ] + working_text = response + for idx, hint in enumerate(hints, start=1): + new_prompt = gen_prefix + hint + inputs = tokenizer(new_prompt, return_tensors="pt").to(device) + with torch.no_grad(): + outputs = model.generate(**inputs, **generation_kwargs) + new_text = tokenizer.decode(outputs[0], skip_special_tokens=True) + new_content = new_text[len(gen_prefix):] if new_text.startswith(gen_prefix) else new_text + new_token_len = len(tokenizer(new_content, add_special_tokens=False)["input_ids"]) + logger.info(f"[Gen][Expand][Step {idx}] tokens={new_token_len}") + if new_token_len > token_len: + # Append incremental expansion content rather than replacing entire text + append_part = new_content + # Avoid duplicating existing content if model echoed original + if append_part.startswith(content.strip()): + append_part = append_part[len(content.strip()):].lstrip() + working_text = working_text.rstrip() + "\n" + append_part.strip() + content = content + " " + append_part + token_len = len(tokenizer(content, add_special_tokens=False)["input_ids"]) + if token_len >= min_tokens: + logger.info("[Gen][Expand] Target reached; stopping expansion.") + break + return working_text + + +results = [] +stage_start = time.time() +for idx, prompt in enumerate(prompts, start=1): + logger.info("=" * 60) + logger.info(f"[Prompt {idx}/{len(prompts)}] {prompt}") + responses = {} + per_prompt_start = time.time() + for name, model in models.items(): + model_total_start = time.time() + gen_start = time.time() + resp = generate_response(model, prompt) + # Fallback expansion only for DPO model + if name == "DPO": + resp = expand_if_too_short(model, prompt, resp, min_tokens=max(80, args.min_new_tokens)) + gen_dur = time.time() - gen_start + model_total = time.time() - model_total_start + logger.info( + f"[Gen] {name} gen_time={gen_dur:.2f}s total_model_time={model_total:.2f}s" + ) + responses[name] = { + "response": resp, + "score": None, + "gen_time_sec": gen_dur, + "reward_time_sec": 0.0, + "model_total_time_sec": model_total, + } + per_prompt_dur = time.time() - per_prompt_start + # No reward scoring -> no delta + responses["duration_sec"] = per_prompt_dur + results.append({"prompt": prompt, "results": responses}) + +total_dur = time.time() - stage_start +logger.info("==================== SUMMARY (Base vs DPO) ====================") +for r in results: + # Scores are always None now; show lengths instead + base_resp = r["results"].get("Base", {}).get("response", "") + dpo_resp = r["results"].get("DPO", {}).get("response", "") + base_len = len(base_resp) + dpo_len = len(dpo_resp) + logger.info( + f"Prompt: {r['prompt'][:60]}... | Base_len={base_len} DPO_len={dpo_len}" + ) + +print("\n" + "=" * 100) +print(" Generation Comparison (no reward model / DPO pure generations)") +print("=" * 100) +for r in results: + prompt_full = r["prompt"] + prompt_short = (prompt_full[:90] + "…") if len(prompt_full) > 93 else prompt_full + base_entry = r["results"].get("Base") + ppo_entry = r["results"].get("DPO") + base_time = base_entry["model_total_time_sec"] if base_entry else float("nan") + ppo_time = ppo_entry["model_total_time_sec"] if ppo_entry else float("nan") + base_len = len(base_entry["response"]) if base_entry else 0 + dpo_len = len(ppo_entry["response"]) if ppo_entry else 0 + print(f"Prompt: {prompt_short}") + print(f" Base: Time: {base_time:.2f}s | Length: {base_len}") + print(f" DPO : Time: {ppo_time:.2f}s | Length: {dpo_len}") + if args.show_text: + print(" --- Base Output ---") + print(base_entry["response"]) + print(" --- DPO Output ---") + print(ppo_entry["response"]) + print() +print(f"Total prompts: {len(results)} | Total elapsed: {total_dur:.2f}s") + +# No aggregate reward stats since no scoring performed + +output_json = Path(args.output_json) +payload = {"total_duration_sec": total_dur, "comparisons": results, "gen_config": generation_kwargs} +def _conciseness_gain(comparisons): + gains = [] + for comp in comparisons: + base_txt = comp["results"]["Base"]["response"] + dpo_txt = comp["results"]["DPO"]["response"] + b_len = len(base_txt) + d_len = len(dpo_txt) + if d_len > 0: + gains.append(b_len / d_len) + if gains: + import statistics + return { + "avg_gain": sum(gains)/len(gains), + "median_gain": statistics.median(gains), + "min_gain": min(gains), + "max_gain": max(gains), + "count": len(gains), + } + return {} + +payload["conciseness_gain"] = _conciseness_gain(results) + +def simple_heuristics(text: str) -> Dict[str, float]: + import re + sentences = [s for s in re.split(r'[.!?]\s+', text.strip()) if s] + words = text.strip().split() + unique_words = set(w.lower() for w in words) + keywords = {"python", "docker", "learning", "performance", "model"} + keyword_hits = sum(1 for w in unique_words if w in keywords) + # Novelty heuristic: ratio of words not in first sentence to total words (rough measure of expansion) + first_sentence_words = set(w.lower() for w in (sentences[0].split() if sentences else [])) + novel_words = [w for w in unique_words if w not in first_sentence_words] + novelty_ratio = len(novel_words) / max(1, len(unique_words)) + return { + "char_len": len(text), + "token_len": len(words), + "sentence_count": len(sentences), + "unique_word_ratio": (len(unique_words) / max(1, len(words))), + "keyword_hits": keyword_hits, + "novelty_ratio": novelty_ratio, + } + +if args.heuristics: + for comp in payload["comparisons"]: + for name in ("Base", "DPO"): + resp = comp["results"][name]["response"] + comp["results"][name]["heuristics"] = simple_heuristics(resp) + # Aggregate averages + import numpy as np + agg = {"Base": {}, "DPO": {}} + keys = ["char_len", "token_len", "sentence_count", "unique_word_ratio", "keyword_hits", "novelty_ratio"] + for model_key in ("Base", "DPO"): + for k in keys: + vals = [c["results"][model_key]["heuristics"][k] for c in payload["comparisons"]] + agg[model_key][k] = float(sum(vals)/len(vals)) if vals else None + payload["heuristic_averages"] = agg + +with output_json.open("w") as f: + json.dump(payload, f, indent=2) +logger.info(f"[Output] JSON written -> {output_json} (time={total_dur:.2f}s)") +print(f"Results JSON: {output_json} | Total time: {total_dur:.2f}s") diff --git a/fine_tuning/alignment/dpo/dpo.py b/fine_tuning/alignment/dpo/dpo.py new file mode 100644 index 0000000..31fa199 --- /dev/null +++ b/fine_tuning/alignment/dpo/dpo.py @@ -0,0 +1,405 @@ +# copy from https://github.com/huggingface/trl/blob/v0.7.6/examples/research_projects/stack_llama_2/scripts/dpo_llama2.py, enable it for Gaudi2 +from dataclasses import dataclass, field +from typing import Dict, List, Optional +import os + +import torch +import warnings +from datasets import Dataset, load_dataset +from peft import LoraConfig +from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser +from transformers.integrations.deepspeed import ( + is_deepspeed_available, +) + +from optimum.habana import GaudiConfig +from optimum.habana.trl import GaudiDPOConfig, GaudiDPOTrainer +from optimum.habana.utils import set_seed + + +# Define and parse arguments. +@dataclass +class ScriptArguments: + """ + The arguments for the DPO training script. + """ + + # data parameters + beta: Optional[float] = field(default=0.1, metadata={"help": "the beta parameter for DPO loss"}) + + # training parameters + model_name_or_path: Optional[str] = field( + default="../sft/results/final_checkpoint", + metadata={"help": "the location of the SFT model name or path"}, + ) + tokenizer_name_or_path: Optional[str] = field( + default="meta-llama/Llama-2-7b-hf", + metadata={"help": "the location of the SFT model name or path"}, + ) + learning_rate: Optional[float] = field(default=5e-4, metadata={"help": "optimizer learning rate"}) + lr_scheduler_type: Optional[str] = field(default="cosine", metadata={"help": "the lr scheduler type"}) + warmup_steps: Optional[int] = field(default=100, metadata={"help": "the number of warmup steps"}) + weight_decay: Optional[float] = field(default=0.05, metadata={"help": "the weight decay"}) + optimizer_type: Optional[str] = field(default="paged_adamw_32bit", metadata={"help": "the optimizer type"}) + + per_device_train_batch_size: Optional[int] = field(default=1, metadata={"help": "train batch size per device"}) + per_device_eval_batch_size: Optional[int] = field(default=1, metadata={"help": "eval batch size per device"}) + gradient_accumulation_steps: Optional[int] = field( + default=4, metadata={"help": "the number of gradient accumulation steps"} + ) + gradient_checkpointing: Optional[bool] = field( + default=False, metadata={"help": "whether to use gradient checkpointing"} + ) + gradient_checkpointing_use_reentrant: Optional[bool] = field( + default=False, metadata={"help": "whether to use reentrant for gradient checkpointing"} + ) + + lora_alpha: Optional[float] = field(default=16, metadata={"help": "the lora alpha parameter"}) + lora_dropout: Optional[float] = field(default=0.05, metadata={"help": "the lora dropout parameter"}) + lora_r: Optional[int] = field(default=16, metadata={"help": "the lora r parameter (raised for higher capacity)"}) + lora_target_modules: List[str] = field( + default_factory=lambda: [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ], + metadata={"help": "Target modules for the LoRA method (expanded for better adaptation)."}, + ) + max_prompt_length: Optional[int] = field(default=512, metadata={"help": "the maximum prompt length"}) + max_length: Optional[int] = field(default=1024, metadata={"help": "the maximum sequence length"}) + max_steps: Optional[int] = field(default=30, metadata={"help": "max number of training steps (small default for fast demo)"}) + logging_steps: Optional[int] = field(default=10, metadata={"help": "the logging frequency"}) + save_steps: Optional[int] = field(default=100, metadata={"help": "the saving frequency"}) + eval_steps: Optional[int] = field(default=100, metadata={"help": "the evaluation frequency"}) + + output_dir: Optional[str] = field(default="./results", metadata={"help": "the output directory"}) + log_freq: Optional[int] = field(default=1, metadata={"help": "the logging frequency"}) + + # Merge options (parity with SFT script for end-to-end artifact convenience) + merge_adapter_after_train: Optional[bool] = field( + default=True, + metadata={"help": "If True, merge trained LoRA adapter into base weights after DPO training"}, + ) + merged_output_dir: Optional[str] = field( + default=None, + metadata={"help": "Destination directory for merged model (defaults to _merged)"}, + ) + merge_overwrite: Optional[bool] = field( + default=False, + metadata={"help": "If True overwrite merged_output_dir if it already exists"}, + ) + + # Pair quality filtering (stronger separation & length to avoid trivial short answers) + min_chosen_chars: Optional[int] = field( + default=60, + metadata={"help": "Require chosen response to have at least this many characters"}, + ) + min_score_delta: Optional[float] = field( + default=1.5, + metadata={"help": "Minimum (best_score - worst_score) needed to keep a pair"}, + ) + + # instrumentation + sanity_check: Optional[bool] = field(default=False, metadata={"help": "only train on 1000 samples"}) + report_to: Optional[str] = field( + default="none", + metadata={ + "help": 'Integrations to report to (e.g. "wandb", "tensorboard", "none"). Default now "none" to avoid requiring extra installs.' + }, + ) + # debug argument for distributed training + ignore_bias_buffers: Optional[bool] = field( + default=False, + metadata={ + "help": "fix for DDP issues with LM bias/mask buffers - invalid scalar type,`inplace operation. See" + "https://github.com/huggingface/transformers/issues/22482#issuecomment-1595790992" + }, + ) + seed: Optional[int] = field( + default=0, metadata={"help": "Random seed that will be set at the beginning of training."} + ) + deepspeed: Optional[str] = field(default=None, metadata={"help": "the deepspeed json config file"}) + num_workers: Optional[int] = field(default=None, metadata={"help": "the number of workers to map the data"}) + + +def get_ultrafeedback_pairs( + start: int = 0, + sample_limit: int = 1000, + num_proc: int = 8, + max_prompt_length: int = 512, + max_total_length: int = 1024, + min_chosen_chars: int = 60, + min_score_delta: float = 1.5, +) -> Dataset: + """Construct (prompt, chosen, rejected) pairs from openbmb/UltraFeedback. + + Strategy: + * Load slice of the single 'train' split using HF slice syntax + * For each instruction, score each completion by averaging available numeric ratings + * Select best (chosen) and worst (rejected) response + * Filter out degenerate or overly short/long samples + """ + end = start + sample_limit + slice_str = f"train[{start}:{end}]" + raw = load_dataset("openbmb/UltraFeedback", split=slice_str) + original_columns = raw.column_names + + def score_completion(comp) -> float: + ann = comp.get("annotations", {}) + ratings = [] + for aspect, value in ann.items(): + if isinstance(value, list): + for item in value: + r = item.get("Rating") if isinstance(item, dict) else None + if r is not None: + try: + ratings.append(int(r)) + except Exception: + pass + elif isinstance(value, dict): + r = value.get("Rating") + if r is not None: + try: + ratings.append(int(r)) + except Exception: + pass + return sum(ratings) / len(ratings) if ratings else 0.0 + + def map_batch(batch) -> Dict[str, List[str]]: + prompts: List[str] = [] + chosens: List[str] = [] + rejecteds: List[str] = [] + for instr, completions in zip(batch.get("instruction", []), batch.get("completions", [])): + if not isinstance(instr, str) or not isinstance(completions, list) or len(completions) < 2: + continue + scored = [] + for c in completions: + if not isinstance(c, dict): + continue + resp = c.get("response", "") + if not isinstance(resp, str) or len(resp.strip()) < 20: + continue + s = score_completion(c) + length_bonus = 0.15 * (len(resp.strip()) / 100) # reward more substantive answers modestly + composite = s + length_bonus + scored.append((composite, resp, s)) + if len(scored) < 2: + continue + scored.sort(key=lambda x: x[0], reverse=True) + best_score, best_resp, raw_best = scored[0] + worst_score, worst_resp, raw_worst = scored[-1] + if best_resp.strip() == worst_resp.strip(): + continue + # Enforce minimum quality & separation + if len(best_resp.strip()) < min_chosen_chars: + continue + if (best_score - worst_score) < min_score_delta: + continue + # Additional constraint: prefer pairs where chosen answer is meaningfully longer than rejected + if len(best_resp) - len(worst_resp) < 40: + continue + prompt_text = instr.strip() + "\n\nAnswer:" + if len(prompt_text) > max_prompt_length: + continue + if (len(prompt_text) + len(best_resp)) > max_total_length or (len(prompt_text) + len(worst_resp)) > max_total_length: + continue + prompts.append(prompt_text) + chosens.append(best_resp) + rejecteds.append(worst_resp) + return {"prompt": prompts, "chosen": chosens, "rejected": rejecteds} + + ds = raw.map(map_batch, batched=True, num_proc=num_proc, remove_columns=original_columns) + ds = ds.filter(lambda x: len(x["prompt"]) > 0 and len(x["chosen"]) > 0 and len(x["rejected"]) > 0) + return ds + + +if __name__ == "__main__": + # Suppress repetitive deprecation warnings about Trainer.tokenizer + warnings.filterwarnings( + "ignore", + message="Trainer.tokenizer is now deprecated. You should use Trainer.processing_class instead.", + category=UserWarning, + ) + parser = HfArgumentParser(ScriptArguments) + script_args = parser.parse_args_into_dataclasses()[0] + + # Determine reporting integrations (disable by default; enable wandb only if requested) + env_enable_wandb = os.environ.get("ENABLE_WANDB") + report_to = script_args.report_to + if env_enable_wandb and report_to in ("none", "None", ""): + report_to = "wandb" + # Normalize disabling values to empty list for HF Trainer + if report_to in ("none", "None", "", None): + report_to = [] + + # 1. initialize training arguments: + training_args = GaudiDPOConfig( + per_device_train_batch_size=script_args.per_device_train_batch_size, + per_device_eval_batch_size=script_args.per_device_eval_batch_size, + max_steps=script_args.max_steps, + logging_steps=script_args.logging_steps, + save_steps=script_args.save_steps, + gradient_accumulation_steps=script_args.gradient_accumulation_steps, + gradient_checkpointing=script_args.gradient_checkpointing, + learning_rate=script_args.learning_rate, + eval_strategy="steps", + eval_steps=script_args.eval_steps, + output_dir=script_args.output_dir, + report_to=report_to, + lr_scheduler_type=script_args.lr_scheduler_type, + warmup_steps=script_args.warmup_steps, + optim=script_args.optimizer_type, + bf16=True, + remove_unused_columns=False, + run_name="dpo_llama2", + gradient_checkpointing_kwargs={"use_reentrant": script_args.gradient_checkpointing_use_reentrant}, + use_habana=True, + use_lazy_mode=True, + use_hpu_graphs_for_training=not script_args.gradient_checkpointing and (not script_args.deepspeed), + use_hpu_graphs_for_inference=not script_args.deepspeed, + seed=script_args.seed, + deepspeed=script_args.deepspeed, + overwrite_output_dir=True, + ) + + # Set seed before initializing model. + set_seed(training_args.seed) + + low_cpu_mem_usage = True + if is_deepspeed_available(): + from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled + + if is_deepspeed_zero3_enabled(): + low_cpu_mem_usage = False + + # 2. load a pretrained model + model = AutoModelForCausalLM.from_pretrained( + script_args.model_name_or_path, + low_cpu_mem_usage=low_cpu_mem_usage, + torch_dtype=torch.bfloat16, + ) + model.config.use_cache = False + model.config.use_fused_rope = False + + if script_args.ignore_bias_buffers: + # torch distributed hack + model._ddp_params_and_buffers_to_ignore = [ + name for name, buffer in model.named_buffers() if buffer.dtype == torch.bool + ] + + model_ref = AutoModelForCausalLM.from_pretrained( + script_args.model_name_or_path, + low_cpu_mem_usage=low_cpu_mem_usage, + torch_dtype=torch.bfloat16, + ) + model_ref.config.use_cache = False + tokenizer = AutoTokenizer.from_pretrained(script_args.tokenizer_name_or_path) + tokenizer.pad_token = tokenizer.eos_token + + # 3. Load UltraFeedback pairs from train split slices + train_dataset = get_ultrafeedback_pairs( + start=0, + sample_limit=1000, + num_proc=script_args.num_workers or 8, + max_prompt_length=script_args.max_prompt_length, + max_total_length=script_args.max_length, + min_chosen_chars=script_args.min_chosen_chars, + min_score_delta=script_args.min_score_delta, + ) + # 4. Evaluation slice (non-overlapping range) + eval_dataset = get_ultrafeedback_pairs( + start=1200, + sample_limit=200, + num_proc=script_args.num_workers or 8, + max_prompt_length=script_args.max_prompt_length, + max_total_length=script_args.max_length, + min_chosen_chars=script_args.min_chosen_chars, + min_score_delta=script_args.min_score_delta, + ) + try: + print(f"[UltraFeedback][Dataset] train_pairs={len(train_dataset)} eval_pairs={len(eval_dataset)}") + except Exception: + pass + + peft_config = LoraConfig( + r=script_args.lora_r, + lora_alpha=script_args.lora_alpha, + lora_dropout=script_args.lora_dropout, + target_modules=script_args.lora_target_modules, + bias="none", + task_type="CAUSAL_LM", + ) + + gaudi_config = GaudiConfig() + gaudi_config.use_fused_adam = True + gaudi_config.use_fused_clip_norm = True + + # 5. initialize the DPO trainer + dpo_trainer = GaudiDPOTrainer( + model, + model_ref, + gaudi_config=gaudi_config, + args=training_args, + beta=script_args.beta, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + tokenizer=tokenizer, + peft_config=peft_config, + max_prompt_length=script_args.max_prompt_length, + max_length=script_args.max_length, + force_use_ref_model=True, + ) + # Assign processing_class to avoid downstream tokenizer property access + try: + setattr(dpo_trainer, 'processing_class', tokenizer) + except Exception: + pass + + # 6. train + train_result = dpo_trainer.train() + + # 7. save + dpo_trainer.save_model(script_args.output_dir) + + # 8. save metric + metrics = train_result.metrics + dpo_trainer.log_metrics("train", metrics) + dpo_trainer.save_metrics("train", metrics) + + # 9. Optional merge of adapter for faster downstream loading + if script_args.merge_adapter_after_train: + try: + from peft import PeftConfig, PeftModel + import os + adapter_dir = script_args.output_dir + peft_conf = PeftConfig.from_pretrained(adapter_dir) + # Only proceed if this is actually a PEFT adapter (safetensors or bin present) + marker_files = [ + os.path.join(adapter_dir, f) for f in ["adapter_model.safetensors", "adapter_model.bin"] + ] + if not any(os.path.exists(m) for m in marker_files): + print(f"[MergeInline][DPO] Adapter markers not found in {adapter_dir}; skipping merge.") + else: + base_loader = AutoModelForCausalLM + base_fresh = base_loader.from_pretrained( + script_args.model_name_or_path, + low_cpu_mem_usage=True, + torch_dtype=torch.bfloat16, + ) + merged_model = PeftModel.from_pretrained(base_fresh, adapter_dir) + print("[MergeInline][DPO] Adapter loaded into fresh base model; merging...") + merged_model = merged_model.merge_and_unload() + out_dir = script_args.merged_output_dir or f"{script_args.output_dir}_merged" + if os.path.exists(out_dir) and not script_args.merge_overwrite: + print(f"[MergeInline][DPO] Output dir {out_dir} exists and merge_overwrite=False -> abort merge") + else: + os.makedirs(out_dir, exist_ok=True) + merged_model.save_pretrained(out_dir) + tokenizer.save_pretrained(out_dir) + print(f"[MergeInline][DPO] Merged model saved to {out_dir}") + except Exception as e: + print(f"[MergeInline][DPO] Merge failed: {e}") \ No newline at end of file diff --git a/fine_tuning/alignment/dpo/dpo_pipeline_sanity.sh b/fine_tuning/alignment/dpo/dpo_pipeline_sanity.sh new file mode 100755 index 0000000..5c1ba5d --- /dev/null +++ b/fine_tuning/alignment/dpo/dpo_pipeline_sanity.sh @@ -0,0 +1,158 @@ +#!/usr/bin/env bash +set -euo pipefail + +START_PIPELINE=$(date +%s) + +# Fast DPO pipeline (override via env if desired) +DPO_MAX_STEPS=${DPO_MAX_STEPS:-140} +DPO_LOGGING_STEPS=${DPO_LOGGING_STEPS:-10} +DPO_LORA_TARGET=${DPO_LORA_TARGET:-"q_proj k_proj v_proj o_proj gate_proj up_proj down_proj"} + +# Base model (now mandatory: pass as first arg OR export MODEL_NAME) +MODEL_NAME="${1:-${MODEL_NAME:-}}" +if [ -z "${MODEL_NAME}" ]; then + echo "ERROR: Base model not specified. Provide as first argument or export MODEL_NAME." + echo "Usage: MODEL_NAME=google/gemma-3-270m ./dpo_pipeline_sanity.sh" >&2 + echo " or: ./dpo_pipeline_sanity.sh google/gemma-3-270m" >&2 + exit 1 +fi +echo "[CONFIG] MODEL_NAME=${MODEL_NAME}" +declare -A STEP_START +declare -A STEP_DURATION + +time_step_begin() { + local key="$1"; STEP_START[$key]=$(date +%s) + echo "[TIMER] BEGIN $key at $(date -u '+%Y-%m-%dT%H:%M:%SZ')" +} + +time_step_end() { + local key="$1"; local end=$(date +%s); local dur=$(( end - STEP_START[$key] )) + STEP_DURATION[$key]=$dur + echo "[TIMER] END $key at $(date -u '+%Y-%m-%dT%H:%M:%SZ') (duration=${dur}s)" +} + +LOGFILE="dpo_pipeline_sanity.log" +exec > >(tee -a "$LOGFILE") 2>&1 + +banner() { + echo + echo "============================================================" + echo ">>> $1" + echo "============================================================" +} + +# Base model (now mandatory: pass as first arg OR export MODEL_NAME) +MODEL_NAME="${1:-${MODEL_NAME:-}}" +if [ -z "${MODEL_NAME}" ]; then + echo "ERROR: Base model not specified. Provide as first argument or export MODEL_NAME." + echo "Usage: MODEL_NAME=google/gemma-3-270m ./ppo_pipeline_sanity.sh" >&2 + echo " or: ./ppo_pipeline_sanity.sh google/gemma-3-270m" >&2 + exit 1 +fi +echo "[CONFIG] MODEL_NAME=${MODEL_NAME}" + +# Common HPU env vars +export PT_HPU_DISABLE_FUSED_ADAMW=1 +export PT_HPU_DISABLE_FUSED_ADAM=1 +export PT_HPU_ENABLE_REFINE_DYNAMIC_SHAPES=0 + +if [ "${SKIP_SFT:-0}" = "1" ]; then + banner "STEP 1: Supervised Fine-Tuning SKIPPED (using base model directly)" + # Create a symlink named sft_sanity_merged pointing to the base model directory (if local) or just reuse MODEL_NAME if remote + if [ -d "${MODEL_NAME}" ]; then + ln -sfn "${MODEL_NAME}" ./sft_sanity_merged + else + # Remote model: create marker directory referencing base + mkdir -p ./sft_sanity_merged + echo "REMOTE_MODEL=${MODEL_NAME}" > ./sft_sanity_merged/BASE_MODEL_REFERENCE + fi +else + banner "STEP 1: Supervised Fine-Tuning (Ultra-Fast)" + time_step_begin SFT + python sft.py \ + --model_name_or_path "${MODEL_NAME}" \ + --dataset_name openbmb/UltraFeedback \ + --subset None \ + --streaming False \ + --size_valid_set 200 \ + --output_dir ./sft_sanity \ + --do_train \ + --max_steps 20 \ + --logging_steps 5 \ + --save_steps 999999 \ + --per_device_train_batch_size 2 \ + --gradient_accumulation_steps 1 \ + --learning_rate 3e-5 \ + --lr_scheduler_type cosine \ + --warmup_steps 2 \ + --weight_decay 0.01 \ + --lora_target_modules "q_proj" "k_proj" "v_proj" "o_proj" "gate_proj" "up_proj" "down_proj" \ + --bf16 \ + --remove_unused_columns False \ + --report_to none \ + --use_habana \ + --use_lazy_mode \ + --demo_fast \ + --merge_adapter_after_train \ + --merged_output_dir ./sft_sanity_merged + time_step_end SFT +fi + +banner "STEP 2: DPO Training" +time_step_begin DPO +PT_HPU_LAZY_MODE=1 python dpo.py \ + --model_name_or_path ./sft_sanity_merged \ + --tokenizer_name_or_path "${MODEL_NAME}" \ + --output_dir ./dpo_sanity \ + --beta 0.05 \ + --max_steps ${DPO_MAX_STEPS} \ + --learning_rate 2e-5 \ + --warmup_steps 10 \ + --per_device_train_batch_size 2 \ + --per_device_eval_batch_size 2 \ + --gradient_accumulation_steps 1 \ + --logging_steps ${DPO_LOGGING_STEPS} \ + --save_steps 999999 \ + --eval_steps 999999 \ + --lora_target_modules ${DPO_LORA_TARGET} \ + --report_to none \ + --seed 42 \ + --merge_adapter_after_train True \ + --merged_output_dir ./dpo_sanity_merged +time_step_end DPO + +banner "STEP 3: Comparing Base and DPO Models" +time_step_begin COMPARE +PT_HPU_LAZY_MODE=1 python compare_base_vs_dpo.py \ + --base_model "${MODEL_NAME}" \ + --finetuned_model $( [ -d ./dpo_sanity_merged ] && echo ./dpo_sanity_merged || echo ./dpo_sanity ) \ + --seed 123 \ + --max_new_tokens 160 \ + --min_new_tokens 80 \ + --temperature 0.85 \ + --repetition_penalty 1.15 \ + --stackexchange_style \ + --heuristics \ + --show_text +time_step_end COMPARE + +PIPELINE_END=$(date +%s) +TOTAL_DUR=$(( PIPELINE_END - START_PIPELINE )) + +echo +echo "=================== PIPELINE TIMING SUMMARY ===================" +printf "%-20s %10s\n" "Stage" "Seconds" +printf "%-20s %10s\n" "-----" "-------" +for k in SFT DPO COMPARE; do + printf "%-20s %10s\n" "$k" "${STEP_DURATION[$k]:-n/a}" +done +printf "%-20s %10s\n" "TOTAL" "$TOTAL_DUR" +echo "===============================================================" + +echo +echo "=================== ARTIFACT SUMMARY ===================" +if [ -d ./sft_sanity_merged ]; then echo "SFT merged policy: ./sft_sanity_merged"; fi +if [ -d ./dpo_sanity ]; then echo "DPO adapter policy: ./dpo_sanity"; fi +echo "========================================================" + +banner "✅ DPO PIPELINE COMPLETED SUCCESSFULLY (WITH COMPARISON)" diff --git a/fine_tuning/alignment/dpo/logging_utils.py b/fine_tuning/alignment/dpo/logging_utils.py new file mode 100644 index 0000000..271a53e --- /dev/null +++ b/fine_tuning/alignment/dpo/logging_utils.py @@ -0,0 +1,24 @@ +import logging +import time + + +class _UTCFormatter(logging.Formatter): + converter = time.gmtime # type: ignore[attr-defined] + + +def setup_logging(level: int = logging.INFO) -> None: + """ + Initialize root logger with consistent UTC ISO-like timestamps if not already configured. + Safe to call multiple times (no duplicate handlers added). + """ + root = logging.getLogger() + if root.handlers: + return + fmt = "[%(asctime)s] %(levelname)s %(name)s: %(message)s" + datefmt = "%Y-%m-%dT%H:%M:%S" + handler = logging.StreamHandler() + handler.setFormatter(_UTCFormatter(fmt, datefmt)) + root.setLevel(level) + root.addHandler(handler) + +__all__ = ["setup_logging"] diff --git a/fine_tuning/alignment/dpo/merge_peft_adapter.py b/fine_tuning/alignment/dpo/merge_peft_adapter.py new file mode 100644 index 0000000..e0f36bb --- /dev/null +++ b/fine_tuning/alignment/dpo/merge_peft_adapter.py @@ -0,0 +1,63 @@ +# Adapted from https://github.com/huggingface/optimum-habana/tree/v1.16.0/examples/trl +# Customized and enabled for Gaudi3 +from dataclasses import dataclass, field +import logging +from typing import Optional + +import torch +from peft import PeftConfig, PeftModel +from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer, HfArgumentParser + + +@dataclass +class ScriptArguments: + """ + The input names representing the Adapter and Base model fine-tuned with PEFT, and the output name representing the + merged model. + """ + + adapter_model_name: Optional[str] = field(default=None, metadata={"help": "the adapter name"}) + base_model_name: Optional[str] = field(default=None, metadata={"help": "the base model name"}) + output_name: Optional[str] = field(default=None, metadata={"help": "the merged model name"}) + + +from logging_utils import setup_logging +parser = HfArgumentParser(ScriptArguments) +script_args = parser.parse_args_into_dataclasses()[0] +setup_logging() +logger = logging.getLogger(__name__) +logger.info(f"ScriptArguments: {script_args}") +assert script_args.adapter_model_name is not None, "please provide the name of the Adapter you would like to merge" +assert script_args.base_model_name is not None, "please provide the name of the Base model" +assert script_args.output_name is not None, "please provide the output name of the merged model" + +if not (hasattr(torch, 'hpu') and torch.hpu.is_available()): + raise RuntimeError('[HPU][Required] Habana HPU not available. merge_peft_adapter enforces HPU use only.') +peft_config = PeftConfig.from_pretrained(script_args.adapter_model_name) +logger.info(f"[Adapter] Loaded PEFT config from {script_args.adapter_model_name} (task_type={peft_config.task_type})") +if peft_config.task_type == "SEQ_CLS": + logger.info("[IF] task_type == SEQ_CLS -> loading AutoModelForSequenceClassification (reward model style)") + # The sequence classification task is used for the reward model in PPO + model = AutoModelForSequenceClassification.from_pretrained( + script_args.base_model_name, num_labels=1, torch_dtype=torch.bfloat16 + ) +else: + logger.info("[IF] task_type != SEQ_CLS -> loading AutoModelForCausalLM") + model = AutoModelForCausalLM.from_pretrained( + script_args.base_model_name, return_dict=True, torch_dtype=torch.bfloat16 + ) + +tokenizer = AutoTokenizer.from_pretrained(script_args.base_model_name) + +# Load the PEFT model +model = PeftModel.from_pretrained(model, script_args.adapter_model_name) +logger.info(f"[Merge] Adapter weights loaded from {script_args.adapter_model_name}") +model.eval() + +model = model.to('hpu').merge_and_unload() +logger.info("[Merge] merge_and_unload complete -> adapter merged into base model") + +model.save_pretrained(f"{script_args.output_name}") +tokenizer.save_pretrained(f"{script_args.output_name}") +logger.info(f"[Save] Merged model + tokenizer saved to {script_args.output_name}") +# model.push_to_hub(f"{script_args.output_name}", use_temp_dir=False) diff --git a/fine_tuning/alignment/dpo/sft.py b/fine_tuning/alignment/dpo/sft.py new file mode 100644 index 0000000..8586e45 --- /dev/null +++ b/fine_tuning/alignment/dpo/sft.py @@ -0,0 +1,480 @@ +# Adapted from https://github.com/huggingface/optimum-habana/tree/v1.16.0/examples/trl +# Customized and enabled for Gaudi3 +import logging +import warnings +import math +from dataclasses import dataclass, field +from typing import List, Optional + +import torch +import transformers +from datasets import load_dataset +from peft import LoraConfig +from tqdm import tqdm +from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser +from transformers.integrations.deepspeed import ( + is_deepspeed_available, +) + +from optimum.habana import GaudiConfig +from optimum.habana.trl import GaudiSFTConfig, GaudiSFTTrainer +from optimum.habana.utils import set_seed +from logging_utils import setup_logging + + +logger = logging.getLogger(__name__) + + +@dataclass +class ScriptArguments: + model_name_or_path: Optional[str] = field(default=None, metadata={"help": "the model name (REQUIRED)"}) + dataset_name: Optional[str] = field(default=None, metadata={"help": "the dataset name"}) + use_peft: Optional[bool] = field(default=True, metadata={"help": "whether to use peft"}) + subset: Optional[str] = field(default=None, metadata={"help": "the subset (local data_dir) if applicable; None for remote HF datasets"}) + split: Optional[str] = field(default="train", metadata={"help": "the split to use"}) + size_valid_set: Optional[int] = field(default=4000, metadata={"help": "the size of the validation set"}) + streaming: Optional[bool] = field(default=True, metadata={"help": "whether to stream the dataset"}) + shuffle_buffer: Optional[int] = field(default=5000, metadata={"help": "the shuffle buffer size"}) + num_workers: Optional[int] = field(default=4, metadata={"help": "the number of workers"}) + num_buckets: Optional[int] = field(default=-1, metadata={"help": "whether to use bucketing for SFTTrainer"}) + validation_split_percentage: Optional[int] = field( + default=5, + metadata={ + "help": "The percentage of the train set used as validation set in case there's no validation split" + }, + ) + use_flash_attention: Optional[bool] = field( + default=False, metadata={"help": "Whether to use Habana flash attention for fine-tuning."} + ) + flash_attention_recompute: Optional[bool] = field( + default=False, metadata={"help": "Whether to enable recompute in Habana flash attention for fine-tuning."} + ) + flash_attention_causal_mask: Optional[bool] = field( + default=False, metadata={"help": "Whether to enable causal mask in Habana flash attention for fine-tuning."} + ) + + # LoraConfig + lora_alpha: Optional[float] = field(default=16, metadata={"help": "the lora alpha parameter"}) + lora_dropout: Optional[float] = field(default=0.05, metadata={"help": "the lora dropout parameter"}) + lora_r: Optional[int] = field(default=8, metadata={"help": "the lora r parameter"}) + lora_target_modules: List[str] = field( + default_factory=lambda: None, + metadata={"help": "Target modules for the LoRA method."}, + ) + + token: str = field( + default=None, + metadata={ + "help": ( + "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token " + "generated when running `huggingface-cli login` (stored in `~/.huggingface`)." + ) + }, + ) + merge_adapter_after_train: bool = field( + default=True, metadata={"help": "If True and PEFT used, merge LoRA adapter into base model after training."} + ) + merged_output_dir: Optional[str] = field( + default=None, metadata={"help": "Directory to save merged model (defaults to _merged)."} + ) + merge_overwrite: bool = field( + default=False, metadata={"help": "If True overwrite merged_output_dir if it exists."} + ) + demo_fast: bool = field( + default=False, metadata={"help": "Enable ultra-fast demo mode: shrink dataset & steps for <20m pipeline."} + ) + + +if __name__ == "__main__": + # Silence deprecated Trainer.tokenizer warnings + warnings.filterwarnings( + "ignore", + message="Trainer.tokenizer is now deprecated. You should use Trainer.processing_class instead.", + category=UserWarning, + ) + parser = HfArgumentParser((ScriptArguments, GaudiSFTConfig)) + script_args, training_args = parser.parse_args_into_dataclasses() + if not script_args.model_name_or_path: + raise ValueError("--model_name_or_path is required (no default). Provide a model path or repo id.") + if not (hasattr(torch, 'hpu') and torch.hpu.is_available()): + raise RuntimeError('[HPU][Required] Habana HPU not available. SFT script enforces HPU-only execution.') + if getattr(training_args, 'use_habana', False) is False: + logging.getLogger(__name__).warning('[HPU][Override] Forcing use_habana=True (Always Use HPU policy)') + training_args.use_habana = True + setup_logging() + + # Validate required arguments + if not script_args.model_name_or_path: + raise ValueError("--model_name_or_path is required. Please specify the base model path.") + + # Demo-fast overrides applied early (before logging full args) + if script_args.demo_fast: + logger.warning("[DemoFast][SFT] Activating ultra-fast overrides (<20m goal)") + # Reduce training steps if large + if training_args.max_steps == -1 or training_args.max_steps > 20: + logger.info("[DemoFast][SFT] Setting max_steps=20") + training_args.max_steps = 20 + # Force single epoch semantics when steps specified + if getattr(training_args,'num_train_epochs',1) > 1: + training_args.num_train_epochs = 1 + logger.info("[DemoFast][SFT] Forcing num_train_epochs=1") + # Reduce warmup + if getattr(training_args,'warmup_steps',0) > 2: + training_args.warmup_steps = 2 + logger.info("[DemoFast][SFT] Setting warmup_steps=2") + # Lower logging cadence + if getattr(training_args,'logging_steps',10) > 5: + training_args.logging_steps = 5 + logger.info("[DemoFast][SFT] Setting logging_steps=5") + # Min gradient accumulation + if getattr(training_args,'gradient_accumulation_steps',1) > 1: + training_args.gradient_accumulation_steps = 1 + logger.info("[DemoFast][SFT] Setting gradient_accumulation_steps=1") + # Limit LoRA rank + if script_args.use_peft and script_args.lora_r > 16: + logger.info("[DemoFast][SFT] Reducing lora_r to 16") + script_args.lora_r = 16 + # Reduce shuffle buffer if streaming + if script_args.streaming and script_args.shuffle_buffer > 1000: + script_args.shuffle_buffer = 1000 + logger.info("[DemoFast][SFT] Setting shuffle_buffer=1000") + # Shrink validation set if huge + if script_args.size_valid_set > 1000: + script_args.size_valid_set = 500 + logger.info("[DemoFast][SFT] Setting size_valid_set=500") + # Tag merged output dir for clarity + if script_args.merged_output_dir is None: + script_args.merged_output_dir = "./sft_sanity_merged" + logger.info(f"ScriptArguments: {script_args}") + logger.info(f"TrainingArguments: {training_args}") + if script_args.use_peft: + logger.info("[IF] use_peft == True -> configuring LoRA") + peft_config = LoraConfig( + r=script_args.lora_r, + lora_alpha=script_args.lora_alpha, + lora_dropout=script_args.lora_dropout, + target_modules=script_args.lora_target_modules, + bias="none", + task_type="CAUSAL_LM", + ) + else: + logger.info("[IF] use_peft == False -> full fine-tune (no LoRA)") + peft_config = None + + if training_args.group_by_length and training_args.packing: + logger.warning("[IF] group_by_length AND packing are both True -> raising ValueError") + raise ValueError("Cannot use both packing and group by length") + + set_seed(training_args.seed) + + def chars_token_ratio(dataset, tokenizer, nb_examples=400): + """ + Estimate the average number of characters per token in the dataset. + """ + total_characters, total_tokens = 0, 0 + for _, example in tqdm(zip(range(nb_examples), iter(dataset)), total=nb_examples): + text = prepare_sample_text(example) + total_characters += len(text) + if tokenizer.is_fast: + total_tokens += len(tokenizer(text).tokens()) + else: + total_tokens += len(tokenizer.tokenize(text)) + + return total_characters / total_tokens + + def prepare_sample_text(example): + """Prepare the text from a sample of supported datasets. + + For stack-exchange-paired: uses question + response_j. + For ultrafeedback_binarized: use prompt + chosen (preferred answer). + Falls back to joining all string fields if expected keys absent. + """ + if 'question' in example and 'response_j' in example: + return f"Question: {example['question']}\n\nAnswer: {example['response_j']}" + if 'prompt' in example and 'chosen' in example: + return f"{example['prompt'].rstrip()}\n{example['chosen']}" + # Fallback: concatenate first two string fields + for k in example: + if isinstance(example[k], str): + base = example[k] + break + else: + base = "" + extras = [v for k,v in example.items() if isinstance(v,str) and v is not base] + return base + ("\n" + "\n".join(extras) if extras else "") + + def create_datasets(tokenizer, args, seed=None): + if args.dataset_name: + logger.info(f"[IF] dataset_name provided -> loading {args.dataset_name}") + load_kwargs = { + 'split': args.split, + 'token': script_args.token, + 'num_proc': args.num_workers if not args.streaming else None, + 'streaming': args.streaming, + } + # Only pass data_dir if explicitly provided and not None (for local custom datasets) + if args.subset not in (None, 'None', ''): + load_kwargs['data_dir'] = args.subset + dataset = load_dataset(args.dataset_name, **load_kwargs) + # UltraFeedback: map to single (instruction, best_response) for SFT + if args.dataset_name == 'openbmb/UltraFeedback': + logger.info('[UltraFeedback][SFT] Mapping completions to best scored response') + def _score_completion(comp): + ann = comp.get('annotations', {}) + ratings = [] + for aspect, val in ann.items(): + if isinstance(val, list): + for item in val: + if isinstance(item, dict) and 'Rating' in item: + try: + ratings.append(int(item['Rating'])) + except Exception: + pass + elif isinstance(val, dict) and 'Rating' in val: + try: + ratings.append(int(val['Rating'])) + except Exception: + pass + return sum(ratings)/len(ratings) if ratings else 0.0 + def _map_batch(batch): + out_instr = [] + out_resp = [] + for instr, comps in zip(batch.get('instruction', []), batch.get('completions', [])): + if not isinstance(instr, str) or not isinstance(comps, list) or len(comps) == 0: + continue + scored = [] + for c in comps: + if not isinstance(c, dict): + continue + resp = c.get('response', '') + if not isinstance(resp, str) or len(resp.strip()) < 10: + continue + scored.append((_score_completion(c), resp)) + if len(scored) == 0: + continue + scored.sort(key=lambda x: x[0], reverse=True) + best_resp = scored[0][1] + # Clean up best_resp to remove embedded prompt repetitions like leading 'Question:' blocks + def _clean(r: str) -> str: + lines = [ln for ln in r.splitlines() if ln.strip()] + cleaned = [] + for ln in lines: + # Drop lines that start with typical prompt markers + if ln.strip().lower().startswith('question:'): + continue + cleaned.append(ln) + r2 = ' '.join(cleaned) + # Collapse excessive spaces + while ' ' in r2: + r2 = r2.replace(' ', ' ') + return r2.strip() + best_resp = _clean(best_resp) + # Guard: ensure we didn't strip everything + if len(best_resp) < 10: + continue + out_instr.append(instr) + out_resp.append(best_resp) + return {'question': out_instr, 'response_j': out_resp} + original_cols = dataset.column_names + dataset = dataset.map(_map_batch, batched=True, remove_columns=original_cols) + logger.info(f"[UltraFeedback][SFT] After mapping rows={len(dataset)} cols={dataset.column_names}") + else: + logger.error("[IF] dataset_name missing -> raising ValueError") + raise ValueError("No dataset_name") + if args.streaming: + logger.info("[IF] streaming == True -> using .take/.skip and shuffle") + logger.info("Loading the dataset in streaming mode") + valid_data = dataset.take(args.size_valid_set) + train_data = dataset.skip(args.size_valid_set) + train_data = train_data.shuffle(buffer_size=args.shuffle_buffer, seed=seed) + if script_args.demo_fast: + # For streaming, further restrict training iterable + fast_limit = 50 + logger.info(f"[DemoFast][SFT] Limiting streaming train_data to first {fast_limit} samples") + train_data = train_data.take(fast_limit) + valid_data = valid_data.take(min(script_args.size_valid_set, 50)) + else: + logger.info("[IF] streaming == False -> train_test_split path") + dataset = dataset.train_test_split(test_size=args.validation_split_percentage * 0.01, seed=seed) + train_data = dataset["train"] + valid_data = dataset["test"] + logger.info(f"Size of the train set: {len(train_data)}. Size of the validation set: {len(valid_data)}") + if script_args.demo_fast: + # Subsample train & valid sets + import random + fast_train_limit = min(80, len(train_data)) + fast_valid_limit = min(20, len(valid_data)) + logger.info(f"[DemoFast][SFT] Subsampling train->{fast_train_limit} valid->{fast_valid_limit}") + train_indices = list(range(len(train_data)))[:fast_train_limit] + valid_indices = list(range(len(valid_data)))[:fast_valid_limit] + train_data = train_data.select(train_indices) + valid_data = valid_data.select(valid_indices) + if args.dataset_name in ("lvwerra/stack-exchange-paired", "openbmb/UltraFeedback"): + logger.info(f"[IF] dataset_name {args.dataset_name} -> computing chars/token ratio and using formatting func") + try: + chars_per_token = chars_token_ratio(train_data, tokenizer) + logger.info(f"[DatasetStats] chars/token ratio: {chars_per_token:.2f}") + except Exception as e: + logger.warning(f"[DatasetStats] Failed chars/token ratio estimation: {e}") + formating_func = prepare_sample_text + else: + logger.info(f"[IF] dataset_name {args.dataset_name} not in specialized list -> no formatting func") + formating_func = None + return train_data, valid_data, formating_func + + low_cpu_mem_usage = True + if is_deepspeed_available(): + logger.info("[IF] Deepspeed available -> checking Zero3") + from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled + + if is_deepspeed_zero3_enabled(): + logger.info("[IF] DeepSpeed Zero3 enabled -> disabling low_cpu_mem_usage") + low_cpu_mem_usage = False + + base_model = AutoModelForCausalLM.from_pretrained( + script_args.model_name_or_path, + low_cpu_mem_usage=low_cpu_mem_usage, + torch_dtype=torch.bfloat16, + token=script_args.token, + ) + mt = getattr(base_model.config, 'model_type', '') + if mt.startswith('gemma3'): + try: + base_model.config.attn_implementation = 'eager' + except Exception: + pass + + # --- Parameter statistics before (and later after) PEFT injection --- + def _param_stats(model): + total = 0 + trainable = 0 + for p in model.parameters(): + n = p.numel() + total += n + if p.requires_grad: + trainable += n + pct = (trainable / total * 100) if total else 0.0 + return total, trainable, pct + + tot, trn, pct = _param_stats(base_model) + logger.info(f"[Model Params][Base] total={tot:,} trainable={trn:,} ({pct:.4f}%)") + + base_model.config.use_cache = False + if not script_args.use_flash_attention and ( + script_args.flash_attention_recompute or script_args.flash_attention_recompute + ): + logger.warning("[IF] flash attention recompute flags set while use_flash_attention is False -> assert") + assert "Need to enable use_flash_attention" + base_model.generation_config.use_flash_attention = script_args.use_flash_attention + base_model.generation_config.flash_attention_recompute = script_args.flash_attention_recompute + base_model.generation_config.flash_attention_causal_mask = script_args.flash_attention_causal_mask + if script_args.use_flash_attention: + logger.info("[IF] use_flash_attention == True -> flash attention settings applied") + else: + logger.info("[IF] use_flash_attention == False -> skipping flash attention optimizations") + + tokenizer = AutoTokenizer.from_pretrained(script_args.model_name_or_path, trust_remote_code=True) + tokenizer.pad_token = tokenizer.eos_token + tokenizer.padding_side = "right" # Fix weird overflow issue with fp16 training + + log_level = training_args.get_process_log_level() + logger.setLevel(log_level) + transformers.utils.logging.set_verbosity(log_level) + transformers.utils.logging.enable_default_handler() + transformers.utils.logging.enable_explicit_format() + + train_dataset, eval_dataset, formatting_func = create_datasets(tokenizer, script_args, seed=training_args.seed) + + gaudi_config = GaudiConfig() + gaudi_config.use_fused_adam = False + gaudi_config.use_fused_clip_norm = False + if training_args.do_train: + logger.info("[IF] do_train == True -> beginning training loop") + trainer = GaudiSFTTrainer( + model=base_model, + gaudi_config=gaudi_config, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + peft_config=peft_config, + tokenizer=tokenizer, + args=training_args, + formatting_func=formatting_func, + num_buckets=script_args.num_buckets, + ) + # Provide processing_class attribute expected by newer Trainer interface + try: + setattr(trainer, 'processing_class', tokenizer) + except Exception: + pass + + # If PEFT active, trainer.model now wrapped; print updated stats + try: + wrapped_tot, wrapped_trn, wrapped_pct = _param_stats(trainer.model) + logger.info(f"[Model Params][After Trainer Init] total={wrapped_tot:,} trainable={wrapped_trn:,} ({wrapped_pct:.4f}%)") + if peft_config is not None: + # Highlight LoRA target modules used + logger.info(f"[LoRA] Target modules: {peft_config.target_modules}") + logger.info(f"[LoRA] r={getattr(peft_config,'r',None)} alpha={getattr(peft_config,'lora_alpha',None)} dropout={getattr(peft_config,'lora_dropout',None)}") + except Exception as e: + logger.warning(f"[Warn] Could not compute PEFT parameter stats: {e}") + train_result = trainer.train() + trainer.save_model(training_args.output_dir) + metrics = train_result.metrics + trainer.log_metrics("train", metrics) + trainer.save_metrics("train", metrics) + logger.info("[INFO] Training completed and model saved") + # Optional inline merge + if script_args.merge_adapter_after_train: + if not script_args.use_peft: + logger.warning("[MergeInline] merge_adapter_after_train=True but use_peft=False -> skipping merge") + else: + try: + from peft import PeftConfig, PeftModel + import os + adapter_dir = training_args.output_dir + peft_conf = PeftConfig.from_pretrained(adapter_dir) + logger.info(f"[MergeInline] Loaded PEFT config task_type={peft_conf.task_type}") + base_loader = AutoModelForCausalLM + base_fresh = base_loader.from_pretrained( + script_args.model_name_or_path, + low_cpu_mem_usage=low_cpu_mem_usage, + torch_dtype=torch.bfloat16, + token=script_args.token, + ) + merged_model = PeftModel.from_pretrained(base_fresh, adapter_dir) + logger.info("[MergeInline] Adapter loaded into fresh base model; merging...") + merged_model = merged_model.merge_and_unload() + out_dir = script_args.merged_output_dir or f"{training_args.output_dir}_merged" + if os.path.exists(out_dir) and not script_args.merge_overwrite: + logger.warning(f"[MergeInline] Output dir {out_dir} exists and merge_overwrite=False -> abort merge") + else: + os.makedirs(out_dir, exist_ok=True) + merged_model.save_pretrained(out_dir) + tokenizer.save_pretrained(out_dir) + logger.info(f"[MergeInline] Merged model saved to {out_dir}") + except Exception as e: + logger.exception(f"[MergeInline] Failed to merge adapter inline: {e}") + else: + logger.info("[IF] do_train == False -> skipping training") + + # Evaluation + if training_args.do_eval: + logger.info("[IF] do_eval == True -> running evaluation") + logger.info("*** Evaluate ***") + metrics = trainer.evaluate() + if isinstance(eval_dataset, torch.utils.data.IterableDataset): + eval_dataset = list(eval_dataset) + + metrics["eval_samples"] = len(eval_dataset) + + try: + perplexity = math.exp(metrics["eval_loss"]) + except OverflowError: + perplexity = float("inf") + metrics["perplexity"] = perplexity + + trainer.log_metrics("eval", metrics) + trainer.save_metrics("eval", metrics) + logger.info("[INFO] Evaluation complete") + else: + logger.info("[IF] do_eval == False -> skipping evaluation") diff --git a/fine_tuning/alignment/dpo_latest/Dockerfile b/fine_tuning/alignment/dpo_latest/Dockerfile new file mode 100644 index 0000000..34e89a6 --- /dev/null +++ b/fine_tuning/alignment/dpo_latest/Dockerfile @@ -0,0 +1,38 @@ +FROM vault.habana.ai/gaudi-docker/1.21.0/ubuntu22.04/habanalabs/pytorch-installer-2.6.0:latest + +# Set working directory +WORKDIR /workspace + +# Upgrade pip and install core Python dependencies +RUN pip install --upgrade pip && \ + pip install \ + optimum-habana==1.18.1 \ + transformers==4.51.0 \ + datasets==2.19.2 \ + pyarrow \ + rouge_score \ + nltk \ + evaluate \ + peft==0.12.0 \ + trl==0.9.6 \ + accelerate>=0.33 \ + sentencepiece!=0.1.92 \ + scipy \ + scikit-learn==1.5.2 \ + protobuf \ + torch>=1.3 \ + tyro +RUN pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.21.0 + +# Create directory structure +RUN mkdir -p /workspace/alignment/dpo_latest + +# Copy all files from current directory (ppo) to container +COPY . /workspace/alignment/dpo_latest/ + +# Set working directory to the copied files +WORKDIR /workspace/alignment/dpo_latest + +# Default command +CMD ["/bin/bash"] + diff --git a/fine_tuning/alignment/dpo_latest/Makefile b/fine_tuning/alignment/dpo_latest/Makefile new file mode 100644 index 0000000..15ae1a6 --- /dev/null +++ b/fine_tuning/alignment/dpo_latest/Makefile @@ -0,0 +1,20 @@ + +# Makefile for Fine-Tuning Workflow + +.PHONY: build run + +DOCKER_IMAGE=rlhf-ppo-llama-3-8b-habana +ENV_FLAGS=\ + -e HF_TOKEN=$(HF_TOKEN) + +DOCKER_RUN=docker run -it --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host $(ENV_FLAGS) $(DOCKER_IMAGE) + +build: + docker build -t $(DOCKER_IMAGE) . + +run: +ifeq ($(HF_TOKEN),) + @echo "ERROR: HF_TOKEN is not provided. Please provide your Hugging Face token: make run HF_TOKEN=your_hf_token_here" && exit 1 +else + $(DOCKER_RUN) +endif diff --git a/fine_tuning/alignment/dpo_latest/README.md b/fine_tuning/alignment/dpo_latest/README.md new file mode 100644 index 0000000..7dcf5e6 --- /dev/null +++ b/fine_tuning/alignment/dpo_latest/README.md @@ -0,0 +1,124 @@ +# DPO Fine-Tuning on Gaudi (Gaudi2 / Gaudi3) — Gemma 3 & Llama Support + +This directory contains scripts and utilities for fine-tuning and reinforcement learning with language models using HuggingFace TRL, adapted for Habana Gaudi3 hardware. + +## ⚠️ Important: Model Arguments Required + +All scripts require explicit model specification. No default model is assumed. + +- Single scripts: pass `--model_name_or_path ` +- Pipeline script: `MODEL_NAME= ./dpo_pipeline_sanity.sh` +- If arguments are missing the script aborts with a clear message. + +## Contents + +- `sft.py` — Supervised fine-tuning (optional pre-step before DPO; can start from a base model). +- `dpo.py` — Direct Preference Optimization training (pairwise preference data). +- `compare_base_vs_dpo.py` — Compare base and DPO-finetuned models via logprob margins & qualitative outputs. +- `merge_peft_adapter.py` — Merge PEFT adapters into base model for deployment. +- `dpo_pipeline_sanity.sh` — One-shot SFT (optional) + DPO + comparison. +- `Dockerfile` — Container setup for Gaudi training. +- `Makefile` — Build and run targets. + +## End-to-End DPO Pipeline Sanity Check + +## Running the DPO Pipeline + +The `dpo_pipeline_sanity.sh` script demonstrates an ultra-fast workflow (<20 min target) on Gaudi: optional SFT warmup, DPO training on a tiny subset or synthetic pairs, and a qualitative comparison. All steps require explicit model specification. + +**To execute:** + +```bash +make build +make run HF_TOKEN=<> +chmod +x dpo_pipeline_sanity.sh +MODEL_NAME=google/gemma-3-270m ./dpo_pipeline_sanity.sh +``` + +### Steps Overview + +To achieve robust and production-quality fine-tuning of Llama models, the pipeline is split into key stages. Each step is essential for building a high-performing RLHF system: + +All steps run sequentially, logging to `dpo_pipeline_sanity.log`. + +1. **Supervised Fine-Tuning (SFT)** + - Trains the base Llama model on curated human demonstration data, teaching it to follow instructions and generate useful responses. LoRA adapters and Habana optimizations accelerate and scale this process. + - Example command (**model_name_or_path is required**): + ```bash + python sft.py \ + --model_name_or_path google/gemma-3-270m \ + --model_name_or_path google/gemma-3-270m \ + --dataset_name lvwerra/stack-exchange-paired \ + --output_dir ./sft_sanity \ + --do_train \ + --max_steps 50 \ + --per_device_train_batch_size 2 \ + --gradient_accumulation_steps 2 \ + --learning_rate 3e-5 \ + --lora_target_modules "q_proj" "k_proj" "v_proj" "o_proj" \ + --bf16 \ + --use_habana \ + --use_lazy_mode + ``` + +2. **Merge SFT Adapters** + - Integrates the learned LoRA weights into the base model, producing a single checkpoint for further training and evaluation. This simplifies deployment and downstream usage. + - Example command (**all arguments are required**): + ```bash + python merge_peft_adapter.py \ + --base_model_name "google/gemma-3-270m" \ + --base_model_name "google/gemma-3-270m" \ + --adapter_model_name "./sft_sanity" \ + --output_name "./sft_sanity_merged" + ``` + +3. **DPO Training** + - Directly optimizes the policy using preference pairs (chosen vs rejected) without a reward model, simplifying RLHF. + - Example command: + ```bash + PT_HPU_LAZY_MODE=1 python dpo.py \ + --model_name_or_path ./sft_sanity_merged \ + --tokenizer_name_or_path google/gemma-3-270m \ + --output_dir ./dpo_sanity \ + --beta 0.1 \ + --max_steps 40 \ + --per_device_train_batch_size 2 \ + --gradient_accumulation_steps 2 \ + --logging_steps 10 \ + --save_steps 999999 \ + --eval_steps 999999 \ + --lora_target_modules "q_proj" "k_proj" "v_proj" "o_proj" \ + --seed 42 + ``` + +4. **Compare Base vs DPO** + - Simple qualitative + logprob margin comparison between the base (or SFT) model and the DPO-finetuned model. + - Example command: + ```bash + PT_HPU_LAZY_MODE=1 python compare_base_vs_dpo.py \ + --base_model google/gemma-3-270m \ + --finetuned_model ./dpo_sanity \ + --seed 123 + ``` + +### Logging + +All output is logged to `ppo_pipeline_sanity.log` for review and debugging. + + + +## Comparing Base and DPO Models + +Use `compare_base_vs_dpo.py` for quick margins & qualitative outputs. Reward model arguments are unnecessary; DPO directly leverages preference pairs. + +```bash +PT_HPU_LAZY_MODE=1 python compare_base_vs_dpo.py \ + --base_model google/gemma-3-270m \ + --finetuned_model ./dpo_sanity \ + --seed 123 +``` + +This outputs logprob margins (chosen minus rejected) and selected generations to confirm improvement. + +--- + diff --git a/fine_tuning/alignment/dpo_latest/compare_base_vs_dpo.py b/fine_tuning/alignment/dpo_latest/compare_base_vs_dpo.py new file mode 100644 index 0000000..12f605f --- /dev/null +++ b/fine_tuning/alignment/dpo_latest/compare_base_vs_dpo.py @@ -0,0 +1,410 @@ +#!/usr/bin/env python +# compare_base_vs_dpo.py (adapted from PPO comparison; reward model optional for qualitative check) + +import argparse +import warnings +import json +import logging +import os +import random +import time +from pathlib import Path +from typing import Dict + +import numpy as np +import torch +from transformers import ( + AutoTokenizer, + AutoModelForCausalLM, +) +try: + from peft import PeftConfig, PeftModel + _PEFT_AVAILABLE = True +except Exception: + _PEFT_AVAILABLE = False + +from logging_utils import setup_logging + + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument("--base_model", required=True, help="Base model path (REQUIRED)") + p.add_argument("--finetuned_model", "--ppo_model", dest="ppo_model", required=True, help="Fine-tuned/DPO model path (REQUIRED)") + p.add_argument("--max_new_tokens", type=int, default=96) + p.add_argument("--seed", type=int, default=None, help="Seed for reproducibility. If unset -> non-deterministic sampling") + p.add_argument("--do_sample", type=lambda v: str(v).lower() in {"1","true","yes"}, default=True) + p.add_argument("--top_p", type=float, default=0.95) + p.add_argument("--top_k", type=int, default=0) + p.add_argument("--temperature", type=float, default=1.05) + p.add_argument("--greedy", action="store_true", help="Override: force greedy generation (do_sample False)") + p.add_argument("--prompts_file", type=str, help="Optional path to a text file with one prompt per line") + p.add_argument( + "--output_json", type=str, default="compare_results.json", help="Where to write structured results" + ) + p.add_argument("--show_text", action="store_true", help="Print full generated texts for both models") + p.add_argument("--heuristics", action="store_true", help="Compute simple heuristic metrics (length, sentence count, keyword hits, diversity)") + p.add_argument("--min_new_tokens", type=int, default=0, help="Force at least this many new tokens by regenerating if shorter (simple safeguard against early EOS)") + p.add_argument("--stackexchange_style", action="store_true", help="Wrap prompts as 'Question: ...\n\nAnswer:' to mirror training format before generation") + return p.parse_args() + + +args = parse_args() +setup_logging() +warnings.filterwarnings( + "ignore", + message="Trainer.tokenizer is now deprecated. You should use Trainer.processing_class instead.", + category=UserWarning, +) +logger = logging.getLogger(__name__) + + +def set_all_seeds(seed: int): + if seed is None: + logger.info("[Seed] No seed provided -> results will vary run-to-run (sampling stochastic).") + return + logger.info(f"[Seed] Setting seed={seed}") + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + # Habana specific (if available) + if hasattr(torch, "hpu"): + try: + torch.hpu.manual_seed(seed) # type: ignore[attr-defined] + except Exception: + pass + os.environ["PYTHONHASHSEED"] = str(seed) + + +set_all_seeds(args.seed) + +base_model = args.base_model +ppo_model = args.ppo_model + +# Validate required arguments +if not base_model or base_model.strip() == "": + raise ValueError("--base_model is required and cannot be empty. Please specify the base model path.") +if not ppo_model or ppo_model.strip() == "": + raise ValueError("--finetuned_model is required and cannot be empty. Please specify the fine-tuned model path.") +logger.info( + f"[Init] base_model={base_model} finetuned_model={ppo_model} seed={args.seed}" +) + +def _hpu_available(): + # Robust check for HPU presence + try: + if not hasattr(torch, "hpu"): + import importlib + importlib.import_module("habana_frameworks.torch.core") # side-effect registers torch.hpu + return hasattr(torch, "hpu") and torch.hpu.is_available() + except Exception as e: + logger.warning(f"[HPU][Diagnose] Import error: {e}") + return False + +# Early diagnostics +logger.info(f"[HPU][Env] HPU env vars summary: " + ", ".join([f"{k}={v}" for k,v in os.environ.items() if k.startswith("HABANA") or k.startswith("PT_HPU")])) +logger.info(f"[HPU][Torch] hasattr(torch,'hpu')={hasattr(torch,'hpu')}") + +if _hpu_available(): + device = "hpu" + dtype = torch.bfloat16 + try: + import habana_frameworks.torch.core.hpu as hpu_mod # noqa: F401 + logger.info("[HPU][Import] habana_frameworks.torch.core.hpu imported successfully") + except Exception as e: + logger.warning(f"[HPU][Import] Failed optional import: {e}") + logger.info("[Device] Using Habana HPU") +elif torch.cuda.is_available(): + device = "cuda" + # Prefer bfloat16 if supported else float16 for small model speed + if torch.cuda.is_bf16_supported(): + dtype = torch.bfloat16 + else: + dtype = torch.float16 + logger.info(f"[Device] Using CUDA ({dtype})") +else: + device = "cpu" + dtype = torch.float32 + logger.info("[Device] Using CPU fallback (dtype=float32). This will be slower.") + +# Prompts +if args.prompts_file and Path(args.prompts_file).is_file(): + prompts = [ln.strip() for ln in Path(args.prompts_file).read_text().splitlines() if ln.strip()] + logger.info(f"[Prompts] Loaded {len(prompts)} prompts from {args.prompts_file}") +else: + prompts = [ + "Why do programmers prefer Python over Java for machine learning?", + "What are the advantages of Docker for deploying applications?", + ] + logger.info(f"[Prompts] Using default {len(prompts)} hard-coded prompts") + +# Tokenizer (shared) +tokenizer = AutoTokenizer.from_pretrained(base_model) +if getattr(tokenizer, "pad_token", None) is None: + tokenizer.pad_token = tokenizer.eos_token +tokenizer.padding_side = "left" # safer for causal LM generation when padding +logger.info("[Tokenizer] Loaded tokenizer & set pad_token -> eos_token") + +logger.info("[RewardModel] Removed reward model logic for pure generation comparison.") + +# Load both causal models once +def load_causal(path: str, is_adapter: bool = False, base_model: str = None): + if is_adapter and not _PEFT_AVAILABLE: + logger.warning("[ModelLoad][PEFT] Adapter indicated but PEFT not available; attempting direct load.") + if is_adapter and _PEFT_AVAILABLE: + if not base_model: + raise ValueError("Base model path required to load PEFT adapter") + logger.info(f"[ModelLoad][PEFT] Loading base model then applying adapter: base={base_model} adapter={path}") + base = AutoModelForCausalLM.from_pretrained( + base_model, + torch_dtype=dtype, + low_cpu_mem_usage=True, + ) + base.to(device) + base.eval() + wrapped = PeftModel.from_pretrained(base, path) + try: + logger.info("[ModelLoad][PEFT] Merging LoRA adapter for faster generation") + merged = wrapped.merge_and_unload() + merged.to(device) + merged.eval() + return merged + except Exception as e: + logger.warning(f"[ModelLoad][PEFT] merge_and_unload failed: {e}; using wrapped model.") + wrapped.to(device) + wrapped.eval() + return wrapped + else: + logger.info(f"[ModelLoad] Loading causal LM: {path}") + m = AutoModelForCausalLM.from_pretrained( + path, + torch_dtype=dtype, + low_cpu_mem_usage=True, + ) + m.to(device) + m.eval() + return m + + +def _is_peft_adapter_dir(path: str) -> bool: + if not os.path.isdir(path): + return False + # Common adapter marker files + markers = ["adapter_config.json", "adapter_model.bin", "adapter_model.safetensors"] + for m in markers: + if os.path.exists(os.path.join(path, m)): + return True + return False + +is_adapter = _is_peft_adapter_dir(ppo_model) +if is_adapter: + logger.info(f"[Detect][PEFT] Finetuned path appears to be a PEFT adapter directory: {ppo_model}") +else: + logger.info(f"[Detect] Finetuned path treated as full model directory: {ppo_model}") + +models: Dict[str, torch.nn.Module] = { + "Base": load_causal(base_model), + "DPO": load_causal(ppo_model, is_adapter=is_adapter, base_model=base_model if is_adapter else None), +} + +if args.greedy: + args.do_sample = False +elif args.stackexchange_style and args.min_new_tokens == 0: + # Auto-raise safeguard when using Q/A wrapper + args.min_new_tokens = 24 + logger.info(f"[GenConfig][Auto] stackexchange_style active -> setting min_new_tokens={args.min_new_tokens}") +generation_kwargs = { + "max_new_tokens": args.max_new_tokens, + "do_sample": args.do_sample, + "top_p": args.top_p, + "top_k": args.top_k, + "temperature": args.temperature, + "pad_token_id": tokenizer.pad_token_id, + "eos_token_id": tokenizer.eos_token_id, +} +logger.info(f"[GenConfig] {generation_kwargs}") + + +def generate_response(model, prompt: str) -> str: + gen_prompt = prompt + if args.stackexchange_style: + gen_prompt = f"Question: {prompt}\n\nAnswer:" + inputs = tokenizer(gen_prompt, return_tensors="pt").to(device) + with torch.no_grad(): + outputs = model.generate(**inputs, **generation_kwargs) + text = tokenizer.decode(outputs[0], skip_special_tokens=True) + # Simple safeguard for min_new_tokens: re-run once if below threshold + if args.min_new_tokens > 0: + new_part = text[len(gen_prompt):] + token_len = len(tokenizer(new_part, add_special_tokens=False)["input_ids"]) + if token_len < args.min_new_tokens: + logger.info(f"[Gen][Retry] Token length {token_len} < min_new_tokens {args.min_new_tokens}; retrying once.") + with torch.no_grad(): + outputs = model.generate(**inputs, **generation_kwargs) + text = tokenizer.decode(outputs[0], skip_special_tokens=True) + return text + +def expand_if_too_short(model, original_prompt: str, response: str, min_tokens: int = 40) -> str: + """Iteratively expand very short responses with domain-specific hints until min_tokens reached.""" + gen_prefix = f"Question: {original_prompt}\n\nAnswer:" + content = response[len(gen_prefix):] if response.startswith(gen_prefix) else response + token_len = len(tokenizer(content, add_special_tokens=False)["input_ids"]) + if token_len >= min_tokens: + return response + logger.info(f"[Gen][Expand] Response tokens {token_len} < {min_tokens}; starting guided expansion loop.") + hints = [ + " Please provide a concise but informative explanation including at least three concrete advantages:", + " List key benefits such as portability, isolation, reproducibility, efficient resource usage, and simplified CI/CD:", + " Add a short practical example of how an engineering team benefits from these features:", + ] + working_text = response + for idx, hint in enumerate(hints, start=1): + new_prompt = gen_prefix + hint + inputs = tokenizer(new_prompt, return_tensors="pt").to(device) + with torch.no_grad(): + outputs = model.generate(**inputs, **generation_kwargs) + new_text = tokenizer.decode(outputs[0], skip_special_tokens=True) + new_content = new_text[len(gen_prefix):] if new_text.startswith(gen_prefix) else new_text + new_token_len = len(tokenizer(new_content, add_special_tokens=False)["input_ids"]) + logger.info(f"[Gen][Expand][Step {idx}] tokens={new_token_len}") + if new_token_len > token_len: + working_text = new_text + token_len = new_token_len + if token_len >= min_tokens: + logger.info("[Gen][Expand] Target reached; stopping expansion.") + break + return working_text + + +results = [] +stage_start = time.time() +for idx, prompt in enumerate(prompts, start=1): + logger.info("=" * 60) + logger.info(f"[Prompt {idx}/{len(prompts)}] {prompt}") + responses = {} + per_prompt_start = time.time() + for name, model in models.items(): + model_total_start = time.time() + gen_start = time.time() + resp = generate_response(model, prompt) + # Fallback expansion only for DPO model + if name == "DPO": + resp = expand_if_too_short(model, prompt, resp, min_tokens=max(80, args.min_new_tokens)) + gen_dur = time.time() - gen_start + model_total = time.time() - model_total_start + logger.info( + f"[Gen] {name} gen_time={gen_dur:.2f}s total_model_time={model_total:.2f}s" + ) + responses[name] = { + "response": resp, + "score": None, + "gen_time_sec": gen_dur, + "reward_time_sec": 0.0, + "model_total_time_sec": model_total, + } + per_prompt_dur = time.time() - per_prompt_start + # No reward scoring -> no delta + responses["duration_sec"] = per_prompt_dur + results.append({"prompt": prompt, "results": responses}) + +total_dur = time.time() - stage_start +logger.info("==================== SUMMARY (Base vs DPO) ====================") +for r in results: + # Scores are always None now; show lengths instead + base_resp = r["results"].get("Base", {}).get("response", "") + dpo_resp = r["results"].get("DPO", {}).get("response", "") + base_len = len(base_resp) + dpo_len = len(dpo_resp) + logger.info( + f"Prompt: {r['prompt'][:60]}... | Base_len={base_len} DPO_len={dpo_len}" + ) + +print("\n" + "=" * 100) +print(" Generation Comparison (no reward model / DPO pure generations)") +print("=" * 100) +for r in results: + prompt_full = r["prompt"] + prompt_short = (prompt_full[:90] + "…") if len(prompt_full) > 93 else prompt_full + base_entry = r["results"].get("Base") + ppo_entry = r["results"].get("DPO") + base_time = base_entry["model_total_time_sec"] if base_entry else float("nan") + ppo_time = ppo_entry["model_total_time_sec"] if ppo_entry else float("nan") + base_len = len(base_entry["response"]) if base_entry else 0 + dpo_len = len(ppo_entry["response"]) if ppo_entry else 0 + print(f"Prompt: {prompt_short}") + print(f" Base: Time: {base_time:.2f}s | Length: {base_len}") + print(f" DPO : Time: {ppo_time:.2f}s | Length: {dpo_len}") + if args.show_text: + print(" --- Base Output ---") + print(base_entry["response"]) + print(" --- DPO Output ---") + print(ppo_entry["response"]) + print() +print(f"Total prompts: {len(results)} | Total elapsed: {total_dur:.2f}s") + +# No aggregate reward stats since no scoring performed + +output_json = Path(args.output_json) +payload = {"total_duration_sec": total_dur, "comparisons": results, "gen_config": generation_kwargs} +def _conciseness_gain(comparisons): + gains = [] + for comp in comparisons: + base_txt = comp["results"]["Base"]["response"] + dpo_txt = comp["results"]["DPO"]["response"] + b_len = len(base_txt) + d_len = len(dpo_txt) + if d_len > 0: + gains.append(b_len / d_len) + if gains: + import statistics + return { + "avg_gain": sum(gains)/len(gains), + "median_gain": statistics.median(gains), + "min_gain": min(gains), + "max_gain": max(gains), + "count": len(gains), + } + return {} + +payload["conciseness_gain"] = _conciseness_gain(results) + +def simple_heuristics(text: str) -> Dict[str, float]: + import re + sentences = [s for s in re.split(r'[.!?]\s+', text.strip()) if s] + words = text.strip().split() + unique_words = set(w.lower() for w in words) + keywords = {"python", "docker", "learning", "performance", "model"} + keyword_hits = sum(1 for w in unique_words if w in keywords) + # Novelty heuristic: ratio of words not in first sentence to total words (rough measure of expansion) + first_sentence_words = set(w.lower() for w in (sentences[0].split() if sentences else [])) + novel_words = [w for w in unique_words if w not in first_sentence_words] + novelty_ratio = len(novel_words) / max(1, len(unique_words)) + return { + "char_len": len(text), + "token_len": len(words), + "sentence_count": len(sentences), + "unique_word_ratio": (len(unique_words) / max(1, len(words))), + "keyword_hits": keyword_hits, + "novelty_ratio": novelty_ratio, + } + +if args.heuristics: + for comp in payload["comparisons"]: + for name in ("Base", "DPO"): + resp = comp["results"][name]["response"] + comp["results"][name]["heuristics"] = simple_heuristics(resp) + # Aggregate averages + import numpy as np + agg = {"Base": {}, "DPO": {}} + keys = ["char_len", "token_len", "sentence_count", "unique_word_ratio", "keyword_hits", "novelty_ratio"] + for model_key in ("Base", "DPO"): + for k in keys: + vals = [c["results"][model_key]["heuristics"][k] for c in payload["comparisons"]] + agg[model_key][k] = float(sum(vals)/len(vals)) if vals else None + payload["heuristic_averages"] = agg + +with output_json.open("w") as f: + json.dump(payload, f, indent=2) +logger.info(f"[Output] JSON written -> {output_json} (time={total_dur:.2f}s)") +print(f"Results JSON: {output_json} | Total time: {total_dur:.2f}s") diff --git a/fine_tuning/alignment/dpo_latest/dpo.py b/fine_tuning/alignment/dpo_latest/dpo.py new file mode 100644 index 0000000..31fa199 --- /dev/null +++ b/fine_tuning/alignment/dpo_latest/dpo.py @@ -0,0 +1,405 @@ +# copy from https://github.com/huggingface/trl/blob/v0.7.6/examples/research_projects/stack_llama_2/scripts/dpo_llama2.py, enable it for Gaudi2 +from dataclasses import dataclass, field +from typing import Dict, List, Optional +import os + +import torch +import warnings +from datasets import Dataset, load_dataset +from peft import LoraConfig +from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser +from transformers.integrations.deepspeed import ( + is_deepspeed_available, +) + +from optimum.habana import GaudiConfig +from optimum.habana.trl import GaudiDPOConfig, GaudiDPOTrainer +from optimum.habana.utils import set_seed + + +# Define and parse arguments. +@dataclass +class ScriptArguments: + """ + The arguments for the DPO training script. + """ + + # data parameters + beta: Optional[float] = field(default=0.1, metadata={"help": "the beta parameter for DPO loss"}) + + # training parameters + model_name_or_path: Optional[str] = field( + default="../sft/results/final_checkpoint", + metadata={"help": "the location of the SFT model name or path"}, + ) + tokenizer_name_or_path: Optional[str] = field( + default="meta-llama/Llama-2-7b-hf", + metadata={"help": "the location of the SFT model name or path"}, + ) + learning_rate: Optional[float] = field(default=5e-4, metadata={"help": "optimizer learning rate"}) + lr_scheduler_type: Optional[str] = field(default="cosine", metadata={"help": "the lr scheduler type"}) + warmup_steps: Optional[int] = field(default=100, metadata={"help": "the number of warmup steps"}) + weight_decay: Optional[float] = field(default=0.05, metadata={"help": "the weight decay"}) + optimizer_type: Optional[str] = field(default="paged_adamw_32bit", metadata={"help": "the optimizer type"}) + + per_device_train_batch_size: Optional[int] = field(default=1, metadata={"help": "train batch size per device"}) + per_device_eval_batch_size: Optional[int] = field(default=1, metadata={"help": "eval batch size per device"}) + gradient_accumulation_steps: Optional[int] = field( + default=4, metadata={"help": "the number of gradient accumulation steps"} + ) + gradient_checkpointing: Optional[bool] = field( + default=False, metadata={"help": "whether to use gradient checkpointing"} + ) + gradient_checkpointing_use_reentrant: Optional[bool] = field( + default=False, metadata={"help": "whether to use reentrant for gradient checkpointing"} + ) + + lora_alpha: Optional[float] = field(default=16, metadata={"help": "the lora alpha parameter"}) + lora_dropout: Optional[float] = field(default=0.05, metadata={"help": "the lora dropout parameter"}) + lora_r: Optional[int] = field(default=16, metadata={"help": "the lora r parameter (raised for higher capacity)"}) + lora_target_modules: List[str] = field( + default_factory=lambda: [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ], + metadata={"help": "Target modules for the LoRA method (expanded for better adaptation)."}, + ) + max_prompt_length: Optional[int] = field(default=512, metadata={"help": "the maximum prompt length"}) + max_length: Optional[int] = field(default=1024, metadata={"help": "the maximum sequence length"}) + max_steps: Optional[int] = field(default=30, metadata={"help": "max number of training steps (small default for fast demo)"}) + logging_steps: Optional[int] = field(default=10, metadata={"help": "the logging frequency"}) + save_steps: Optional[int] = field(default=100, metadata={"help": "the saving frequency"}) + eval_steps: Optional[int] = field(default=100, metadata={"help": "the evaluation frequency"}) + + output_dir: Optional[str] = field(default="./results", metadata={"help": "the output directory"}) + log_freq: Optional[int] = field(default=1, metadata={"help": "the logging frequency"}) + + # Merge options (parity with SFT script for end-to-end artifact convenience) + merge_adapter_after_train: Optional[bool] = field( + default=True, + metadata={"help": "If True, merge trained LoRA adapter into base weights after DPO training"}, + ) + merged_output_dir: Optional[str] = field( + default=None, + metadata={"help": "Destination directory for merged model (defaults to _merged)"}, + ) + merge_overwrite: Optional[bool] = field( + default=False, + metadata={"help": "If True overwrite merged_output_dir if it already exists"}, + ) + + # Pair quality filtering (stronger separation & length to avoid trivial short answers) + min_chosen_chars: Optional[int] = field( + default=60, + metadata={"help": "Require chosen response to have at least this many characters"}, + ) + min_score_delta: Optional[float] = field( + default=1.5, + metadata={"help": "Minimum (best_score - worst_score) needed to keep a pair"}, + ) + + # instrumentation + sanity_check: Optional[bool] = field(default=False, metadata={"help": "only train on 1000 samples"}) + report_to: Optional[str] = field( + default="none", + metadata={ + "help": 'Integrations to report to (e.g. "wandb", "tensorboard", "none"). Default now "none" to avoid requiring extra installs.' + }, + ) + # debug argument for distributed training + ignore_bias_buffers: Optional[bool] = field( + default=False, + metadata={ + "help": "fix for DDP issues with LM bias/mask buffers - invalid scalar type,`inplace operation. See" + "https://github.com/huggingface/transformers/issues/22482#issuecomment-1595790992" + }, + ) + seed: Optional[int] = field( + default=0, metadata={"help": "Random seed that will be set at the beginning of training."} + ) + deepspeed: Optional[str] = field(default=None, metadata={"help": "the deepspeed json config file"}) + num_workers: Optional[int] = field(default=None, metadata={"help": "the number of workers to map the data"}) + + +def get_ultrafeedback_pairs( + start: int = 0, + sample_limit: int = 1000, + num_proc: int = 8, + max_prompt_length: int = 512, + max_total_length: int = 1024, + min_chosen_chars: int = 60, + min_score_delta: float = 1.5, +) -> Dataset: + """Construct (prompt, chosen, rejected) pairs from openbmb/UltraFeedback. + + Strategy: + * Load slice of the single 'train' split using HF slice syntax + * For each instruction, score each completion by averaging available numeric ratings + * Select best (chosen) and worst (rejected) response + * Filter out degenerate or overly short/long samples + """ + end = start + sample_limit + slice_str = f"train[{start}:{end}]" + raw = load_dataset("openbmb/UltraFeedback", split=slice_str) + original_columns = raw.column_names + + def score_completion(comp) -> float: + ann = comp.get("annotations", {}) + ratings = [] + for aspect, value in ann.items(): + if isinstance(value, list): + for item in value: + r = item.get("Rating") if isinstance(item, dict) else None + if r is not None: + try: + ratings.append(int(r)) + except Exception: + pass + elif isinstance(value, dict): + r = value.get("Rating") + if r is not None: + try: + ratings.append(int(r)) + except Exception: + pass + return sum(ratings) / len(ratings) if ratings else 0.0 + + def map_batch(batch) -> Dict[str, List[str]]: + prompts: List[str] = [] + chosens: List[str] = [] + rejecteds: List[str] = [] + for instr, completions in zip(batch.get("instruction", []), batch.get("completions", [])): + if not isinstance(instr, str) or not isinstance(completions, list) or len(completions) < 2: + continue + scored = [] + for c in completions: + if not isinstance(c, dict): + continue + resp = c.get("response", "") + if not isinstance(resp, str) or len(resp.strip()) < 20: + continue + s = score_completion(c) + length_bonus = 0.15 * (len(resp.strip()) / 100) # reward more substantive answers modestly + composite = s + length_bonus + scored.append((composite, resp, s)) + if len(scored) < 2: + continue + scored.sort(key=lambda x: x[0], reverse=True) + best_score, best_resp, raw_best = scored[0] + worst_score, worst_resp, raw_worst = scored[-1] + if best_resp.strip() == worst_resp.strip(): + continue + # Enforce minimum quality & separation + if len(best_resp.strip()) < min_chosen_chars: + continue + if (best_score - worst_score) < min_score_delta: + continue + # Additional constraint: prefer pairs where chosen answer is meaningfully longer than rejected + if len(best_resp) - len(worst_resp) < 40: + continue + prompt_text = instr.strip() + "\n\nAnswer:" + if len(prompt_text) > max_prompt_length: + continue + if (len(prompt_text) + len(best_resp)) > max_total_length or (len(prompt_text) + len(worst_resp)) > max_total_length: + continue + prompts.append(prompt_text) + chosens.append(best_resp) + rejecteds.append(worst_resp) + return {"prompt": prompts, "chosen": chosens, "rejected": rejecteds} + + ds = raw.map(map_batch, batched=True, num_proc=num_proc, remove_columns=original_columns) + ds = ds.filter(lambda x: len(x["prompt"]) > 0 and len(x["chosen"]) > 0 and len(x["rejected"]) > 0) + return ds + + +if __name__ == "__main__": + # Suppress repetitive deprecation warnings about Trainer.tokenizer + warnings.filterwarnings( + "ignore", + message="Trainer.tokenizer is now deprecated. You should use Trainer.processing_class instead.", + category=UserWarning, + ) + parser = HfArgumentParser(ScriptArguments) + script_args = parser.parse_args_into_dataclasses()[0] + + # Determine reporting integrations (disable by default; enable wandb only if requested) + env_enable_wandb = os.environ.get("ENABLE_WANDB") + report_to = script_args.report_to + if env_enable_wandb and report_to in ("none", "None", ""): + report_to = "wandb" + # Normalize disabling values to empty list for HF Trainer + if report_to in ("none", "None", "", None): + report_to = [] + + # 1. initialize training arguments: + training_args = GaudiDPOConfig( + per_device_train_batch_size=script_args.per_device_train_batch_size, + per_device_eval_batch_size=script_args.per_device_eval_batch_size, + max_steps=script_args.max_steps, + logging_steps=script_args.logging_steps, + save_steps=script_args.save_steps, + gradient_accumulation_steps=script_args.gradient_accumulation_steps, + gradient_checkpointing=script_args.gradient_checkpointing, + learning_rate=script_args.learning_rate, + eval_strategy="steps", + eval_steps=script_args.eval_steps, + output_dir=script_args.output_dir, + report_to=report_to, + lr_scheduler_type=script_args.lr_scheduler_type, + warmup_steps=script_args.warmup_steps, + optim=script_args.optimizer_type, + bf16=True, + remove_unused_columns=False, + run_name="dpo_llama2", + gradient_checkpointing_kwargs={"use_reentrant": script_args.gradient_checkpointing_use_reentrant}, + use_habana=True, + use_lazy_mode=True, + use_hpu_graphs_for_training=not script_args.gradient_checkpointing and (not script_args.deepspeed), + use_hpu_graphs_for_inference=not script_args.deepspeed, + seed=script_args.seed, + deepspeed=script_args.deepspeed, + overwrite_output_dir=True, + ) + + # Set seed before initializing model. + set_seed(training_args.seed) + + low_cpu_mem_usage = True + if is_deepspeed_available(): + from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled + + if is_deepspeed_zero3_enabled(): + low_cpu_mem_usage = False + + # 2. load a pretrained model + model = AutoModelForCausalLM.from_pretrained( + script_args.model_name_or_path, + low_cpu_mem_usage=low_cpu_mem_usage, + torch_dtype=torch.bfloat16, + ) + model.config.use_cache = False + model.config.use_fused_rope = False + + if script_args.ignore_bias_buffers: + # torch distributed hack + model._ddp_params_and_buffers_to_ignore = [ + name for name, buffer in model.named_buffers() if buffer.dtype == torch.bool + ] + + model_ref = AutoModelForCausalLM.from_pretrained( + script_args.model_name_or_path, + low_cpu_mem_usage=low_cpu_mem_usage, + torch_dtype=torch.bfloat16, + ) + model_ref.config.use_cache = False + tokenizer = AutoTokenizer.from_pretrained(script_args.tokenizer_name_or_path) + tokenizer.pad_token = tokenizer.eos_token + + # 3. Load UltraFeedback pairs from train split slices + train_dataset = get_ultrafeedback_pairs( + start=0, + sample_limit=1000, + num_proc=script_args.num_workers or 8, + max_prompt_length=script_args.max_prompt_length, + max_total_length=script_args.max_length, + min_chosen_chars=script_args.min_chosen_chars, + min_score_delta=script_args.min_score_delta, + ) + # 4. Evaluation slice (non-overlapping range) + eval_dataset = get_ultrafeedback_pairs( + start=1200, + sample_limit=200, + num_proc=script_args.num_workers or 8, + max_prompt_length=script_args.max_prompt_length, + max_total_length=script_args.max_length, + min_chosen_chars=script_args.min_chosen_chars, + min_score_delta=script_args.min_score_delta, + ) + try: + print(f"[UltraFeedback][Dataset] train_pairs={len(train_dataset)} eval_pairs={len(eval_dataset)}") + except Exception: + pass + + peft_config = LoraConfig( + r=script_args.lora_r, + lora_alpha=script_args.lora_alpha, + lora_dropout=script_args.lora_dropout, + target_modules=script_args.lora_target_modules, + bias="none", + task_type="CAUSAL_LM", + ) + + gaudi_config = GaudiConfig() + gaudi_config.use_fused_adam = True + gaudi_config.use_fused_clip_norm = True + + # 5. initialize the DPO trainer + dpo_trainer = GaudiDPOTrainer( + model, + model_ref, + gaudi_config=gaudi_config, + args=training_args, + beta=script_args.beta, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + tokenizer=tokenizer, + peft_config=peft_config, + max_prompt_length=script_args.max_prompt_length, + max_length=script_args.max_length, + force_use_ref_model=True, + ) + # Assign processing_class to avoid downstream tokenizer property access + try: + setattr(dpo_trainer, 'processing_class', tokenizer) + except Exception: + pass + + # 6. train + train_result = dpo_trainer.train() + + # 7. save + dpo_trainer.save_model(script_args.output_dir) + + # 8. save metric + metrics = train_result.metrics + dpo_trainer.log_metrics("train", metrics) + dpo_trainer.save_metrics("train", metrics) + + # 9. Optional merge of adapter for faster downstream loading + if script_args.merge_adapter_after_train: + try: + from peft import PeftConfig, PeftModel + import os + adapter_dir = script_args.output_dir + peft_conf = PeftConfig.from_pretrained(adapter_dir) + # Only proceed if this is actually a PEFT adapter (safetensors or bin present) + marker_files = [ + os.path.join(adapter_dir, f) for f in ["adapter_model.safetensors", "adapter_model.bin"] + ] + if not any(os.path.exists(m) for m in marker_files): + print(f"[MergeInline][DPO] Adapter markers not found in {adapter_dir}; skipping merge.") + else: + base_loader = AutoModelForCausalLM + base_fresh = base_loader.from_pretrained( + script_args.model_name_or_path, + low_cpu_mem_usage=True, + torch_dtype=torch.bfloat16, + ) + merged_model = PeftModel.from_pretrained(base_fresh, adapter_dir) + print("[MergeInline][DPO] Adapter loaded into fresh base model; merging...") + merged_model = merged_model.merge_and_unload() + out_dir = script_args.merged_output_dir or f"{script_args.output_dir}_merged" + if os.path.exists(out_dir) and not script_args.merge_overwrite: + print(f"[MergeInline][DPO] Output dir {out_dir} exists and merge_overwrite=False -> abort merge") + else: + os.makedirs(out_dir, exist_ok=True) + merged_model.save_pretrained(out_dir) + tokenizer.save_pretrained(out_dir) + print(f"[MergeInline][DPO] Merged model saved to {out_dir}") + except Exception as e: + print(f"[MergeInline][DPO] Merge failed: {e}") \ No newline at end of file diff --git a/fine_tuning/alignment/dpo_latest/dpo_pipeline_sanity.sh b/fine_tuning/alignment/dpo_latest/dpo_pipeline_sanity.sh new file mode 100755 index 0000000..128ec1d --- /dev/null +++ b/fine_tuning/alignment/dpo_latest/dpo_pipeline_sanity.sh @@ -0,0 +1,156 @@ +#!/usr/bin/env bash +set -euo pipefail + +START_PIPELINE=$(date +%s) + +# Fast DPO pipeline (override via env if desired) +DPO_MAX_STEPS=${DPO_MAX_STEPS:-140} +DPO_LOGGING_STEPS=${DPO_LOGGING_STEPS:-10} +DPO_LORA_TARGET=${DPO_LORA_TARGET:-"q_proj k_proj v_proj o_proj gate_proj up_proj down_proj"} + +# Base model (now mandatory: pass as first arg OR export MODEL_NAME) +MODEL_NAME="${1:-${MODEL_NAME:-}}" +if [ -z "${MODEL_NAME}" ]; then + echo "ERROR: Base model not specified. Provide as first argument or export MODEL_NAME." + echo "Usage: MODEL_NAME=google/gemma-3-270m ./dpo_pipeline_sanity.sh" >&2 + echo " or: ./dpo_pipeline_sanity.sh google/gemma-3-270m" >&2 + exit 1 +fi +echo "[CONFIG] MODEL_NAME=${MODEL_NAME}" +declare -A STEP_START +declare -A STEP_DURATION + +time_step_begin() { + local key="$1"; STEP_START[$key]=$(date +%s) + echo "[TIMER] BEGIN $key at $(date -u '+%Y-%m-%dT%H:%M:%SZ')" +} + +time_step_end() { + local key="$1"; local end=$(date +%s); local dur=$(( end - STEP_START[$key] )) + STEP_DURATION[$key]=$dur + echo "[TIMER] END $key at $(date -u '+%Y-%m-%dT%H:%M:%SZ') (duration=${dur}s)" +} + +LOGFILE="dpo_pipeline_sanity.log" +exec > >(tee -a "$LOGFILE") 2>&1 + +banner() { + echo + echo "============================================================" + echo ">>> $1" + echo "============================================================" +} + +# Base model (now mandatory: pass as first arg OR export MODEL_NAME) +MODEL_NAME="${1:-${MODEL_NAME:-}}" +if [ -z "${MODEL_NAME}" ]; then + echo "ERROR: Base model not specified. Provide as first argument or export MODEL_NAME." + echo "Usage: MODEL_NAME=google/gemma-3-270m ./ppo_pipeline_sanity.sh" >&2 + echo " or: ./ppo_pipeline_sanity.sh google/gemma-3-270m" >&2 + exit 1 +fi +echo "[CONFIG] MODEL_NAME=${MODEL_NAME}" + +# Common HPU env vars +export PT_HPU_DISABLE_FUSED_ADAMW=1 +export PT_HPU_DISABLE_FUSED_ADAM=1 +export PT_HPU_ENABLE_REFINE_DYNAMIC_SHAPES=0 + +if [ "${SKIP_SFT:-0}" = "1" ]; then + banner "STEP 1: Supervised Fine-Tuning SKIPPED (using base model directly)" + # Create a symlink named sft_sanity_merged pointing to the base model directory (if local) or just reuse MODEL_NAME if remote + if [ -d "${MODEL_NAME}" ]; then + ln -sfn "${MODEL_NAME}" ./sft_sanity_merged + else + # Remote model: create marker directory referencing base + mkdir -p ./sft_sanity_merged + echo "REMOTE_MODEL=${MODEL_NAME}" > ./sft_sanity_merged/BASE_MODEL_REFERENCE + fi +else + banner "STEP 1: Supervised Fine-Tuning (Ultra-Fast)" + time_step_begin SFT + python sft.py \ + --model_name_or_path "${MODEL_NAME}" \ + --dataset_name openbmb/UltraFeedback \ + --subset None \ + --streaming False \ + --size_valid_set 200 \ + --output_dir ./sft_sanity \ + --do_train \ + --max_steps 20 \ + --logging_steps 5 \ + --save_steps 999999 \ + --per_device_train_batch_size 2 \ + --gradient_accumulation_steps 1 \ + --learning_rate 3e-5 \ + --lr_scheduler_type cosine \ + --warmup_steps 2 \ + --weight_decay 0.01 \ + --lora_target_modules "q_proj" "k_proj" "v_proj" "o_proj" "gate_proj" "up_proj" "down_proj" \ + --bf16 \ + --remove_unused_columns False \ + --report_to none \ + --use_habana \ + --use_lazy_mode \ + --demo_fast \ + --merge_adapter_after_train \ + --merged_output_dir ./sft_sanity_merged + time_step_end SFT +fi + +banner "STEP 2: DPO Training" +time_step_begin DPO +PT_HPU_LAZY_MODE=1 python dpo.py \ + --model_name_or_path ./sft_sanity_merged \ + --tokenizer_name_or_path "${MODEL_NAME}" \ + --output_dir ./dpo_sanity \ + --beta 0.05 \ + --max_steps ${DPO_MAX_STEPS} \ + --learning_rate 2e-5 \ + --warmup_steps 10 \ + --per_device_train_batch_size 2 \ + --per_device_eval_batch_size 2 \ + --gradient_accumulation_steps 1 \ + --logging_steps ${DPO_LOGGING_STEPS} \ + --save_steps 999999 \ + --eval_steps 999999 \ + --lora_target_modules ${DPO_LORA_TARGET} \ + --report_to none \ + --seed 42 \ + --merge_adapter_after_train True \ + --merged_output_dir ./dpo_sanity_merged +time_step_end DPO + +banner "STEP 3: Comparing Base and DPO Models" +time_step_begin COMPARE +PT_HPU_LAZY_MODE=1 python compare_base_vs_dpo.py \ + --base_model "${MODEL_NAME}" \ + --finetuned_model $( [ -d ./dpo_sanity_merged ] && echo ./dpo_sanity_merged || echo ./dpo_sanity ) \ + --seed 123 \ + --max_new_tokens 128 \ + --min_new_tokens 60 \ + --stackexchange_style \ + --heuristics \ + --show_text +time_step_end COMPARE + +PIPELINE_END=$(date +%s) +TOTAL_DUR=$(( PIPELINE_END - START_PIPELINE )) + +echo +echo "=================== PIPELINE TIMING SUMMARY ===================" +printf "%-20s %10s\n" "Stage" "Seconds" +printf "%-20s %10s\n" "-----" "-------" +for k in SFT DPO COMPARE; do + printf "%-20s %10s\n" "$k" "${STEP_DURATION[$k]:-n/a}" +done +printf "%-20s %10s\n" "TOTAL" "$TOTAL_DUR" +echo "===============================================================" + +echo +echo "=================== ARTIFACT SUMMARY ===================" +if [ -d ./sft_sanity_merged ]; then echo "SFT merged policy: ./sft_sanity_merged"; fi +if [ -d ./dpo_sanity ]; then echo "DPO adapter policy: ./dpo_sanity"; fi +echo "========================================================" + +banner "✅ DPO PIPELINE COMPLETED SUCCESSFULLY (WITH COMPARISON)" diff --git a/fine_tuning/alignment/dpo_latest/logging_utils.py b/fine_tuning/alignment/dpo_latest/logging_utils.py new file mode 100644 index 0000000..271a53e --- /dev/null +++ b/fine_tuning/alignment/dpo_latest/logging_utils.py @@ -0,0 +1,24 @@ +import logging +import time + + +class _UTCFormatter(logging.Formatter): + converter = time.gmtime # type: ignore[attr-defined] + + +def setup_logging(level: int = logging.INFO) -> None: + """ + Initialize root logger with consistent UTC ISO-like timestamps if not already configured. + Safe to call multiple times (no duplicate handlers added). + """ + root = logging.getLogger() + if root.handlers: + return + fmt = "[%(asctime)s] %(levelname)s %(name)s: %(message)s" + datefmt = "%Y-%m-%dT%H:%M:%S" + handler = logging.StreamHandler() + handler.setFormatter(_UTCFormatter(fmt, datefmt)) + root.setLevel(level) + root.addHandler(handler) + +__all__ = ["setup_logging"] diff --git a/fine_tuning/alignment/dpo_latest/merge_peft_adapter.py b/fine_tuning/alignment/dpo_latest/merge_peft_adapter.py new file mode 100644 index 0000000..e0f36bb --- /dev/null +++ b/fine_tuning/alignment/dpo_latest/merge_peft_adapter.py @@ -0,0 +1,63 @@ +# Adapted from https://github.com/huggingface/optimum-habana/tree/v1.16.0/examples/trl +# Customized and enabled for Gaudi3 +from dataclasses import dataclass, field +import logging +from typing import Optional + +import torch +from peft import PeftConfig, PeftModel +from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer, HfArgumentParser + + +@dataclass +class ScriptArguments: + """ + The input names representing the Adapter and Base model fine-tuned with PEFT, and the output name representing the + merged model. + """ + + adapter_model_name: Optional[str] = field(default=None, metadata={"help": "the adapter name"}) + base_model_name: Optional[str] = field(default=None, metadata={"help": "the base model name"}) + output_name: Optional[str] = field(default=None, metadata={"help": "the merged model name"}) + + +from logging_utils import setup_logging +parser = HfArgumentParser(ScriptArguments) +script_args = parser.parse_args_into_dataclasses()[0] +setup_logging() +logger = logging.getLogger(__name__) +logger.info(f"ScriptArguments: {script_args}") +assert script_args.adapter_model_name is not None, "please provide the name of the Adapter you would like to merge" +assert script_args.base_model_name is not None, "please provide the name of the Base model" +assert script_args.output_name is not None, "please provide the output name of the merged model" + +if not (hasattr(torch, 'hpu') and torch.hpu.is_available()): + raise RuntimeError('[HPU][Required] Habana HPU not available. merge_peft_adapter enforces HPU use only.') +peft_config = PeftConfig.from_pretrained(script_args.adapter_model_name) +logger.info(f"[Adapter] Loaded PEFT config from {script_args.adapter_model_name} (task_type={peft_config.task_type})") +if peft_config.task_type == "SEQ_CLS": + logger.info("[IF] task_type == SEQ_CLS -> loading AutoModelForSequenceClassification (reward model style)") + # The sequence classification task is used for the reward model in PPO + model = AutoModelForSequenceClassification.from_pretrained( + script_args.base_model_name, num_labels=1, torch_dtype=torch.bfloat16 + ) +else: + logger.info("[IF] task_type != SEQ_CLS -> loading AutoModelForCausalLM") + model = AutoModelForCausalLM.from_pretrained( + script_args.base_model_name, return_dict=True, torch_dtype=torch.bfloat16 + ) + +tokenizer = AutoTokenizer.from_pretrained(script_args.base_model_name) + +# Load the PEFT model +model = PeftModel.from_pretrained(model, script_args.adapter_model_name) +logger.info(f"[Merge] Adapter weights loaded from {script_args.adapter_model_name}") +model.eval() + +model = model.to('hpu').merge_and_unload() +logger.info("[Merge] merge_and_unload complete -> adapter merged into base model") + +model.save_pretrained(f"{script_args.output_name}") +tokenizer.save_pretrained(f"{script_args.output_name}") +logger.info(f"[Save] Merged model + tokenizer saved to {script_args.output_name}") +# model.push_to_hub(f"{script_args.output_name}", use_temp_dir=False) diff --git a/fine_tuning/alignment/dpo_latest/sft.py b/fine_tuning/alignment/dpo_latest/sft.py new file mode 100644 index 0000000..4ec016f --- /dev/null +++ b/fine_tuning/alignment/dpo_latest/sft.py @@ -0,0 +1,462 @@ +# Adapted from https://github.com/huggingface/optimum-habana/tree/v1.16.0/examples/trl +# Customized and enabled for Gaudi3 +import logging +import warnings +import math +from dataclasses import dataclass, field +from typing import List, Optional + +import torch +import transformers +from datasets import load_dataset +from peft import LoraConfig +from tqdm import tqdm +from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser +from transformers.integrations.deepspeed import ( + is_deepspeed_available, +) + +from optimum.habana import GaudiConfig +from optimum.habana.trl import GaudiSFTConfig, GaudiSFTTrainer +from optimum.habana.utils import set_seed +from logging_utils import setup_logging + + +logger = logging.getLogger(__name__) + + +@dataclass +class ScriptArguments: + model_name_or_path: Optional[str] = field(default=None, metadata={"help": "the model name (REQUIRED)"}) + dataset_name: Optional[str] = field(default=None, metadata={"help": "the dataset name"}) + use_peft: Optional[bool] = field(default=True, metadata={"help": "whether to use peft"}) + subset: Optional[str] = field(default=None, metadata={"help": "the subset (local data_dir) if applicable; None for remote HF datasets"}) + split: Optional[str] = field(default="train", metadata={"help": "the split to use"}) + size_valid_set: Optional[int] = field(default=4000, metadata={"help": "the size of the validation set"}) + streaming: Optional[bool] = field(default=True, metadata={"help": "whether to stream the dataset"}) + shuffle_buffer: Optional[int] = field(default=5000, metadata={"help": "the shuffle buffer size"}) + num_workers: Optional[int] = field(default=4, metadata={"help": "the number of workers"}) + num_buckets: Optional[int] = field(default=-1, metadata={"help": "whether to use bucketing for SFTTrainer"}) + validation_split_percentage: Optional[int] = field( + default=5, + metadata={ + "help": "The percentage of the train set used as validation set in case there's no validation split" + }, + ) + use_flash_attention: Optional[bool] = field( + default=False, metadata={"help": "Whether to use Habana flash attention for fine-tuning."} + ) + flash_attention_recompute: Optional[bool] = field( + default=False, metadata={"help": "Whether to enable recompute in Habana flash attention for fine-tuning."} + ) + flash_attention_causal_mask: Optional[bool] = field( + default=False, metadata={"help": "Whether to enable causal mask in Habana flash attention for fine-tuning."} + ) + + # LoraConfig + lora_alpha: Optional[float] = field(default=16, metadata={"help": "the lora alpha parameter"}) + lora_dropout: Optional[float] = field(default=0.05, metadata={"help": "the lora dropout parameter"}) + lora_r: Optional[int] = field(default=8, metadata={"help": "the lora r parameter"}) + lora_target_modules: List[str] = field( + default_factory=lambda: None, + metadata={"help": "Target modules for the LoRA method."}, + ) + + token: str = field( + default=None, + metadata={ + "help": ( + "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token " + "generated when running `huggingface-cli login` (stored in `~/.huggingface`)." + ) + }, + ) + merge_adapter_after_train: bool = field( + default=True, metadata={"help": "If True and PEFT used, merge LoRA adapter into base model after training."} + ) + merged_output_dir: Optional[str] = field( + default=None, metadata={"help": "Directory to save merged model (defaults to _merged)."} + ) + merge_overwrite: bool = field( + default=False, metadata={"help": "If True overwrite merged_output_dir if it exists."} + ) + demo_fast: bool = field( + default=False, metadata={"help": "Enable ultra-fast demo mode: shrink dataset & steps for <20m pipeline."} + ) + + +if __name__ == "__main__": + # Silence deprecated Trainer.tokenizer warnings + warnings.filterwarnings( + "ignore", + message="Trainer.tokenizer is now deprecated. You should use Trainer.processing_class instead.", + category=UserWarning, + ) + parser = HfArgumentParser((ScriptArguments, GaudiSFTConfig)) + script_args, training_args = parser.parse_args_into_dataclasses() + if not script_args.model_name_or_path: + raise ValueError("--model_name_or_path is required (no default). Provide a model path or repo id.") + if not (hasattr(torch, 'hpu') and torch.hpu.is_available()): + raise RuntimeError('[HPU][Required] Habana HPU not available. SFT script enforces HPU-only execution.') + if getattr(training_args, 'use_habana', False) is False: + logging.getLogger(__name__).warning('[HPU][Override] Forcing use_habana=True (Always Use HPU policy)') + training_args.use_habana = True + setup_logging() + + # Validate required arguments + if not script_args.model_name_or_path: + raise ValueError("--model_name_or_path is required. Please specify the base model path.") + + # Demo-fast overrides applied early (before logging full args) + if script_args.demo_fast: + logger.warning("[DemoFast][SFT] Activating ultra-fast overrides (<20m goal)") + # Reduce training steps if large + if training_args.max_steps == -1 or training_args.max_steps > 20: + logger.info("[DemoFast][SFT] Setting max_steps=20") + training_args.max_steps = 20 + # Force single epoch semantics when steps specified + if getattr(training_args,'num_train_epochs',1) > 1: + training_args.num_train_epochs = 1 + logger.info("[DemoFast][SFT] Forcing num_train_epochs=1") + # Reduce warmup + if getattr(training_args,'warmup_steps',0) > 2: + training_args.warmup_steps = 2 + logger.info("[DemoFast][SFT] Setting warmup_steps=2") + # Lower logging cadence + if getattr(training_args,'logging_steps',10) > 5: + training_args.logging_steps = 5 + logger.info("[DemoFast][SFT] Setting logging_steps=5") + # Min gradient accumulation + if getattr(training_args,'gradient_accumulation_steps',1) > 1: + training_args.gradient_accumulation_steps = 1 + logger.info("[DemoFast][SFT] Setting gradient_accumulation_steps=1") + # Limit LoRA rank + if script_args.use_peft and script_args.lora_r > 16: + logger.info("[DemoFast][SFT] Reducing lora_r to 16") + script_args.lora_r = 16 + # Reduce shuffle buffer if streaming + if script_args.streaming and script_args.shuffle_buffer > 1000: + script_args.shuffle_buffer = 1000 + logger.info("[DemoFast][SFT] Setting shuffle_buffer=1000") + # Shrink validation set if huge + if script_args.size_valid_set > 1000: + script_args.size_valid_set = 500 + logger.info("[DemoFast][SFT] Setting size_valid_set=500") + # Tag merged output dir for clarity + if script_args.merged_output_dir is None: + script_args.merged_output_dir = "./sft_sanity_merged" + logger.info(f"ScriptArguments: {script_args}") + logger.info(f"TrainingArguments: {training_args}") + if script_args.use_peft: + logger.info("[IF] use_peft == True -> configuring LoRA") + peft_config = LoraConfig( + r=script_args.lora_r, + lora_alpha=script_args.lora_alpha, + lora_dropout=script_args.lora_dropout, + target_modules=script_args.lora_target_modules, + bias="none", + task_type="CAUSAL_LM", + ) + else: + logger.info("[IF] use_peft == False -> full fine-tune (no LoRA)") + peft_config = None + + if training_args.group_by_length and training_args.packing: + logger.warning("[IF] group_by_length AND packing are both True -> raising ValueError") + raise ValueError("Cannot use both packing and group by length") + + set_seed(training_args.seed) + + def chars_token_ratio(dataset, tokenizer, nb_examples=400): + """ + Estimate the average number of characters per token in the dataset. + """ + total_characters, total_tokens = 0, 0 + for _, example in tqdm(zip(range(nb_examples), iter(dataset)), total=nb_examples): + text = prepare_sample_text(example) + total_characters += len(text) + if tokenizer.is_fast: + total_tokens += len(tokenizer(text).tokens()) + else: + total_tokens += len(tokenizer.tokenize(text)) + + return total_characters / total_tokens + + def prepare_sample_text(example): + """Prepare the text from a sample of supported datasets. + + For stack-exchange-paired: uses question + response_j. + For ultrafeedback_binarized: use prompt + chosen (preferred answer). + Falls back to joining all string fields if expected keys absent. + """ + if 'question' in example and 'response_j' in example: + return f"Question: {example['question']}\n\nAnswer: {example['response_j']}" + if 'prompt' in example and 'chosen' in example: + return f"{example['prompt'].rstrip()}\n{example['chosen']}" + # Fallback: concatenate first two string fields + for k in example: + if isinstance(example[k], str): + base = example[k] + break + else: + base = "" + extras = [v for k,v in example.items() if isinstance(v,str) and v is not base] + return base + ("\n" + "\n".join(extras) if extras else "") + + def create_datasets(tokenizer, args, seed=None): + if args.dataset_name: + logger.info(f"[IF] dataset_name provided -> loading {args.dataset_name}") + load_kwargs = { + 'split': args.split, + 'token': script_args.token, + 'num_proc': args.num_workers if not args.streaming else None, + 'streaming': args.streaming, + } + # Only pass data_dir if explicitly provided and not None (for local custom datasets) + if args.subset not in (None, 'None', ''): + load_kwargs['data_dir'] = args.subset + dataset = load_dataset(args.dataset_name, **load_kwargs) + # UltraFeedback: map to single (instruction, best_response) for SFT + if args.dataset_name == 'openbmb/UltraFeedback': + logger.info('[UltraFeedback][SFT] Mapping completions to best scored response') + def _score_completion(comp): + ann = comp.get('annotations', {}) + ratings = [] + for aspect, val in ann.items(): + if isinstance(val, list): + for item in val: + if isinstance(item, dict) and 'Rating' in item: + try: + ratings.append(int(item['Rating'])) + except Exception: + pass + elif isinstance(val, dict) and 'Rating' in val: + try: + ratings.append(int(val['Rating'])) + except Exception: + pass + return sum(ratings)/len(ratings) if ratings else 0.0 + def _map_batch(batch): + out_instr = [] + out_resp = [] + for instr, comps in zip(batch.get('instruction', []), batch.get('completions', [])): + if not isinstance(instr, str) or not isinstance(comps, list) or len(comps) == 0: + continue + scored = [] + for c in comps: + if not isinstance(c, dict): + continue + resp = c.get('response', '') + if not isinstance(resp, str) or len(resp.strip()) < 10: + continue + scored.append((_score_completion(c), resp)) + if len(scored) == 0: + continue + scored.sort(key=lambda x: x[0], reverse=True) + best_resp = scored[0][1] + out_instr.append(instr) + out_resp.append(best_resp) + return {'question': out_instr, 'response_j': out_resp} + original_cols = dataset.column_names + dataset = dataset.map(_map_batch, batched=True, remove_columns=original_cols) + logger.info(f"[UltraFeedback][SFT] After mapping rows={len(dataset)} cols={dataset.column_names}") + else: + logger.error("[IF] dataset_name missing -> raising ValueError") + raise ValueError("No dataset_name") + if args.streaming: + logger.info("[IF] streaming == True -> using .take/.skip and shuffle") + logger.info("Loading the dataset in streaming mode") + valid_data = dataset.take(args.size_valid_set) + train_data = dataset.skip(args.size_valid_set) + train_data = train_data.shuffle(buffer_size=args.shuffle_buffer, seed=seed) + if script_args.demo_fast: + # For streaming, further restrict training iterable + fast_limit = 50 + logger.info(f"[DemoFast][SFT] Limiting streaming train_data to first {fast_limit} samples") + train_data = train_data.take(fast_limit) + valid_data = valid_data.take(min(script_args.size_valid_set, 50)) + else: + logger.info("[IF] streaming == False -> train_test_split path") + dataset = dataset.train_test_split(test_size=args.validation_split_percentage * 0.01, seed=seed) + train_data = dataset["train"] + valid_data = dataset["test"] + logger.info(f"Size of the train set: {len(train_data)}. Size of the validation set: {len(valid_data)}") + if script_args.demo_fast: + # Subsample train & valid sets + import random + fast_train_limit = min(80, len(train_data)) + fast_valid_limit = min(20, len(valid_data)) + logger.info(f"[DemoFast][SFT] Subsampling train->{fast_train_limit} valid->{fast_valid_limit}") + train_indices = list(range(len(train_data)))[:fast_train_limit] + valid_indices = list(range(len(valid_data)))[:fast_valid_limit] + train_data = train_data.select(train_indices) + valid_data = valid_data.select(valid_indices) + if args.dataset_name in ("lvwerra/stack-exchange-paired", "openbmb/UltraFeedback"): + logger.info(f"[IF] dataset_name {args.dataset_name} -> computing chars/token ratio and using formatting func") + try: + chars_per_token = chars_token_ratio(train_data, tokenizer) + logger.info(f"[DatasetStats] chars/token ratio: {chars_per_token:.2f}") + except Exception as e: + logger.warning(f"[DatasetStats] Failed chars/token ratio estimation: {e}") + formating_func = prepare_sample_text + else: + logger.info(f"[IF] dataset_name {args.dataset_name} not in specialized list -> no formatting func") + formating_func = None + return train_data, valid_data, formating_func + + low_cpu_mem_usage = True + if is_deepspeed_available(): + logger.info("[IF] Deepspeed available -> checking Zero3") + from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled + + if is_deepspeed_zero3_enabled(): + logger.info("[IF] DeepSpeed Zero3 enabled -> disabling low_cpu_mem_usage") + low_cpu_mem_usage = False + + base_model = AutoModelForCausalLM.from_pretrained( + script_args.model_name_or_path, + low_cpu_mem_usage=low_cpu_mem_usage, + torch_dtype=torch.bfloat16, + token=script_args.token, + ) + mt = getattr(base_model.config, 'model_type', '') + if mt.startswith('gemma3'): + try: + base_model.config.attn_implementation = 'eager' + except Exception: + pass + + # --- Parameter statistics before (and later after) PEFT injection --- + def _param_stats(model): + total = 0 + trainable = 0 + for p in model.parameters(): + n = p.numel() + total += n + if p.requires_grad: + trainable += n + pct = (trainable / total * 100) if total else 0.0 + return total, trainable, pct + + tot, trn, pct = _param_stats(base_model) + logger.info(f"[Model Params][Base] total={tot:,} trainable={trn:,} ({pct:.4f}%)") + + base_model.config.use_cache = False + if not script_args.use_flash_attention and ( + script_args.flash_attention_recompute or script_args.flash_attention_recompute + ): + logger.warning("[IF] flash attention recompute flags set while use_flash_attention is False -> assert") + assert "Need to enable use_flash_attention" + base_model.generation_config.use_flash_attention = script_args.use_flash_attention + base_model.generation_config.flash_attention_recompute = script_args.flash_attention_recompute + base_model.generation_config.flash_attention_causal_mask = script_args.flash_attention_causal_mask + if script_args.use_flash_attention: + logger.info("[IF] use_flash_attention == True -> flash attention settings applied") + else: + logger.info("[IF] use_flash_attention == False -> skipping flash attention optimizations") + + tokenizer = AutoTokenizer.from_pretrained(script_args.model_name_or_path, trust_remote_code=True) + tokenizer.pad_token = tokenizer.eos_token + tokenizer.padding_side = "right" # Fix weird overflow issue with fp16 training + + log_level = training_args.get_process_log_level() + logger.setLevel(log_level) + transformers.utils.logging.set_verbosity(log_level) + transformers.utils.logging.enable_default_handler() + transformers.utils.logging.enable_explicit_format() + + train_dataset, eval_dataset, formatting_func = create_datasets(tokenizer, script_args, seed=training_args.seed) + + gaudi_config = GaudiConfig() + gaudi_config.use_fused_adam = False + gaudi_config.use_fused_clip_norm = False + if training_args.do_train: + logger.info("[IF] do_train == True -> beginning training loop") + trainer = GaudiSFTTrainer( + model=base_model, + gaudi_config=gaudi_config, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + peft_config=peft_config, + tokenizer=tokenizer, + args=training_args, + formatting_func=formatting_func, + num_buckets=script_args.num_buckets, + ) + # Provide processing_class attribute expected by newer Trainer interface + try: + setattr(trainer, 'processing_class', tokenizer) + except Exception: + pass + + # If PEFT active, trainer.model now wrapped; print updated stats + try: + wrapped_tot, wrapped_trn, wrapped_pct = _param_stats(trainer.model) + logger.info(f"[Model Params][After Trainer Init] total={wrapped_tot:,} trainable={wrapped_trn:,} ({wrapped_pct:.4f}%)") + if peft_config is not None: + # Highlight LoRA target modules used + logger.info(f"[LoRA] Target modules: {peft_config.target_modules}") + logger.info(f"[LoRA] r={getattr(peft_config,'r',None)} alpha={getattr(peft_config,'lora_alpha',None)} dropout={getattr(peft_config,'lora_dropout',None)}") + except Exception as e: + logger.warning(f"[Warn] Could not compute PEFT parameter stats: {e}") + train_result = trainer.train() + trainer.save_model(training_args.output_dir) + metrics = train_result.metrics + trainer.log_metrics("train", metrics) + trainer.save_metrics("train", metrics) + logger.info("[INFO] Training completed and model saved") + # Optional inline merge + if script_args.merge_adapter_after_train: + if not script_args.use_peft: + logger.warning("[MergeInline] merge_adapter_after_train=True but use_peft=False -> skipping merge") + else: + try: + from peft import PeftConfig, PeftModel + import os + adapter_dir = training_args.output_dir + peft_conf = PeftConfig.from_pretrained(adapter_dir) + logger.info(f"[MergeInline] Loaded PEFT config task_type={peft_conf.task_type}") + base_loader = AutoModelForCausalLM + base_fresh = base_loader.from_pretrained( + script_args.model_name_or_path, + low_cpu_mem_usage=low_cpu_mem_usage, + torch_dtype=torch.bfloat16, + token=script_args.token, + ) + merged_model = PeftModel.from_pretrained(base_fresh, adapter_dir) + logger.info("[MergeInline] Adapter loaded into fresh base model; merging...") + merged_model = merged_model.merge_and_unload() + out_dir = script_args.merged_output_dir or f"{training_args.output_dir}_merged" + if os.path.exists(out_dir) and not script_args.merge_overwrite: + logger.warning(f"[MergeInline] Output dir {out_dir} exists and merge_overwrite=False -> abort merge") + else: + os.makedirs(out_dir, exist_ok=True) + merged_model.save_pretrained(out_dir) + tokenizer.save_pretrained(out_dir) + logger.info(f"[MergeInline] Merged model saved to {out_dir}") + except Exception as e: + logger.exception(f"[MergeInline] Failed to merge adapter inline: {e}") + else: + logger.info("[IF] do_train == False -> skipping training") + + # Evaluation + if training_args.do_eval: + logger.info("[IF] do_eval == True -> running evaluation") + logger.info("*** Evaluate ***") + metrics = trainer.evaluate() + if isinstance(eval_dataset, torch.utils.data.IterableDataset): + eval_dataset = list(eval_dataset) + + metrics["eval_samples"] = len(eval_dataset) + + try: + perplexity = math.exp(metrics["eval_loss"]) + except OverflowError: + perplexity = float("inf") + metrics["perplexity"] = perplexity + + trainer.log_metrics("eval", metrics) + trainer.save_metrics("eval", metrics) + logger.info("[INFO] Evaluation complete") + else: + logger.info("[IF] do_eval == False -> skipping evaluation")