diff --git a/astrai/trainer/strategy.py b/astrai/trainer/strategy.py index 2085967..55c5468 100644 --- a/astrai/trainer/strategy.py +++ b/astrai/trainer/strategy.py @@ -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 diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index d2df191..933ac4c 100644 --- a/astrai/trainer/train_context.py +++ b/astrai/trainer/train_context.py @@ -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 diff --git a/scripts/tools/train.py b/scripts/tools/train.py index a1a6559..b80c048 100644 --- a/scripts/tools/train.py +++ b/scripts/tools/train.py @@ -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 = { diff --git a/tests/trainer/test_grpo_strategy.py b/tests/trainer/test_grpo_strategy.py index 0c3e1a0..99e683a 100644 --- a/tests/trainer/test_grpo_strategy.py +++ b/tests/trainer/test_grpo_strategy.py @@ -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