cd ~/unifolm_ws/unifolm-vla

# Tìm file config
find . -name "*.yaml" -path "*/config/*" | head -10

# Sửa path checkpoint
# Thay:  pretrained_model_path: "/path/to/Unifolm-VLM-0"
# Thành: pretrained_model_path: "$HOME/models/Qwen2.5-VL-7B-Instruct"

# src/unifolm_vla/config/train_config.yaml (ví dụ)

model:
  # THAY ĐỔI: dùng Qwen2.5-VL-7B-Instruct thay Unifolm-VLM-0
  pretrained_model_path: "/home/user/models/Qwen2.5-VL-7B-Instruct"
  
dataset:
  rlds_data_dir: "/home/user/datasets/g1_pickplace_rlds"
  dataset_name: "g1_pickplace"
  robot_type: "g1_dex3"
  
training:
  num_epochs: 120          # tăng từ ~80 lên 120 vì bắt đầu từ checkpoint không chuyên
  batch_size: 4            # per GPU
  learning_rate: 1e-4
  warmup_steps: 200        # tăng warmup vì model cần adapt nhiều hơn
  save_every_n_epochs: 10
  
  # DeepSpeed settings (do NOT change)
  gradient_checkpointing: true
  bf16: true

conda activate unifolm
cd ~/unifolm_ws/unifolm-vla

accelerate launch \
  --config_file src/unifolm_vla/config/deepseeds/deepspeed_zero2.yaml \
  --num_processes 8 \
  src/unifolm_vla/training/train_unifolm_vla.py

[DeepSpeed] Using ZeRO Optimization Stage 2
Epoch 1/120 | Step 25 | Loss: 2.847 | LR: 5.0e-5
Epoch 1/120 | Step 50 | Loss: 2.341 | LR: 1.0e-4
...
Epoch 10/120 | Val Loss: 1.234 | Saved checkpoint: ./checkpoints/epoch_10/
...
Epoch 120/120 | Val Loss: 0.487 | Best checkpoint: epoch_110

conda activate unifolm
pip install peft bitsandbytes transformers accelerate

# Verify
python -c "import peft; print('peft version:', peft.__version__)"

"""
Custom single-GPU LoRA fine-tuning for unifolm-vla.
Trains Qwen2.5-VL-7B-Instruct with QLoRA on G1 robot data.
NOTE: This is a workaround — not official unifolm-vla API.
"""

import torch
import json
import numpy as np
from pathlib import Path
from transformers import (
    Qwen2VLForConditionalGeneration,
    AutoProcessor,
    BitsAndBytesConfig,
)
from peft import LoraConfig, get_peft_model, TaskType
import h5py
import glob
from torch.utils.data import Dataset, DataLoader

# ── 1. Model setup với QLoRA ──────────────────────────────────────────────────

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
)

model = Qwen2VLForConditionalGeneration.from_pretrained(
    "/home/user/models/Qwen2.5-VL-7B-Instruct",
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

processor = AutoProcessor.from_pretrained(
    "/home/user/models/Qwen2.5-VL-7B-Instruct"
)

# ── 2. LoRA config ────────────────────────────────────────────────────────────

lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=16,                    # LoRA rank — tăng lên 32 nếu có đủ VRAM
    lora_alpha=32,
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    lora_dropout=0.05,
    bias="none",
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# Kết quả mong đợi: trainable params: ~40M / 7B (0.6%) — rất nhỏ, fit single GPU

# ── 3. Dataset ────────────────────────────────────────────────────────────────

class G1RobotDataset(Dataset):
    """Load G1 robot demos từ HDF5 format."""
    
    def __init__(self, hdf5_dir: str, instruction: str = "pick up the red cup"):
        self.files = glob.glob(f"{hdf5_dir}/train/*.hdf5")
        self.instruction = instruction
    
    def __len__(self):
        return len(self.files)
    
    def __getitem__(self, idx):
        with h5py.File(self.files[idx], 'r') as f:
            # Load first frame as observation image
            frames = f['obs/left_wrist_rgb'][:]     # shape: (T, H, W, 3)
            actions = f['action'][:]                 # shape: (T, 28)
            instruction = f['language_instruction'][()].decode()
            
        # Dùng frame giữa episode làm example
        mid_frame = frames[len(frames) // 2]  # (H, W, 3), uint8
        mid_action = actions[len(actions) // 2]  # (28,)
        
        return {
            "image": mid_frame,
            "instruction": instruction or self.instruction,
            "action": torch.tensor(mid_action, dtype=torch.float32),
        }

dataset = G1RobotDataset(
    hdf5_dir="/home/user/datasets/g1_pickplace_hdf5",
    instruction="pick up the red cup"
)
dataloader = DataLoader(dataset, batch_size=1, shuffle=True)

# ── 4. Training loop ──────────────────────────────────────────────────────────

optimizer = torch.optim.AdamW(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=2e-4,
    weight_decay=0.01,
)

# Simple regression head để predict actions từ hidden states
action_head = torch.nn.Linear(3584, 28).to("cuda")  # 3584 = Qwen2.5-VL-7B hidden dim
head_optimizer = torch.optim.AdamW(action_head.parameters(), lr=2e-4)

num_epochs = 50  # cần nhiều hơn full fine-tune, nhưng nhanh hơn nhiều mỗi epoch

for epoch in range(num_epochs):
    total_loss = 0
    for batch in dataloader:
        # Prepare input cho VLM
        messages = [{
            "role": "user",
            "content": [
                {"type": "image", "image": batch["image"][0].numpy()},
                {"type": "text", "text": f"Robot task: {batch['instruction'][0]}. Predict action."}
            ]
        }]
        
        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = processor(
            text=[text],
            images=[batch["image"][0].numpy()],
            return_tensors="pt",
            padding=True
        ).to("cuda")
        
        # Forward pass — lấy hidden states
        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
            outputs = model(**inputs, output_hidden_states=True)
            hidden = outputs.hidden_states[-1][:, -1, :]  # last token hidden state
            
            # Predict action
            predicted_action = action_head(hidden.float())
            target_action = batch["action"].to("cuda")
            
            loss = torch.nn.functional.mse_loss(predicted_action, target_action)
        
        optimizer.zero_grad()
        head_optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        head_optimizer.step()
        
        total_loss += loss.item()
    
    if (epoch + 1) % 10 == 0:
        print(f"Epoch {epoch+1}/{num_epochs} | Loss: {total_loss/len(dataloader):.4f}")

# ── 5. Save checkpoint ────────────────────────────────────────────────────────

save_dir = Path("./checkpoints/lora_g1_pickplace")
model.save_pretrained(save_dir / "lora_weights")
torch.save(action_head.state_dict(), save_dir / "action_head.pt")
processor.save_pretrained(save_dir / "processor")
print(f"Saved to {save_dir}")

conda activate unifolm
python train_lora_single_gpu.py

# VRAM usage:
# Qwen2.5-VL-7B 4-bit quantized: ~6GB
# LoRA adapter: ~0.5GB
# Action head: negligible
# Activations + optimizer: ~8GB
# TỔNG: ~14-15GB → fit RTX 4090 (24GB) tốt

trainable params: 41,943,040 || all params: 7,615,832,064 || trainable%: 0.55
Epoch 10/50 | Loss: 0.8341
Epoch 20/50 | Loss: 0.4123
Epoch 30/50 | Loss: 0.2847
Epoch 40/50 | Loss: 0.2156
Epoch 50/50 | Loss: 0.1934
Saved to ./checkpoints/lora_g1_pickplace

# RunPod: 8× A100 SXM4 80GB, khoảng $8/giờ
# Training 100 demos ≈ 3 giờ = $24

# Sau khi thuê:
rsync -avz $HOME/datasets/g1_pickplace_rlds/ runpod:/workspace/datasets/
rsync -avz ~/unifolm_ws/unifolm-vla/ runpod:/workspace/unifolm-vla/

# Chạy training trên cloud
accelerate launch \
  --config_file src/unifolm_vla/config/deepseeds/deepspeed_zero2.yaml \
  --num_processes 8 \
  src/unifolm_vla/training/train_unifolm_vla.py

# Download checkpoint về local
rsync -avz runpod:/workspace/unifolm-vla/checkpoints/ $HOME/checkpoints/

tensorboard --logdir ./runs/

# Nếu không có TensorBoard, theo dõi qua terminal output
# Loss hội tụ tốt:
# Epoch 10: 2.1 → Epoch 50: 1.2 → Epoch 100: 0.7 → Epoch 120: 0.5