File size: 6,408 Bytes
e791b16 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | """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()
|