| |
|
|
| import os |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch.utils.data import Dataset, DataLoader |
| from fastapi import FastAPI, HTTPException |
| from pydantic import BaseModel |
| from typing import List, Dict, Tuple |
| import uvicorn |
| import copy |
|
|
| from configs.model_config import SystemConfig, PathConfig |
| from src.model.llm import DecoderOnlyTransformer |
| from src.model.config import ModelConfig |
| from tokenizers import Tokenizer |
|
|
| app = FastAPI( |
| title="Generative AI Project - Alignment Service", |
| description="SFT & DPO Microservice supporting multi-GPU training", |
| version="1.0.0" |
| ) |
|
|
|
|
| @app.get("/health") |
| def health(): |
| """Lightweight liveness probe used by the gateway.""" |
| return {"status": "ok", "service": "alignment"} |
|
|
| class SFTTrainRequest(BaseModel): |
| sft_data: List[Dict[str, str]] |
| epochs: int = 3 |
| lr: float = 2e-5 |
| batch_size: int = 2 |
|
|
| class DPOTrainRequest(BaseModel): |
| dpo_data: List[Dict[str, str]] |
| epochs: int = 3 |
| lr: float = 5e-6 |
| beta: float = 0.1 |
| batch_size: int = 2 |
|
|
|
|
| class SFTDataset(Dataset): |
| def __init__(self, data: List[Dict[str, str]], tokenizer: Tokenizer, max_len: int = 512): |
| from src.common.prompts import format_prompt, resolve_eos_id, resolve_pad_id |
| self.samples = [] |
| eos_id = resolve_eos_id(tokenizer) |
| pad_id = resolve_pad_id(tokenizer) |
| for item in data: |
| prompt = format_prompt(item["prompt"]) |
| response = item["response"] |
|
|
| enc_prompt = tokenizer.encode(prompt).ids |
| enc_response = tokenizer.encode(response).ids |
| if eos_id is not None: |
| enc_response = enc_response + [eos_id] |
|
|
| input_ids = enc_prompt + enc_response |
| labels = [-100] * len(enc_prompt) + enc_response |
|
|
| if len(input_ids) > max_len: |
| input_ids = input_ids[:max_len] |
| labels = labels[:max_len] |
| else: |
| pad_len = max_len - len(input_ids) |
| input_ids = input_ids + [pad_id] * pad_len |
| labels = labels + [-100] * pad_len |
|
|
| self.samples.append((torch.tensor(input_ids, dtype=torch.long), torch.tensor(labels, dtype=torch.long))) |
|
|
| def __len__(self): |
| return len(self.samples) |
|
|
| def __getitem__(self, idx): |
| return self.samples[idx] |
|
|
|
|
| class DPODataset(Dataset): |
| def __init__(self, data: List[Dict[str, str]], tokenizer: Tokenizer, max_len: int = 512): |
| from src.common.prompts import format_prompt, resolve_eos_id, resolve_pad_id |
| self.samples = [] |
| pad_token = resolve_pad_id(tokenizer) |
| eos_id = resolve_eos_id(tokenizer) |
|
|
| for item in data: |
| prompt = format_prompt(item["prompt"]) |
| chosen = item["chosen"] |
| rejected = item["rejected"] |
|
|
| enc_prompt = tokenizer.encode(prompt).ids |
| enc_chosen = tokenizer.encode(chosen).ids |
| enc_rejected = tokenizer.encode(rejected).ids |
| if eos_id is not None: |
| enc_chosen = enc_chosen + [eos_id] |
| enc_rejected = enc_rejected + [eos_id] |
| |
| def pad(ids, labels): |
| if len(ids) > max_len: |
| return ids[:max_len], labels[:max_len] |
| pad_len = max_len - len(ids) |
| return ids + [pad_token] * pad_len, labels + [-100] * pad_len |
| |
| c_input_ids, c_labels = pad(enc_prompt + enc_chosen, [-100] * len(enc_prompt) + enc_chosen) |
| r_input_ids, r_labels = pad(enc_prompt + enc_rejected, [-100] * len(enc_prompt) + enc_rejected) |
| |
| self.samples.append({ |
| "c_ids": torch.tensor(c_input_ids, dtype=torch.long), |
| "c_labels": torch.tensor(c_labels, dtype=torch.long), |
| "r_ids": torch.tensor(r_input_ids, dtype=torch.long), |
| "r_labels": torch.tensor(r_labels, dtype=torch.long) |
| }) |
|
|
| def __len__(self): |
| return len(self.samples) |
|
|
| def __getitem__(self, idx): |
| return self.samples[idx] |
|
|
|
|
| def get_log_probs(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: |
| shift_logits = logits[..., :-1, :].contiguous() |
| shift_labels = labels[..., 1:].contiguous() |
| loss_mask = (shift_labels != -100) |
| |
| dummy_labels = shift_labels.clone() |
| dummy_labels[shift_labels == -100] = 0 |
| |
| log_probs = F.log_softmax(shift_logits, dim=-1) |
| per_token_log_probs = torch.gather(log_probs, dim=-1, index=dummy_labels.unsqueeze(-1)).squeeze(-1) |
| return (per_token_log_probs * loss_mask).sum(dim=-1) |
|
|
|
|
| def wrap_model_for_device(model: nn.Module) -> nn.Module: |
| """Wraps model in DataParallel dynamically if multiple GPUs are detected.""" |
| model = model.to(SystemConfig.DEVICE) |
| if SystemConfig.DEVICE == "cuda" and SystemConfig.NUM_GPUS > 1: |
| print(f"[ACCELERATION] Wrapping model in torch.nn.DataParallel across {SystemConfig.NUM_GPUS} GPUs.") |
| model = nn.DataParallel(model) |
| return model |
|
|
|
|
| @app.post("/sft") |
| def train_sft(payload: SFTTrainRequest): |
| if not os.path.exists(PathConfig.PRETRAINED_BASE_MODEL): |
| raise HTTPException(status_code=400, detail="Base model checkpoint not found. Run pretraining first.") |
| |
| tokenizer = Tokenizer.from_file(PathConfig.TOKENIZER_PATH) |
| checkpoint = torch.load(PathConfig.PRETRAINED_BASE_MODEL, map_location=SystemConfig.DEVICE) |
| model_config_dict = checkpoint.get('model_config', checkpoint.get('args')) |
| model_config = ModelConfig(**model_config_dict) if isinstance(model_config_dict, dict) else model_config_dict |
| |
| raw_model = DecoderOnlyTransformer(model_config) |
| model_state = checkpoint.get('model', checkpoint.get('model_state_dict')) |
| raw_model.load_state_dict(model_state) |
| model = wrap_model_for_device(raw_model) |
| model.train() |
| |
| dataset = SFTDataset(payload.sft_data, tokenizer, max_len=model_config.block_size) |
| dataloader = DataLoader(dataset, batch_size=payload.batch_size, shuffle=True) |
| optimizer = torch.optim.AdamW(model.parameters(), lr=payload.lr) |
| |
| print("Starting SFT training stage...") |
| for epoch in range(payload.epochs): |
| epoch_loss = 0.0 |
| for inputs, labels in dataloader: |
| inputs, labels = inputs.to(SystemConfig.DEVICE), labels.to(SystemConfig.DEVICE) |
| optimizer.zero_grad() |
| |
| logits, _ = model(inputs) |
| |
| shift_logits = logits[..., :-1, :].contiguous().view(-1, model_config.vocab_size) |
| shift_labels = labels[..., 1:].contiguous().view(-1) |
| loss = F.cross_entropy(shift_logits, shift_labels, ignore_index=-100) |
| |
| loss.backward() |
| optimizer.step() |
| epoch_loss += loss.item() |
| |
| print(f"SFT Epoch {epoch+1}/{payload.epochs} | Loss: {epoch_loss/len(dataloader):.4f}") |
| |
| state_to_save = model.module.state_dict() if isinstance(model, nn.DataParallel) else model.state_dict() |
| torch.save({"model_state_dict": state_to_save, "args": model_config_dict}, PathConfig.SFT_MODEL_CHECKPOINT) |
| return {"status": "SFT Complete", "checkpoint_path": PathConfig.SFT_MODEL_CHECKPOINT} |
|
|
|
|
| @app.post("/dpo") |
| def train_dpo(payload: DPOTrainRequest): |
| if not os.path.exists(PathConfig.SFT_MODEL_CHECKPOINT): |
| raise HTTPException(status_code=400, detail="SFT checkpoint not found. Run SFT first.") |
| |
| tokenizer = Tokenizer.from_file(PathConfig.TOKENIZER_PATH) |
| checkpoint = torch.load(PathConfig.SFT_MODEL_CHECKPOINT, map_location=SystemConfig.DEVICE) |
| model_config_dict = checkpoint.get('model_config', checkpoint.get('args')) |
| model_config = ModelConfig(**model_config_dict) if isinstance(model_config_dict, dict) else model_config_dict |
| |
| policy_raw = DecoderOnlyTransformer(model_config) |
| model_state = checkpoint.get('model', checkpoint.get('model_state_dict')) |
| policy_raw.load_state_dict(model_state) |
| policy_model = wrap_model_for_device(policy_raw) |
| policy_model.train() |
| |
| ref_raw = DecoderOnlyTransformer(model_config) |
| ref_raw.load_state_dict(copy.deepcopy(model_state)) |
| ref_model = wrap_model_for_device(ref_raw) |
| ref_model.eval() |
| |
| dataset = DPODataset(payload.dpo_data, tokenizer, max_len=model_config.block_size) |
| dataloader = DataLoader(dataset, batch_size=payload.batch_size, shuffle=True) |
| optimizer = torch.optim.AdamW(policy_model.parameters(), lr=payload.lr) |
| |
| print("Starting DPO training stage...") |
| for epoch in range(payload.epochs): |
| epoch_loss = 0.0 |
| for batch in dataloader: |
| optimizer.zero_grad() |
| |
| c_ids, c_labels = batch["c_ids"].to(SystemConfig.DEVICE), batch["c_labels"].to(SystemConfig.DEVICE) |
| r_ids, r_labels = batch["r_ids"].to(SystemConfig.DEVICE), batch["r_labels"].to(SystemConfig.DEVICE) |
| |
| p_c_logits, _ = policy_model(c_ids) |
| p_r_logits, _ = policy_model(r_ids) |
| policy_chosen_logps = get_log_probs(p_c_logits, c_labels) |
| policy_rejected_logps = get_log_probs(p_r_logits, r_labels) |
| |
| with torch.no_grad(): |
| ref_c_logits, _ = ref_model(c_ids) |
| ref_r_logits, _ = ref_model(r_ids) |
| ref_chosen_logps = get_log_probs(ref_c_logits, c_labels) |
| ref_rejected_logps = get_log_probs(ref_r_logits, r_labels) |
| |
| pi_log_ratio = policy_chosen_logps - policy_rejected_logps |
| ref_log_ratio = ref_chosen_logps - ref_rejected_logps |
| |
| logits_diff = pi_log_ratio - ref_log_ratio |
| loss = -F.logsigmoid(payload.beta * logits_diff).mean() |
| |
| loss.backward() |
| optimizer.step() |
| epoch_loss += loss.item() |
| |
| print(f"DPO Epoch {epoch+1}/{payload.epochs} | Loss: {epoch_loss/len(dataloader):.4f}") |
| |
| state_to_save = policy_model.module.state_dict() if isinstance(policy_model, nn.DataParallel) else policy_model.state_dict() |
| torch.save({"model_state_dict": state_to_save, "args": model_config_dict}, PathConfig.DPO_MODEL_CHECKPOINT) |
| return {"status": "DPO Alignment Complete", "checkpoint_path": PathConfig.DPO_MODEL_CHECKPOINT} |
|
|
|
|
| if __name__ == "__main__": |
| uvicorn.run(app, host="0.0.0.0", port=8002) |
|
|
