Add training pipeline v2.0 — DeepSeek-VL2-tiny × ChartQA LoRA

train_pipeline.py

@@ -0,0 +1,423 @@
+#!/usr/bin/env python3
+"""
+DeepSeek-VL2-tiny Chart Fine-tuning Pipeline v2.0
+═══════════════════════════════════════════════════
+Dataset  : HuggingFaceM4/ChartQA (fallback from YUNGHUI2024/deepseek-ocr2-chart-finetune)
+Model    : deepseek-ai/deepseek-vl2-tiny  (1B active / 3B total, bf16 ≈6.3 GB)
+Method   : LoRA (r=16, target q/k/v/o_proj) + gradient checkpointing
+VRAM     : Tested on RTX 3060 12 GB  (batch=1, grad_accum=16)
+Tracking : Trackio (optional) — set env vars:
+             TRACKIO_SPACE_ID, TRACKIO_PROJECT
+Output   : YUNGHUI2024/deepseek-vl2-tiny-chartqa-lora
+
+═══ 本機快速開始 ═══════════════════════════════════════════════════
+pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
+pip install "transformers>=4.40" "datasets>=2.18" peft accelerate trackio huggingface_hub pillow
+git clone https://github.com/deepseek-ai/DeepSeek-VL2 && cd DeepSeek-VL2 && pip install -e . && cd ..
+
+# 登入 HF Hub (push 用)
+huggingface-cli login
+
+python train_pipeline.py
+"""
+
+import os, sys, subprocess, logging, math
+from pathlib import Path
+import torch
+from datasets import load_dataset
+from PIL import Image
+from torch.utils.data import DataLoader, Dataset
+from transformers import AutoModelForCausalLM, get_cosine_schedule_with_warmup
+from torch.optim import AdamW
+from peft import LoraConfig, get_peft_model, TaskType
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s | %(levelname)s | %(message)s",
+    datefmt="%H:%M:%S",
+)
+log = logging.getLogger(__name__)
+
+# ─── optional: auto-install deepseek_vl if not found ─────────────────────────
+try:
+    from deepseek_vl.models import DeepseekVLV2Processor, DeepseekVLV2ForCausalLM
+except ImportError:
+    log.info("deepseek_vl not found — installing from GitHub …")
+    subprocess.run(
+        [sys.executable, "-m", "pip", "install", "-q",
+         "git+https://github.com/deepseek-ai/DeepSeek-VL2.git"],
+        check=True,
+    )
+    from deepseek_vl.models import DeepseekVLV2Processor, DeepseekVLV2ForCausalLM
+
+# ─── optional Trackio ─────────────────────────────────────────────────────────
+_USE_TRACKIO = bool(os.getenv("TRACKIO_SPACE_ID") or os.getenv("TRACKIO_PROJECT"))
+if _USE_TRACKIO:
+    import trackio
+
+def tlog(metrics: dict):
+    if _USE_TRACKIO:
+        trackio.log(metrics)
+
+def talert(title: str, text: str, level: str = "INFO"):
+    if _USE_TRACKIO:
+        trackio.alert(title=title, text=text, level=level)
+    log.info(f"[ALERT {level}] {title}: {text}")
+
+# ──────────────────────────────────────────────────────────────────────────────
+# ███  CONFIG  ████████████████████████████████████████████████████████████████
+# ──────────────────────────────────────────────────────────────────────────────
+MODEL_ID      = "deepseek-ai/deepseek-vl2-tiny"
+DATASET_ID    = "HuggingFaceM4/ChartQA"
+HUB_MODEL_ID  = "YUNGHUI2024/deepseek-vl2-tiny-chartqa-lora"
+OUTPUT_DIR    = "./output-deepseek-vl2-chartqa"   # local folder
+
+# LoRA
+LORA_R        = 16
+LORA_ALPHA    = 32
+LORA_DROPOUT  = 0.05
+LORA_TARGETS  = ["q_proj", "k_proj", "v_proj", "o_proj"]
+
+# Training — tuned for 12 GB VRAM (RTX 3060)
+LR            = 2e-4
+NUM_EPOCHS    = 2
+BATCH_SIZE    = 1          # per-GPU
+GRAD_ACCUM    = 16         # effective batch = 16
+LOG_EVERY     = 20         # opt-steps
+SAVE_STEPS    = 200        # opt-steps
+
+# Set to small int (e.g. 50) for a quick smoke-test; None = full dataset
+MAX_TRAIN     = None
+MAX_VAL       = 200        # cap val for speed
+
+# Trackio
+TRACKIO_SPACE = os.getenv("TRACKIO_SPACE_ID", "YUNGHUI2024/ml-intern-chartqa")
+TRACKIO_PROJ  = os.getenv("TRACKIO_PROJECT",  "deepseek-vl2-chartqa")
+RUN_NAME      = f"vl2tiny_lora_r{LORA_R}_lr{LR}"
+
+# ──────────────────────────────────────────────────────────────────────────────
+# ███  TRACKIO INIT  ██████████████████████████████████████████████████████████
+# ────────────────���─────────────────────────────────────────────────────────────
+if _USE_TRACKIO:
+    trackio.init(
+        project=TRACKIO_PROJ,
+        name=RUN_NAME,
+        space_id=TRACKIO_SPACE,
+        config={
+            "model": MODEL_ID, "dataset": DATASET_ID,
+            "lora_r": LORA_R, "lora_alpha": LORA_ALPHA,
+            "lr": LR, "epochs": NUM_EPOCHS,
+            "batch_size": BATCH_SIZE, "grad_accum": GRAD_ACCUM,
+        },
+    )
+    log.info(f"Trackio init — project={TRACKIO_PROJ}  run={RUN_NAME}")
+
+# ──────────────────────────────────────────────────────────────────────────────
+# ███  PROCESSOR & MODEL  █████████████████████████████████████████████████████
+# ──────────────────────────────────────────────────────────────────────────────
+log.info(f"Loading processor from {MODEL_ID} …")
+processor: DeepseekVLV2Processor = DeepseekVLV2Processor.from_pretrained(MODEL_ID)
+tokenizer = processor.tokenizer
+tokenizer.padding_side = "right"
+
+log.info(f"Loading model {MODEL_ID} → bf16 …")
+model: DeepseekVLV2ForCausalLM = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+)
+model.config.use_cache = False
+
+# Gradient checkpointing BEFORE LoRA wrapping (saves ~30–40% VRAM)
+if hasattr(model, "gradient_checkpointing_enable"):
+    model.gradient_checkpointing_enable()
+elif hasattr(model, "language_model"):
+    model.language_model.gradient_checkpointing_enable()
+
+# ──────────────────────────────────────────────────────────────────────────────
+# ███  LoRA  ██████████████████████████████████████████████████████████████████
+# ──────────────────────────────────────────────────────────────────────────────
+lora_cfg = LoraConfig(
+    task_type=TaskType.CAUSAL_LM,
+    r=LORA_R,
+    lora_alpha=LORA_ALPHA,
+    lora_dropout=LORA_DROPOUT,
+    target_modules=LORA_TARGETS,
+    bias="none",
+)
+model = get_peft_model(model, lora_cfg)
+model.print_trainable_parameters()
+
+trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+total     = sum(p.numel() for p in model.parameters())
+log.info(f"LoRA trainable: {trainable/1e6:.2f}M / {total/1e6:.0f}M "
+         f"({100*trainable/total:.2f}%)")
+
+# Move to GPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model  = model.to(device)
+log.info(f"Model on {device}")
+if torch.cuda.is_available():
+    mem = torch.cuda.memory_reserved() / 1e9
+    log.info(f"VRAM reserved after model load: {mem:.1f} GB")
+
+# ──────────────────────────────────────────────────────────────────────────────
+# ███  DATASET  ███████████████████████████████████████████████████████████████
+# ──────────────────────────────────────────────────────────────────────────────
+log.info(f"Loading {DATASET_ID} …")
+raw      = load_dataset(DATASET_ID)
+train_raw = raw["train"]
+val_raw   = raw["val"]
+
+if MAX_TRAIN:
+    train_raw = train_raw.select(range(MAX_TRAIN))
+val_raw = val_raw.select(range(min(len(val_raw), MAX_VAL)))
+log.info(f"Train: {len(train_raw):,}  Val: {len(val_raw):,}")
+tlog({"dataset/train_samples": len(train_raw), "dataset/val_samples": len(val_raw)})
+
+
+class ChartQADataset(Dataset):
+    def __init__(self, hf_ds): self.data = hf_ds
+    def __len__(self): return len(self.data)
+
+    def __getitem__(self, idx):
+        row    = self.data[idx]
+        image  = row["image"]
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(image)
+        image  = image.convert("RGB")
+        question = str(row["query"])
+        answer   = row["label"][0] if isinstance(row["label"], list) else str(row["label"])
+        conversation = [
+            {"role": "<|User|>",      "content": f"<image>\n{question}", "images": [image]},
+            {"role": "<|Assistant|>", "content": answer},
+        ]
+        return conversation, [image]
+
+
+def _find_asst_start(ids, asst_tok_ids):
+    """Return index just AFTER the <|Assistant|> token sequence."""
+    for j in range(len(ids) - len(asst_tok_ids) + 1):
+        if ids[j: j + len(asst_tok_ids)] == asst_tok_ids:
+            return j + len(asst_tok_ids)
+    return None
+
+
+_ASST_TOKEN_IDS = tokenizer.encode("<|Assistant|>", add_special_tokens=False)
+
+
+def collate_fn(batch):
+    conversations, images_list = zip(*batch)
+    all_images = [img for imgs in images_list for img in imgs]
+    inputs = processor(
+        conversations=list(conversations),
+        images=all_images,
+        force_batchify=True,
+        system_prompt=(
+            "You are a helpful assistant that answers questions "
+            "about charts and graphs accurately and concisely."
+        ),
+    )
+    input_ids = inputs["input_ids"]
+    labels    = input_ids.clone()
+    labels[input_ids == tokenizer.pad_token_id] = -100
+    # Mask user/system tokens — only compute loss on assistant reply
+    for i in range(labels.shape[0]):
+        asst_start = _find_asst_start(input_ids[i].tolist(), _ASST_TOKEN_IDS)
+        if asst_start is not None:
+            labels[i, :asst_start] = -100
+    inputs["labels"] = labels
+    return inputs
+
+
+train_ds = ChartQADataset(train_raw)
+val_ds   = ChartQADataset(val_raw)
+train_dl = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True,
+                      collate_fn=collate_fn, num_workers=2, pin_memory=True)
+val_dl   = DataLoader(val_ds,   batch_size=BATCH_SIZE, shuffle=False,
+                      collate_fn=collate_fn, num_workers=2, pin_memory=True)
+log.info(f"DataLoaders ready — {len(train_dl)} train steps/epoch")
+
+# ──────────────────────────────────────────────────────────────────────────────
+# ███  OPTIMIZER & SCHEDULER  █████████████████████████████████████████████████
+# ──────────────────────────────────────────────────────────────────────────────
+optimizer        = AdamW(
+    [p for p in model.parameters() if p.requires_grad],
+    lr=LR, weight_decay=0.01,
+)
+total_opt_steps  = math.ceil(len(train_dl) / GRAD_ACCUM) * NUM_EPOCHS
+warmup_steps     = min(50, total_opt_steps // 10)
+scheduler        = get_cosine_schedule_with_warmup(
+    optimizer,
+    num_warmup_steps=warmup_steps,
+    num_training_steps=total_opt_steps,
+)
+log.info(f"Optimiser ready — total opt_steps={total_opt_steps}  warmup={warmup_steps}")
+
+# ──────────────────────────────────────────────────────────────────────────────
+# ███  TRAINING LOOP  █████████████████████████████████████████████████████████
+# ──────────────────────────────────────────────────────────────────────────────
+Path(OUTPUT_DIR).mkdir(parents=True, exist_ok=True)
+global_step    = 0
+opt_step       = 0
+best_val_loss  = float("inf")
+running_loss   = 0.0
+
+log.info("=" * 60)
+log.info("Training start")
+log.info("=" * 60)
+
+for epoch in range(1, NUM_EPOCHS + 1):
+    model.train()
+    running_loss = 0.0
+
+    for batch in train_dl:
+        batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v
+                 for k, v in batch.items()}
+
+        with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
+            inputs_embeds = model.prepare_inputs_embeds(**batch)
+            out = model.language_model(
+                inputs_embeds=inputs_embeds,
+                attention_mask=batch["attention_mask"],
+                labels=batch["labels"],
+            )
+        loss = out.loss / GRAD_ACCUM
+        loss.backward()
+        running_loss += out.loss.item()
+        global_step  += 1
+
+        if global_step % GRAD_ACCUM == 0:
+            torch.nn.utils.clip_grad_norm_(
+                [p for p in model.parameters() if p.requires_grad], 1.0
+            )
+            optimizer.step()
+            scheduler.step()
+            optimizer.zero_grad()
+            opt_step += 1
+
+            if opt_step % LOG_EVERY == 0:
+                avg_loss = running_loss / (LOG_EVERY * GRAD_ACCUM)
+                lr_now   = scheduler.get_last_lr()[0]
+                vram_gb  = torch.cuda.memory_reserved() / 1e9 if torch.cuda.is_available() else 0
+                log.info(
+                    f"Epoch {epoch} | step {opt_step}/{total_opt_steps} | "
+                    f"loss={avg_loss:.4f} | lr={lr_now:.2e} | VRAM={vram_gb:.1f}GB"
+                )
+                tlog({
+                    "train/loss": avg_loss, "train/lr": lr_now,
+                    "train/vram_gb": vram_gb,
+                    "epoch": epoch, "opt_step": opt_step,
+                })
+                if avg_loss > 10.0:
+                    talert("loss_diverging",
+                           f"loss={avg_loss:.2f} at step {opt_step} — reduce lr by 10x",
+                           "ERROR")
+                elif avg_loss > 3.5 and opt_step > 100:
+                    talert("slow_convergence",
+                           f"loss={avg_loss:.2f} at step {opt_step} — check lr schedule",
+                           "WARN")
+                running_loss = 0.0
+
+            if opt_step % SAVE_STEPS == 0:
+                ckpt = f"{OUTPUT_DIR}/checkpoint-{opt_step}"
+                model.save_pretrained(ckpt)
+                log.info(f"Checkpoint -> {ckpt}")
+                talert("checkpoint_saved", f"step={opt_step} -> {ckpt}", "INFO")
+
+    # ── Validation ────────────────────────────────────────────────────────────
+    model.eval()
+    val_loss, val_steps = 0.0, 0
+    with torch.no_grad():
+        for batch in val_dl:
+            batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v
+                     for k, v in batch.items()}
+            with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
+                inputs_embeds = model.prepare_inputs_embeds(**batch)
+                out = model.language_model(
+                    inputs_embeds=inputs_embeds,
+                    attention_mask=batch["attention_mask"],
+                    labels=batch["labels"],
+                )
+            val_loss  += out.loss.item()
+            val_steps += 1
+
+    avg_val = val_loss / max(val_steps, 1)
+    log.info(f"---- Epoch {epoch} val_loss={avg_val:.4f} best={best_val_loss:.4f} ----")
+    tlog({"val/loss": avg_val, "epoch": epoch})
+
+    if avg_val < best_val_loss:
+        best_val_loss = avg_val
+        model.save_pretrained(f"{OUTPUT_DIR}/best")
+        log.info(f"New best -> val_loss={best_val_loss:.4f}")
+        talert("new_best", f"val_loss={best_val_loss:.4f} epoch={epoch}", "INFO")
+    elif epoch > 1 and avg_val > best_val_loss * 1.05:
+        talert("val_degrading",
+               f"val_loss={avg_val:.4f} > best*1.05={best_val_loss*1.05:.4f} — possible overfit",
+               "WARN")
+
+# ──────────────────────────────────────────────────────────────────────────────
+# ███  PUSH TO HUB  ███████████████████████████████████████████████████████████
+# ──────────────────────────────────────────────────────────────────────────────
+log.info(f"Pushing best checkpoint to {HUB_MODEL_ID} ...")
+best_path = Path(f"{OUTPUT_DIR}/best")
+
+# Model card
+card = f"""---
+license: other
+tags:
+  - deepseek-vl2
+  - chart-qa
+  - vision-language
+  - lora
+  - peft
+base_model: {MODEL_ID}
+datasets:
+  - HuggingFaceM4/ChartQA
+---
+
+# DeepSeek-VL2-tiny x ChartQA LoRA
+
+Fine-tuned [`{MODEL_ID}`]({MODEL_ID}) on
+[ChartQA](https://huggingface.co/datasets/HuggingFaceM4/ChartQA)
+with LoRA (r={LORA_R}, a={LORA_ALPHA}).
+
+| | |
+|--|--|
+| Base | `{MODEL_ID}` |
+| LoRA r / a | {LORA_R} / {LORA_ALPHA} |
+| Target modules | {', '.join(LORA_TARGETS)} |
+| LR | {LR} |
+| Epochs | {NUM_EPOCHS} |
+| Effective batch | {BATCH_SIZE * GRAD_ACCUM} |
+| Best val loss | {best_val_loss:.4f} |
+
+## Load adapter
+
+```python
+from deepseek_vl.models import DeepseekVLV2Processor, DeepseekVLV2ForCausalLM
+from peft import PeftModel
+from transformers import AutoModelForCausalLM
+import torch
+
+model_id   = "{MODEL_ID}"
+adapter_id = "{HUB_MODEL_ID}"
+
+base  = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True,
+                                             torch_dtype=torch.bfloat16)
+model = PeftModel.from_pretrained(base, adapter_id).eval().cuda()
+proc  = DeepseekVLV2Processor.from_pretrained(model_id)
+```
+"""
+(best_path / "README.md").write_text(card, encoding="utf-8")
+
+model.push_to_hub(HUB_MODEL_ID, commit_message="LoRA adapter — ChartQA fine-tune")
+processor.tokenizer.push_to_hub(HUB_MODEL_ID, commit_message="Add tokenizer")
+
+log.info(f"Done!  https://huggingface.co/{HUB_MODEL_ID}")
+talert("training_complete",
+       f"best_val_loss={best_val_loss:.4f}  model -> https://huggingface.co/{HUB_MODEL_ID}",
+       "INFO")
+
+if _USE_TRACKIO:
+    trackio.finish()