mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-09 02:51:56 +00:00
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:
+2
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user