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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user