fix(train): fix VF scheduler config and reward model hub push

Add missing peak_lr and decay_lr to DistributionalVFConfig scheduler
preset. Fix push_model_to_hub call for reward models.
This commit is contained in:
Khalil Meftah
2026-07-08 12:08:17 +02:00
parent 9a846c4fca
commit 582e953676
2 changed files with 4 additions and 0 deletions
@@ -98,6 +98,8 @@ class DistributionalVFConfig(RewardModelConfig):
return CosineDecayWithWarmupSchedulerConfig(
num_warmup_steps=500,
num_decay_steps=50000,
peak_lr=3e-4,
decay_lr=1e-6,
)
def validate_features(self) -> None:
+2
View File
@@ -741,6 +741,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
# PEFT only applies when training a policy — reward models use the plain path.
if not cfg.is_reward_model_training and cfg.policy.use_peft:
unwrapped_model.push_model_to_hub(cfg, peft_model=unwrapped_model, dataset_meta=dataset.meta)
elif cfg.is_reward_model_training:
unwrapped_model.push_model_to_hub(cfg)
else:
unwrapped_model.push_model_to_hub(cfg, state_dict=model_state_dict, dataset_meta=dataset.meta)
preprocessor.push_to_hub(active_cfg.repo_id)