File size: 4,057 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 | # Trainer worker implementation (Ray actor or standalone worker)
from typing import Any, List, Optional
import torch
try:
import ray
ray_remote = ray.remote(num_gpus=1)
except ImportError:
ray = None
def ray_remote(cls):
return cls
from ..buffer.replay_buffer import BoundedReplayBuffer, Experience
from .ppo_loss import compute_ppo_loss, compute_m2po_loss
from .grpo_loss import compute_grpo_loss
@ray_remote
class TrainerWorker:
"""Ray actor or standalone worker that samples experiences and runs PPO/M2PO/GRPO updates.
Arguments
---------
model_path: str or torch.nn.Module, optional
Path to model or model module.
buffer: Any
Shared replay buffer or actor handle.
use_grpo: bool, default False
Whether to train with GRPO loss.
use_m2po: bool, default False
Whether to train with M2PO loss.
device: str, optional
Device to use ('cuda' or 'cpu').
"""
def __init__(
self,
model_path: Optional[Any] = None,
buffer: Optional[Any] = None,
use_grpo: bool = False,
use_m2po: bool = False,
device: Optional[str] = None,
lr: float = 1e-4,
):
self.buffer = buffer
self.use_grpo = use_grpo
self.use_m2po = use_m2po
self.device = torch.device(
device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu")
)
if isinstance(model_path, torch.nn.Module):
self.model = model_path.to(self.device)
else:
self.model = torch.nn.Sequential(
torch.nn.Linear(128, 128),
torch.nn.ReLU(),
torch.nn.Linear(128, 128),
).to(self.device)
self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr)
self.version = 0
def step(self, batch_size: int = 64) -> Optional[float]:
"""Perform a single training step."""
if self.buffer is None:
return None
# Sample from buffer (actor handle or local instance)
if hasattr(self.buffer, "sample"):
if hasattr(self.buffer.sample, "remote") and ray is not None:
exps: List[Experience] = ray.get(self.buffer.sample.remote(batch_size))
else:
exps: List[Experience] = self.buffer.sample(batch_size)
else:
return None
if not exps:
return None
# Pad or stack tensors
max_prompt = max(e.prompt_ids.shape[-1] for e in exps)
max_gen = max(e.generated_ids.shape[-1] for e in exps)
padded_old_logprobs = []
padded_advantages = []
for e in exps:
lp = e.log_probs.to(self.device)
if lp.shape[-1] < max_gen:
lp = torch.cat([lp, torch.zeros(max_gen - lp.shape[-1], device=self.device)])
padded_old_logprobs.append(lp)
r = float(e.reward)
padded_advantages.append(torch.full((max_gen,), r, device=self.device))
old_log_probs = torch.stack(padded_old_logprobs)
advantages = torch.stack(padded_advantages)
# Compute simulated policy log probs with small perturbation for optimization
dummy_in = torch.randn(len(exps), 128, device=self.device)
rep = self.model(dummy_in)
policy_log_probs = old_log_probs + 0.01 * rep.mean(dim=-1, keepdim=True)
if self.use_grpo:
loss = compute_grpo_loss(
policy_log_probs.unsqueeze(0),
old_log_probs.unsqueeze(0),
advantages.unsqueeze(0),
)
elif self.use_m2po:
loss = compute_m2po_loss(policy_log_probs, old_log_probs, advantages)
else:
loss = compute_ppo_loss(policy_log_probs, old_log_probs, advantages)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
self.version += 1
return float(loss.item())
def get_version(self) -> int:
return self.version
|