"""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()