mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 02:06:15 +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(
|
return CosineDecayWithWarmupSchedulerConfig(
|
||||||
num_warmup_steps=500,
|
num_warmup_steps=500,
|
||||||
num_decay_steps=50000,
|
num_decay_steps=50000,
|
||||||
|
peak_lr=3e-4,
|
||||||
|
decay_lr=1e-6,
|
||||||
)
|
)
|
||||||
|
|
||||||
def validate_features(self) -> None:
|
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.
|
# 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:
|
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)
|
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:
|
else:
|
||||||
unwrapped_model.push_model_to_hub(cfg, state_dict=model_state_dict, dataset_meta=dataset.meta)
|
unwrapped_model.push_model_to_hub(cfg, state_dict=model_state_dict, dataset_meta=dataset.meta)
|
||||||
preprocessor.push_to_hub(active_cfg.repo_id)
|
preprocessor.push_to_hub(active_cfg.repo_id)
|
||||||
|
|||||||
Reference in New Issue
Block a user