scripts/verify_gpu.py

"""Verification and benchmark script for AsyncTensorRLHF on CUDA GPU.

Validates:
1. GPU hardware detection and VRAM status
2. Tensor-native reward computation directly on CUDA
3. PPO, M2PO, and GRPO loss calculation on CUDA
4. Group-aware buffer advantage normalization
5. Versioned buffer staleness handling
6. End-to-end closed-loop async rollout and policy optimization on CUDA
"""

import asyncio
import os
import pathlib
import sys
import time

# Ensure project root is in sys.path
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent))

import torch

from src.reward.tensor_native import tensor_native_reward, gpu_reward_simple
from src.trainer.ppo_loss import compute_ppo_loss, compute_m2po_loss
from src.trainer.grpo_loss import compute_grpo_loss
from src.buffer.replay_buffer import BoundedReplayBuffer, VersionedReplayBuffer, VersionedExperience
from src.buffer.group_buffer import GroupAwareReplayBuffer
from src.rollout.vllm_engine import HFEngine
from src.rollout.async_engine import AsyncEngine
from src.rollout.version_manager import VersionManager
from src.trainer.trainer_worker import TrainerWorker


def main():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print("=" * 60)
    print("AsyncTensorRLHF GPU Verification & Benchmark")
    print("=" * 60)
    print(f"Target Device: {device.upper()}")
    if device == "cuda":
        print(f"GPU Model: {torch.cuda.get_device_name(0)}")
        print(f"CUDA Capability: {torch.cuda.get_device_capability(0)}")
        allocated_mb = torch.cuda.memory_allocated(0) / (1024 * 1024)
        print(f"Initial Allocated VRAM: {allocated_mb:.2f} MB")
    print("-" * 60)

    # 1. Tensor-Native Reward on GPU
    print("[1/5] Testing Tensor-Native Reward on CUDA...")
    batch_size = 64
    seq_len = 256
    gen_ids = torch.randint(0, 1000, (batch_size, seq_len), device=device)
    answer_patterns = [torch.randint(0, 1000, (4,), device=device) for _ in range(batch_size)]

    t0 = time.perf_counter()
    rewards = tensor_native_reward(gen_ids, answer_patterns, eos_token_id=99, device=device)
    t1 = time.perf_counter()
    print(f"  -> Batched rewards computed for B={batch_size}, L={seq_len} in {(t1-t0)*1000:.2f} ms")
    assert rewards.device.type == device
    assert rewards.shape == (batch_size,)
    print("  -> Tensor-native reward: PASSED")

    # 2. PPO and M2PO Losses on GPU
    print("\n[2/5] Testing Policy Losses (PPO & M2PO) on CUDA...")
    B, L = 16, 64
    policy_lp = torch.randn(B, L, device=device, requires_grad=True)
    old_lp = torch.randn(B, L, device=device)
    advantages = torch.randn(B, L, device=device)

    loss_ppo = compute_ppo_loss(policy_lp, old_lp, advantages)
    loss_ppo.backward()
    assert torch.isfinite(loss_ppo)
    print(f"  -> PPO Loss: {loss_ppo.item():.4f} (Finite: True)")

    policy_lp2 = torch.randn(B, L, device=device, requires_grad=True)
    loss_m2po = compute_m2po_loss(policy_lp2, old_lp, advantages)
    loss_m2po.backward()
    assert torch.isfinite(loss_m2po)
    print(f"  -> M2PO Loss: {loss_m2po.item():.4f} (Finite: True)")
    print("  -> Policy losses: PASSED")

    # 3. Group-Aware Replay Buffer for GRPO
    print("\n[3/5] Testing Group-Aware Buffer for GRPO...")
    group_buf = GroupAwareReplayBuffer(group_size=4, max_groups=8)
    for i in range(4):
        group_buf.add_response(
            prompt_id=42,
            response=torch.tensor([i, i+1], device=device),
            log_prob=torch.tensor([-0.1, -0.2], device=device),
            reward=float(i),
            version=0,
        )
    assert group_buf.ready.qsize() == 1
    completed_group = group_buf.ready.get_nowait()
    assert completed_group.is_complete is True
    assert completed_group.advantages.shape == (4,)
    print(f"  -> GRPO Group Advantages standardized: {completed_group.advantages.tolist()}")
    print("  -> Group-aware buffer: PASSED")

    # 4. Versioned Replay Buffer with Staleness Eviction
    print("\n[4/5] Testing Versioned Replay Buffer & Staleness...")
    vbuf = VersionedReplayBuffer(max_size=100, max_staleness=3)
    vbuf.current_version = 10
    # Stale experience (version 5 is older than 10 - 3 = 7)
    vbuf.push(VersionedExperience(
        prompt_ids=torch.tensor([1], device=device),
        generated_ids=torch.tensor([1, 2], device=device),
        log_probs=torch.tensor([-0.1], device=device),
        reward=1.0,
        policy_version=5,
        generation_step=1,
    ))
    assert len(vbuf.buffer) == 0, "Stale experience was not evicted"
    # Fresh experience (version 8 >= 7)
    vbuf.push(VersionedExperience(
        prompt_ids=torch.tensor([1], device=device),
        generated_ids=torch.tensor([1, 2], device=device),
        log_probs=torch.tensor([-0.1], device=device),
        reward=1.0,
        policy_version=8,
        generation_step=1,
    ))
    assert len(vbuf.buffer) == 1, "Fresh experience was not kept"
    print("  -> Versioned buffer staleness eviction: PASSED")

    # 5. Closed-Loop Async Rollout + GPU Training
    print("\n[5/5] Testing End-to-End Closed-Loop Rollout & Training on GPU...")
    buf = BoundedReplayBuffer(max_size=50)
    vm = VersionManager()
    hf_engine = HFEngine(model_or_path=None, device=device, max_new_tokens=8, vocab_size=64)
    async_engine = AsyncEngine(buffer=buf, version_manager=vm, engine=hf_engine)
    trainer = TrainerWorker(buffer=buf, device=device)

    async def run_training():
        for step in range(3):
            prompts = [
                {"input_ids": torch.tensor([1, 2], device=device), "gt_ids": torch.tensor([2], device=device), "eos_token_id": 99},
                {"input_ids": torch.tensor([3, 4], device=device), "gt_ids": torch.tensor([4], device=device), "eos_token_id": 99},
            ]
            await async_engine.rollout(prompts)
            loss = trainer.step(batch_size=2)
            vm.bump()
            print(f"  -> Step {step+1}: Buffer Size={buf.size()}, Policy Version={vm.current()}, Loss={loss:.4f}")

    asyncio.run(run_training())
    print("  -> Closed-loop rollout + training: PASSED")

    print("-" * 60)
    if device == "cuda":
        final_mb = torch.cuda.memory_allocated(0) / (1024 * 1024)
        print(f"Final Allocated VRAM: {final_mb:.2f} MB")
    print("ALL VERIFICATIONS AND BENCHMARKS COMPLETED SUCCESSFULLY (EXIT 0)")
    print("=" * 60)


if __name__ == "__main__":
    main()