fix: on-policy distillation missing temperature parameter fix (#103)

Co-authored-by: John Schulman <[email protected]>

Author: Andrewzh112

tinker_cookbook/distillation/train_on_policy.py

@@ -132,6 +132,7 @@ class Config:
     dataset_configs: List[DistillationDatasetConfig]
     model_name: str
     max_tokens: int
+    temperature: float = 1.0
     compute_post_kl: bool = False
     evaluator_builders: list[SamplingClientEvaluatorBuilder] = chz.field(default_factory=list)
     lora_rank: int = 32
@@ -308,6 +309,7 @@ async def do_sync_training(
                         do_group_rollout_and_filter_constant_reward(
                             sampling_client,
                             builder,
+                            temperature=cfg.temperature,
                             max_tokens=cfg.max_tokens,
                             do_remove_constant_reward_groups=False,
                         ),

tinker_cookbook/recipes/distillation/on_policy_distillation.py

@@ -64,6 +64,7 @@ class CLIConfig:
     groups_per_batch: int = 1024
     learning_rate: float = 1e-4
     max_tokens: int = 4096
+    temperature: float = 1.0
     kl_penalty_coef: float = 1.0
     kl_discount_factor: float = 0.0
 

tinker_cookbook/recipes/distillation/on_policy_multi_teacher.py

@@ -58,6 +58,7 @@ class CLIConfig:
     group_size: int = 4  # Number of rollouts per prompt
     learning_rate: float = 1e-4
     max_tokens: int = 4096
+    temperature: float = 1.0
     kl_penalty_coef: float = 1.0
     kl_discount_factor: float = 0.0