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:
2026-07-14 20:03:45 +08:00
parent 3e0007fc91
commit 2c7a71a9c0
4 changed files with 80 additions and 58 deletions
+16 -2
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