Write out label_mask instead of labels in olmocore tokenization script

open_instruct/grpo_fast.py

@@ -72,7 +72,7 @@
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
 from rich.pretty import pprint
 from torch.utils.tensorboard import SummaryWriter
-from transformers import AutoModelForCausalLM, PreTrainedModel, PreTrainedTokenizer, get_scheduler
+from transformers import AutoModelForCausalLM, GenerationConfig, PreTrainedModel, PreTrainedTokenizer, get_scheduler
 from transformers.integrations import HfDeepSpeedConfig
 from vllm import SamplingParams
 
@@ -983,8 +983,19 @@ def save_checkpoint_state(self, checkpoint_state_dir: str, client_state: Dict[st
                     checkpoint_state_dir, args.gs_checkpoint_state_dir
                 )
 
-    def save_model(self, output_dir: str) -> None:
+    def save_model(self, output_dir: str, chat_template_name: str, tokenizer: PreTrainedTokenizer) -> None:
         model_to_save = self.model
+        if "olmo" in chat_template_name:
+            # New chat template has no bos token, and two eos tokens: <|im_end|> and <|endoftext|>
+            model_to_save.generation_config = GenerationConfig(
+                temperature=None,
+                top_p=None,
+                eos_token_id=[
+                    tokenizer.convert_tokens_to_ids("<|im_end|>"),
+                    tokenizer.convert_tokens_to_ids("<|endoftext|>"),
+                ],
+            )
+
         if self.rank == 0:
             os.makedirs(output_dir, exist_ok=True)
 
@@ -1774,6 +1785,7 @@ def one_training_step(
     train_dataset,
     writer,
     wandb_url,
+    chat_template_name,
 ):
     """Train the model for one step."""
     update_ref_policy_future = []
@@ -1820,7 +1832,12 @@ def one_training_step(
                 checkpoint_dir = f"{args.output_dir}_checkpoints"
                 step_dir = os.path.join(checkpoint_dir, f"step_{training_step}")
                 logger.info(f"Saving model at step {training_step} to {step_dir}")
-                ray.get([policy_group.models[i].save_model.remote(step_dir) for i in range(args.world_size)])
+                ray.get(
+                    [
+                        policy_group.models[i].save_model.remote(step_dir, chat_template_name, tokenizer)
+                        for i in range(args.world_size)
+                    ]
+                )
                 if args.try_launch_beaker_eval_jobs_on_weka and is_beaker_job():
                     leaderboard_name = f"{args.hf_repo_revision}_step_{training_step}"
                     for i in range(args.world_size):
@@ -1917,11 +1934,23 @@ def maybe_evaluate(
         logger.warning("[Main Thread] 🙈 Evaluation responses not received")
 
 
-def save_final_model(args: Args, policy_group: ModelGroup, training_step: int, wandb_url: str):
+def save_final_model(
+    args: Args,
+    policy_group: ModelGroup,
+    tokenizer: PreTrainedTokenizer,
+    training_step: int,
+    wandb_url: str,
+    chat_template_name: str,
+):
     """Save the final model and launch evaluation jobs if configured."""
     logger.info(f"Saving final model at step {training_step} to {args.output_dir}")
     with Timer("[Main Thread] 🗡️ Saving model"):
-        ray.get([policy_group.models[i].save_model.remote(args.output_dir) for i in range(args.world_size)])
+        ray.get(
+            [
+                policy_group.models[i].save_model.remote(args.output_dir, chat_template_name, tokenizer)
+                for i in range(args.world_size)
+            ]
+        )
         if args.try_launch_beaker_eval_jobs_on_weka and is_beaker_job():
             leaderboard_name = args.hf_repo_revision
             for i in range(args.world_size):
@@ -2189,6 +2218,7 @@ def main(args: Args, tc: TokenizerConfig, model_config: ModelConfig, num_eval_sa
                 train_dataset,
                 writer,
                 wandb_url,
+                tc.chat_template_name,
             )
 
             maybe_evaluate(
@@ -2204,7 +2234,7 @@ def main(args: Args, tc: TokenizerConfig, model_config: ModelConfig, num_eval_sa
                 writer,
             )
 
-        save_final_model(args, policy_group, training_step, wandb_url)
+        save_final_model(args, policy_group, tokenizer, training_step, wandb_url, tc.chat_template_name)
 
     except Exception as e:
         logger.error(f"Training error occurred: {str(e)}\n{traceback.format_exc()}")

open_instruct/grpo_vllm_thread_ray_gtrl.py

@@ -79,6 +79,7 @@
 from transformers import (
     AutoModelForCausalLM,
     AutoModelForSequenceClassification,
+    GenerationConfig,
     PreTrainedModel,
     PreTrainedTokenizer,
     get_scheduler,
@@ -791,6 +792,7 @@ def train(
         train_dataset: Dataset,
         eval_dataset: Dataset,
         tokenizer: PreTrainedTokenizer,
+        tc: TokenizerConfig,
         vllm_engines: List[ray.actor.ActorHandle],
         metrics_queue: RayQueue,
         data_collator: Callable,
@@ -1378,7 +1380,7 @@ def generate_with_engines(prompts: List[List[int]], sampling_params: SamplingPar
                 checkpoint_dir = f"{args.output_dir}_checkpoints"
                 step_dir = os.path.join(checkpoint_dir, f"step_{training_step}")
                 print(f"Saving model at step {training_step} to {step_dir}")
-                self.save_model(self.model, step_dir)
+                self.save_model(self.model, tc.chat_template_name, tokenizer, step_dir)
                 if args.try_launch_beaker_eval_jobs_on_weka:
                     leaderboard_name = f"{args.hf_repo_revision}_step_{training_step}"
                     if self.rank == 0 and is_beaker_job():
@@ -1404,7 +1406,7 @@ def generate_with_engines(prompts: List[List[int]], sampling_params: SamplingPar
                                 print(f"Eval future {eval_futures[0]} is done")
                                 eval_futures.popleft()
         print(f"Saving final model at step {training_step} to {args.output_dir}")
-        self.save_model(self.model, args.output_dir)
+        self.save_model(self.model, tc.chat_template_name, tokenizer, args.output_dir)
         if args.try_launch_beaker_eval_jobs_on_weka:
             leaderboard_name = args.hf_repo_revision
             if self.rank == 0 and is_beaker_job():
@@ -1438,14 +1440,27 @@ def generate_with_engines(prompts: List[List[int]], sampling_params: SamplingPar
             shutil.copytree(args.output_dir, "/output", dirs_exist_ok=True)
         print("finished training")
 
-    def save_model(self, model_to_save: PreTrainedModel, output_dir: str) -> None:
+    def save_model(
+        self, model_to_save: PreTrainedModel, chat_template_name: str, tokenizer: PreTrainedTokenizer, output_dir: str
+    ) -> None:
         if self.rank == 0:
             os.makedirs(output_dir, exist_ok=True)
 
         # save model weights for ZeRO2/3
         if hasattr(model_to_save, "module"):
             model_to_save = model_to_save.module
 
+        if "olmo" in chat_template_name:
+            # New chat template has no bos token, and two eos tokens: <|im_end|> and <|endoftext|>
+            model_to_save.generation_config = GenerationConfig(
+                temperature=None,
+                top_p=None,
+                eos_token_id=[
+                    tokenizer.convert_tokens_to_ids("<|im_end|>"),
+                    tokenizer.convert_tokens_to_ids("<|endoftext|>"),
+                ],
+            )
+
         # gather parameters
         output_state_dict = {}
         for k, v in model_to_save.named_parameters():

open_instruct/model_utils.py

@@ -418,6 +418,7 @@ def save_with_accelerate(
     # otherwise, we get an error thrown at save time.
     if "olmo" in chat_template_name:
         # New chat template has no bos token, and two eos tokens: <|im_end|> and <|endoftext|>
+        logger.log(f"Detected olmo chat template: {chat_template_name}, updating model generation config.")
         model.generation_config = transformers.GenerationConfig(
             temperature=None,
             top_p=None,

open_instruct/ppo_fast.py

@@ -76,6 +76,7 @@
     AutoConfig,
     AutoModelForCausalLM,
     AutoModelForSequenceClassification,
+    GenerationConfig,
     PreTrainedModel,
     PreTrainedTokenizer,
     get_scheduler,
@@ -1074,7 +1075,7 @@ def train(
             self.offload_to_cpu(self.model)
         return metrics_list
 
-    def save_model(self, output_dir: str) -> None:
+    def save_model(self, output_dir: str, chat_template_name: str, tokenizer: PreTrainedTokenizer) -> None:
         model_to_save = self.model
         if self.rank == 0:
             os.makedirs(output_dir, exist_ok=True)
@@ -1083,6 +1084,17 @@ def save_model(self, output_dir: str) -> None:
         if hasattr(model_to_save, "module"):
             model_to_save = model_to_save.module
 
+        if "olmo" in chat_template_name:
+            # New chat template has no bos token, and two eos tokens: <|im_end|> and <|endoftext|>
+            model_to_save.generation_config = GenerationConfig(
+                temperature=None,
+                top_p=None,
+                eos_token_id=[
+                    tokenizer.convert_tokens_to_ids("<|im_end|>"),
+                    tokenizer.convert_tokens_to_ids("<|endoftext|>"),
+                ],
+            )
+
         # gather parameters
         output_state_dict = {}
         for k, v in model_to_save.named_parameters():
@@ -1819,7 +1831,12 @@ def main(args: Args, tc: TokenizerConfig, model_config: ModelConfig, reward_fn:
                         checkpoint_dir = f"{args.output_dir}_checkpoints"
                         step_dir = os.path.join(checkpoint_dir, f"step_{training_step}")
                         print(f"Saving model at step {training_step} to {step_dir}")
-                        ray.get([policy_group.models[i].save_model.remote(step_dir) for i in range(args.world_size)])
+                        ray.get(
+                            [
+                                policy_group.models[i].save_model.remote(step_dir, tc.chat_template_name, tokenizer)
+                                for i in range(args.world_size)
+                            ]
+                        )
                         if args.try_launch_beaker_eval_jobs_on_weka and is_beaker_job():
                             leaderboard_name = f"{args.hf_repo_revision}_step_{training_step}"
                             for i in range(args.world_size):
@@ -1889,7 +1906,12 @@ def main(args: Args, tc: TokenizerConfig, model_config: ModelConfig, reward_fn:
 
         print(f"Saving final model at step {training_step} to {args.output_dir}")
         with Timer("[Main Thread] 🗡️ Saving model"):
-            ray.get([policy_group.models[i].save_model.remote(args.output_dir) for i in range(args.world_size)])
+            ray.get(
+                [
+                    policy_group.models[i].save_model.remote(args.output_dir, tc.chat_template_name, tokenizer)
+                    for i in range(args.world_size)
+                ]
+            )
             if args.try_launch_beaker_eval_jobs_on_weka and is_beaker_job():
                 leaderboard_name = args.hf_repo_revision
                 for i in range(args.world_size):

open_instruct/ppo_vllm_thread_ray_gtrl.py

@@ -77,6 +77,7 @@
 from transformers import (
     AutoModelForCausalLM,
     AutoModelForSequenceClassification,
+    GenerationConfig,
     PreTrainedModel,
     PreTrainedTokenizer,
     get_scheduler,
@@ -1513,14 +1514,27 @@ def generate_with_engines(prompts: List[List[int]], sampling_params: SamplingPar
             shutil.copytree(args.output_dir, "/output", dirs_exist_ok=True)
         print("finished training")
 
-    def save_model(self, model_to_save: PreTrainedModel, output_dir: str) -> None:
+    def save_model(
+        self, model_to_save: PreTrainedModel, chat_template_name: str, tokenizer: PreTrainedTokenizer, output_dir: str
+    ) -> None:
         if self.rank == 0:
             os.makedirs(output_dir, exist_ok=True)
 
         # save model weights for ZeRO2/3
         if hasattr(model_to_save, "module"):
             model_to_save = model_to_save.module
 
+        if "olmo" in chat_template_name:
+            # New chat template has no bos token, and two eos tokens: <|im_end|> and <|endoftext|>
+            model_to_save.generation_config = GenerationConfig(
+                temperature=None,
+                top_p=None,
+                eos_token_id=[
+                    tokenizer.convert_tokens_to_ids("<|im_end|>"),
+                    tokenizer.convert_tokens_to_ids("<|endoftext|>"),
+                ],
+            )
+
         # gather parameters
         output_state_dict = {}
         for k, v in model_to_save.named_parameters():