ViperEkura
/
AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: separate old policy and ref model in GRPO strategy 

- Split single ref_model into old_model (importance sampling ratio) and ref_model (frozen KL regularizer)
- Move ref_model/old_model creation from strategy __init__ to TrainContextBuilder, pass as explicit parameters
- Remove periodic sync_ref_model + sync_interval; add sync_old_model for external rollout loop to call
- DPOStrategy also receives ref_model from builder
- Fix std to use unbiased=False (population std per GRPO paper)
- Remove redundant tests (test_grpo_kl_zero_at_init, test_grpo_no_sync_interval_param)
- Remove --grpo_sync_interval CLI arg
This commit is contained in:
ViperEkura 2026-07-14 20:00:54 +08:00
parent 3e0007fc91
commit 2c7a71a9c0
4 changed files with 80 additions and 58 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

53
astrai/trainer/strategy.py
View File

@ -223,14 +223,13 @@ class DPOStrategy(BaseStrategy):
self,
model: nn.Module,
device: str,
ref_model: nn.Module,
beta: float = 0.1,
reduction: str = "mean",
**kwargs,
):
super().__init__(model, device, **kwargs)
self.ref_model = create_ref_model(
self.model_fn, self.executor.unwrap_model(model)
).to(device=self.device)
self.ref_model = ref_model
self.beta = beta
self.reduction = reduction
@ -272,40 +271,40 @@ class GRPOStrategy(BaseStrategy):
broadcast across all response tokens. The loss is computed **only on
response tokens** — prompt tokens are masked out.
The strategy expects offline-collected batches (``responses`` / ``rewards``
pre-generated by the current or a recent policy). Call ``sync_ref_model()``
after each data-generation round so ``ref_model`` tracks the sampling policy.
Three model roles are distinguished:
* **Policy** ``self.model`` — the model being trained.
* **Old policy** ``self.old_model`` — the behaviour policy that generated
the responses. Used for the importance sampling ratio
``ρ = π_θ / π_old``. Synced externally after each data-generation round.
* **Reference model** ``self.ref_model`` — a frozen copy of the initial
policy (typically the SFT checkpoint) used **only** for the KL
regularisation term. It is never updated during training.
"""
def __init__(
self,
model: nn.Module,
device: str,
old_model: nn.Module,
ref_model: nn.Module,
clip_eps: float = 0.2,
kl_coef: float = 0.01,
group_size: int = 4,
sync_interval: int = 200,
**kwargs,
):
super().__init__(model, device, **kwargs)
self.ref_model = create_ref_model(
self.model_fn, self.executor.unwrap_model(model)
).to(device=self.device)
self.old_model = old_model
self.ref_model = ref_model
self.clip_eps = clip_eps
self.kl_coef = kl_coef
self.group_size = group_size
self.sync_interval = sync_interval
self._step = 0
def sync_ref_model(self):
"""Copy current model weights to ref model."""
self.ref_model.load_state_dict(self.executor.unwrap_model(self.model))
def sync_old_model(self):
"""Copy current policy weights to old model."""
self.old_model.load_state_dict(self.executor.unwrap_model(self.model))
def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
self._step += 1
if self._step % self.sync_interval == 0:
self.sync_ref_model()
batch = move_to_device(batch, self.device)
prompts = batch["prompts"]
responses = batch["responses"]
@ -332,25 +331,29 @@ class GRPOStrategy(BaseStrategy):
self.model, full_sequences, full_masks, "none"
)[:, prompt_len - 1 :]
with torch.no_grad():
token_log_probs_old = get_logprobs(
self.old_model, full_sequences, full_masks, "none"
)[:, prompt_len - 1 :]
token_log_probs_ref = get_logprobs(
self.ref_model, full_sequences, full_masks, "none"
)[:, prompt_len - 1 :]
# Reshape to [B, G, response_len]
token_log_probs_policy = token_log_probs_policy.view(batch_size, group_size, -1)
token_log_probs_old = token_log_probs_old.view(batch_size, group_size, -1)
token_log_probs_ref = token_log_probs_ref.view(batch_size, group_size, -1)
token_masks = masks_flat.view(batch_size, group_size, -1).float()
# Group-normalized advantages from scalar per-response rewards.
eps = 1e-8
mean = rewards.mean(dim=-1, keepdim=True)
std = rewards.std(dim=-1, keepdim=True)
std = rewards.std(dim=-1, keepdim=True, unbiased=False)
advantages = (rewards - mean) / (std + eps)
# Broadcast scalar advantage to every response token: [B, G, 1]
advantages = advantages.unsqueeze(-1)
# Token-level ratio and PPO clipping.
log_ratio = token_log_probs_policy - token_log_probs_ref
# Token-level ratio (π_θ / π_old) and PPO clipping.
log_ratio = token_log_probs_policy - token_log_probs_old
ratio = torch.exp(log_ratio)
surr1 = ratio * advantages
@ -359,8 +362,10 @@ class GRPOStrategy(BaseStrategy):
token_count = token_masks.sum().clamp(min=1.0)
policy_loss = (per_token_policy_loss * token_masks).sum() / token_count
# KL penalty with k1 estimator (non-negative): r - log(r) - 1, r=π_ref/π_θ.
r = torch.exp(-log_ratio)
# KL penalty to frozen reference model with k1 estimator (non-negative):
# k1 = π_ref / π_θ - log(π_ref / π_θ) - 1, where π_ref / π_θ = exp(log_ref - log_policy).
log_ref_ratio = token_log_probs_ref - token_log_probs_policy
r = torch.exp(log_ref_ratio)
kl_per_token = r - torch.log(r + eps) - 1.0
kl_penalty = self.kl_coef * (kl_per_token * token_masks).sum() / token_count

18
astrai/trainer/train_context.py
View File

@ -13,7 +13,7 @@ from astrai.parallel.executor import BaseExecutor, ExecutorFactory
from astrai.parallel.setup import get_current_device, get_rank, get_world_size
from astrai.protocols import OptimizerProtocol, SchedulerProtocol
from astrai.serialization import Checkpoint, load_json
from astrai.trainer.strategy import BaseStrategy, StrategyFactory
from astrai.trainer.strategy import BaseStrategy, StrategyFactory, create_ref_model
@dataclass
@ -177,13 +177,27 @@ class TrainContextBuilder:
if obj is not None:
obj.load_state_dict(extra[name])
strategy_kwargs = dict(cfg.extra_kwargs)
if cfg.strategy in ("dpo", "grpo"):
ref_model = create_ref_model(
cfg.model_fn, executor.unwrap_model(context.model)
).to(device=device)
strategy_kwargs["ref_model"] = ref_model
if cfg.strategy == "grpo":
old_model = create_ref_model(
cfg.model_fn, executor.unwrap_model(context.model)
).to(device=device)
strategy_kwargs["old_model"] = old_model
context.strategy = StrategyFactory.create(
cfg.strategy,
model=context.model,
device=device,
executor=executor,
model_fn=cfg.model_fn,
**cfg.extra_kwargs,
**strategy_kwargs,
)
return context

7
scripts/tools/train.py
View File

@ -252,12 +252,6 @@ def parse_args() -> argparse.Namespace:
default="checkpoint/logs",
help="Directory for metric logs.",
)
parser.add_argument(
"--grpo_sync_interval",
type=int,
default=200,
help="GRPO ref model sync interval (steps).",
)
parser.add_argument(
"--start_epoch", type=int, default=0, help="Start epoch for training."
)
@ -444,7 +438,6 @@ def train(
"clip_eps": kwargs.pop("grpo_clip_eps"),
"kl_coef": kwargs.pop("grpo_kl_coef"),
"group_size": kwargs.pop("group_size"),
"sync_interval": kwargs.pop("grpo_sync_interval"),
}
executor_kwargs = {

60
tests/trainer/test_grpo_strategy.py
View File

@ -56,19 +56,32 @@ def _make_batch(
}
def _make_frozen_copy(model, device):
"""Create a frozen copy of ``model`` with independent weights loaded."""
config = _make_config()
copy = AutoRegressiveLM(config).to(device=device)
copy.load_state_dict(model.state_dict())
copy.requires_grad_(False)
copy.eval()
return copy
@pytest.fixture
def grpo_strategy():
"""Build a GRPOStrategy with a small real model and fake executor."""
device = "cuda" if torch.cuda.is_available() else "cpu"
model, config = _make_model(device)
old_model = _make_frozen_copy(model, device)
ref_model = _make_frozen_copy(model, device)
strategy = GRPOStrategy(
model=model,
device=device,
old_model=old_model,
ref_model=ref_model,
clip_eps=0.2,
kl_coef=0.01,
group_size=4,
sync_interval=200,
model_fn=lambda c=config: AutoRegressiveLM(c).to(device=device),
executor=_FakeExecutor(),
)
@ -108,6 +121,16 @@ def test_grpo_ref_model_not_updated(grpo_strategy):
assert p.grad is None
def test_grpo_old_model_not_updated(grpo_strategy):
"""Backward should not populate gradients on old_model."""
strategy, device = grpo_strategy
batch = _make_batch(device=device)
loss = strategy.compute_loss(batch)
loss.backward()
for p in strategy.old_model.parameters():
assert p.grad is None
def test_grpo_prompt_tokens_masked(grpo_strategy):
"""When only prompt-equivalent tokens are unmasked (response mask all 0),
the policy loss should be zero (no valid tokens contribute)."""
@ -127,45 +150,32 @@ def test_grpo_identical_rewards_zero_advantage(grpo_strategy):
batch = _make_batch(device=device)
batch["rewards"] = torch.ones(batch["rewards"].shape, device=device)
loss = strategy.compute_loss(batch)
# At init policy == ref, so ratio == 1, KL == 0; advantage == 0.
# At init policy == old == ref, so ratio == 1, KL == 0; advantage == 0.
assert loss.item() == pytest.approx(0.0, abs=1e-5)
def test_grpo_kl_zero_at_init(grpo_strategy):
"""At initialization policy == ref_model, so KL penalty must be 0."""
def test_grpo_sync_old_model(grpo_strategy):
"""sync_old_model copies current policy weights into old_model."""
strategy, device = grpo_strategy
batch = _make_batch(device=device)
# Make rewards distinct so advantages are non-zero (isolates KL term).
loss = strategy.compute_loss(batch)
# KL term is 0 at init; loss is purely policy surrogate.
# With ratio==1, surr1==surr2==advantage, so policy_loss = -mean(|adv|).
# Just assert KL portion is negligible by checking loss is finite and
# re-running after a model update increases loss magnitude.
assert torch.isfinite(loss).item()
def test_grpo_sync_ref_model(grpo_strategy):
"""sync_ref_model copies current policy weights into ref_model."""
strategy, device = grpo_strategy
# Perturb policy model so it differs from ref.
# Perturb policy model so it differs from old.
with torch.no_grad():
for p in strategy.model.parameters():
p.add_(0.05)
# ref_model should still hold original weights (differ from policy).
# old_model should still hold original weights (differ from policy).
policy_sd = strategy.model.state_dict()
ref_sd = strategy.ref_model.state_dict()
old_sd = strategy.old_model.state_dict()
differs_before = any(
not torch.allclose(policy_sd[k], ref_sd[k]) for k in policy_sd if k in ref_sd
not torch.allclose(policy_sd[k], old_sd[k]) for k in policy_sd if k in old_sd
)
assert differs_before
strategy.sync_ref_model()
strategy.sync_old_model()
ref_sd_after = strategy.ref_model.state_dict()
old_sd_after = strategy.old_model.state_dict()
matches = all(
torch.allclose(policy_sd[k], ref_sd_after[k])
torch.allclose(policy_sd[k], old_sd_after[k])
for k in policy_sd
if k in ref_sd_after
if k in old_sd_after
)
assert matches