Merge pull request #3037 from AI-Hypercomputer:anisha-dapo2

PiperOrigin-RevId: 862813823

Author: Google-ML-Automation

src/MaxText/configs/rl.yml

@@ -92,7 +92,7 @@ enable_tunix_perf_metrics: False
 batch_size: 1
 # Increase `batch_size` and `MAX_STEPS` for better results.
 # num_batches: 3738
-num_batches: 4  # 200
+num_batches: 4
 # A batch can be split into multiple micro batches for memory management
 # and/or async sampling and training.
 micro_batch_size: -1
@@ -171,7 +171,8 @@ skip_jax_distributed_system: True
 
 # # TODO(@mazumdera): fix this
 # Dataset Configuration
-dataset_name: 'gsm8k'
+dataset_name: 'gsm8k' # huggingface:open-r1/DAPO-Math-17k-Processed
+eval_dataset_name: 'gsm8k' # huggingface:BytedTsinghua-SIA/AIME-2024
 train_split: 'train'
 eval_split: 'test'
 tokenizer_type: 'huggingface'

src/MaxText/rl/train_rl.py

@@ -104,50 +104,47 @@ def get_maxtext_model(config, devices=None):
   return tunix_model, mesh
 
 
-def get_dataset(model_tokenizer, tmvp_config, data_dir, split="train") -> grain.MapDataset:
+def get_dataset(
+    model_tokenizer, tmvp_config, data_dir, split="train", data_files=None, dataset_name=None
+) -> grain.MapDataset:
   """Download data"""
   if not os.path.exists(data_dir):
     os.makedirs(data_dir)
 
-  data = tfds.data_source(
-      tmvp_config.dataset_name,
-      split=split,
-      data_dir=data_dir,
-      builder_kwargs={"file_format": tfds.core.FileFormat.ARRAY_RECORD},
-      download=True,
-  )
+  if dataset_name is None:
+    raise ValueError("dataset_name must be provided")
+
+  if dataset_name.startswith("huggingface:"):
+    import datasets  # pylint: disable=import-outside-toplevel
+
+    if data_files is None:
+      hf_dataset_name = dataset_name.replace("huggingface:", "")
+      data = datasets.load_dataset(hf_dataset_name, split=split, cache_dir=data_dir)
+      if tmvp_config.debug.rl:
+        max_logging.log(f"Loaded Hugging Face dataset {hf_dataset_name} with split {split}. Size: {len(data)}")
+    else:  # data_files have been provided, useful for using slices of large datasets like nvidia/OpenMathInstruct-2
+      data = datasets.load_dataset(
+          "parquet",
+          data_files={tmvp_config.train_split: data_files},
+          split=split,
+          cache_dir=data_dir,
+      )
+  else:
+    builder_kwargs = {"file_format": tfds.core.FileFormat.ARRAY_RECORD}
+    data = tfds.data_source(
+        dataset_name,
+        split=split,
+        data_dir=data_dir,
+        builder_kwargs=builder_kwargs,
+        download=True,
+    )
 
   template_config = load_template_from_file(tmvp_config.chat_template_path)
+
   loaded_dataset = (
       grain.MapDataset.source(data)
       .shuffle(seed=tmvp_config.data_shuffle_seed)
-      .map(
-          lambda x: {
-              # passed to model forward pass
-              "prompts": model_tokenizer.apply_chat_template(
-                  [
-                      {
-                          "role": "user",
-                          "content": template_config["TEMPLATE"].format(
-                              system_prompt=template_config["SYSTEM_PROMPT"].format(
-                                  reasoning_start_token=tmvp_config.reasoning_start_token,
-                                  reasoning_end_token=tmvp_config.reasoning_end_token,
-                                  solution_start_token=tmvp_config.solution_start_token,
-                                  solution_end_token=tmvp_config.solution_end_token,
-                              ),
-                              question=x["question"].decode("utf-8"),
-                          ),
-                      },
-                  ],
-                  tokenize=False,
-                  add_generation_prompt=True,
-              ),
-              # passed to reward functions
-              "question": x["question"].decode("utf-8"),
-              # passed to reward functions
-              "answer": utils_rl.extract_hash_answer(x["answer"].decode("utf-8")),
-          }
-      )
+      .map(lambda x: utils_rl.process_data(dataset_name, model_tokenizer, template_config, tmvp_config, x))
   )
   return loaded_dataset
 
@@ -290,19 +287,33 @@ def rl_train(trainer_config, sampler_config, trainer_devices, sampler_devices):
   model_tokenizer = AutoTokenizer.from_pretrained(trainer_config.tokenizer_path)
 
   # Load datasets
-  dataset = get_dataset(model_tokenizer, trainer_config, train_data_dir, trainer_config.train_split).batch(
-      trainer_config.batch_size
-  )[: trainer_config.num_batches]
+  dataset = get_dataset(
+      model_tokenizer,
+      trainer_config,
+      train_data_dir,
+      trainer_config.train_split,
+      data_files=trainer_config.hf_train_files,
+      dataset_name=trainer_config.dataset_name,
+  ).batch(trainer_config.batch_size)[: trainer_config.num_batches]
 
   if trainer_config.train_fraction == 1.0:
     train_dataset = dataset.repeat(trainer_config.num_epoch)
   else:
     train_dataset = dataset[: int(len(dataset) * trainer_config.train_fraction)]
     train_dataset = train_dataset.repeat(trainer_config.num_epoch)
 
-  test_dataset = get_dataset(model_tokenizer, trainer_config, test_data_dir, trainer_config.eval_split).batch(
-      trainer_config.batch_size
-  )[: trainer_config.num_test_batches]
+  eval_dataset_name = getattr(trainer_config, "eval_dataset_name", None)
+  if not eval_dataset_name:
+    eval_dataset_name = trainer_config.dataset_name
+
+  test_dataset = get_dataset(
+      model_tokenizer,
+      trainer_config,
+      test_data_dir,
+      trainer_config.eval_split,
+      data_files=trainer_config.hf_eval_files,
+      dataset_name=eval_dataset_name,
+  ).batch(trainer_config.batch_size)[: trainer_config.num_test_batches]
 
   # Let's see how one batch of the dataset looks like!
   if trainer_config.debug.rl:

src/MaxText/rl/utils_rl.py

@@ -19,6 +19,65 @@
 from MaxText import max_logging
 
 
+# Constants for normalization
+SUBSTITUTIONS = [
+    ("an ", ""),
+    ("a ", ""),
+    (".$", "$"),
+    ("\\$", ""),
+    (r"\ ", ""),
+    (" ", ""),
+    ("mbox", "text"),
+    (",\\text{and}", ","),
+    ("\\text{and}", ","),
+    ("\\text{m}", "\\text{}"),
+]
+
+REMOVED_EXPRESSIONS = [
+    "square",
+    "ways",
+    "integers",
+    "dollars",
+    "mph",
+    "inches",
+    "hours",
+    "km",
+    "units",
+    "\\ldots",
+    "sue",
+    "points",
+    "feet",
+    "minutes",
+    "digits",
+    "cents",
+    "degrees",
+    "cm",
+    "gm",
+    "pounds",
+    "meters",
+    "meals",
+    "edges",
+    "students",
+    "childrentickets",
+    "multiples",
+    "\\text{s}",
+    "\\text{.}",
+    "\\text{\ns}",
+    "\\text{}^2",
+    "\\text{}^3",
+    "\\text{\n}",
+    "\\text{}",
+    r"\mathrm{th}",
+    r"^\circ",
+    r"^{\circ}",
+    r"\;",
+    r",\!",
+    "{,}",
+    '"',
+    "\\dots",
+]
+
+
 # Let's define a RegEx for checking whether the format matches.
 #
 def get_match_format_regex(tmvp_config):
@@ -90,6 +149,47 @@ def match_format_approximately(prompts, completions, tmvp_config, **kargs):
   return scores
 
 
+def normalize_final_answer(final_answer: str) -> str:
+  """Normalize a final answer to a quantitative reasoning question.
+
+  Args:
+      final_answer: The answer string to normalize
+
+  Returns:
+      Normalized answer string
+  """
+  final_answer = final_answer.split("=")[-1]
+
+  # Apply substitutions and removals
+  for before, after in SUBSTITUTIONS:
+    final_answer = final_answer.replace(before, after)
+  for expr in REMOVED_EXPRESSIONS:
+    final_answer = final_answer.replace(expr, "")
+
+  # Extract and normalize LaTeX math
+  final_answer = re.sub(r"(.*?)(\$)(.*?)(\$)(.*)", "$\\3$", final_answer)
+  final_answer = re.sub(r"(\\text\{)(.*?)(\})", "\\2", final_answer)
+  final_answer = re.sub(r"(\\textbf\{)(.*?)(\})", "\\2", final_answer)
+  final_answer = re.sub(r"(\\overline\{)(.*?)(\})", "\\2", final_answer)
+  final_answer = re.sub(r"(\\boxed\{)(.*)(\})", "\\2", final_answer)
+
+  # Normalize shorthand TeX:
+  #  \fracab -> \frac{a}{b}
+  #  \frac{abc}{bef} -> \frac{abc}{bef}
+  #  \fracabc -> \frac{a}{b}c
+  #  \sqrta -> \sqrt{a}
+  #  \sqrtab -> sqrt{a}b
+  final_answer = re.sub(r"(frac)([^{])(.)", "frac{\\2}{\\3}", final_answer)
+  final_answer = re.sub(r"(sqrt)([^{])", "sqrt{\\2}", final_answer)
+  final_answer = final_answer.replace("$", "")
+
+  # Normalize numbers
+  if final_answer.replace(",", "").isdigit():
+    final_answer = final_answer.replace(",", "")
+
+  return final_answer.strip()
+
+
 def check_answer(prompts, completions, answer, tmvp_config, **kargs):
   """
   Reward the model if the answer is correct. A reward is also given if the answer
@@ -105,6 +205,9 @@ def check_answer(prompts, completions, answer, tmvp_config, **kargs):
     if guess is None:
       scores.append(0)
       continue
+    if "DAPO" in tmvp_config.dataset_name:
+      guess = normalize_final_answer(guess)
+      true_answer = normalize_final_answer(true_answer)
     # Correct answer gets tmvp_config.reward_exact_format_match points!
     if guess == true_answer:
       score += tmvp_config.reward_exact_format_match
@@ -207,3 +310,68 @@ def get_optimizer(tmvp_config, max_train_steps):
         optimizer,
     )
   return optimizer
+
+
+def process_data(dataset_name, model_tokenizer, template_config, tmvp_config, x):
+  """Function to process input dataset"""
+
+  def _to_str(val):
+    if isinstance(val, bytes):
+      return val.decode("utf-8")
+    return str(val)
+
+  # Handle DAPO dataset schema
+  # originally (prompt is a list, answer is in reward_model)
+  # https://huggingface.co/datasets/BytedTsinghua-SIA/DAPO-Math-17k/viewer/default/train?row=0
+  # but using https://huggingface.co/datasets/open-r1/DAPO-Math-17k-Processed/viewer/all/train?row=1
+  # so question is prompt and answer is solution
+
+  question = x.get("question", x.get("prompt"))
+  answer = x.get("answer")
+  if answer is None and "solution" in x:
+    answer = x["solution"]
+
+  # Handle OpenMathInstruct-2
+  if "problem" in x:
+    question = x["problem"]
+  if "expected_answer" in x:
+    answer = x["expected_answer"]
+
+  # Handle AIME-2024
+  if "extra_info" in x and isinstance(x["extra_info"], dict) and "raw_problem" in x["extra_info"]:
+    question = x["extra_info"]["raw_problem"]
+
+  if "reward_model" in x and isinstance(x["reward_model"], dict) and "ground_truth" in x["reward_model"]:
+    answer = x["reward_model"]["ground_truth"]
+
+  question = _to_str(question)
+  answer = _to_str(answer)
+
+  if dataset_name == "gsm8k":
+    answer = extract_hash_answer(answer)
+
+  return {
+      # passed to model forward pass
+      "prompts": model_tokenizer.apply_chat_template(
+          [
+              {
+                  "role": "user",
+                  "content": template_config["TEMPLATE"].format(
+                      system_prompt=template_config["SYSTEM_PROMPT"].format(
+                          reasoning_start_token=tmvp_config.reasoning_start_token,
+                          reasoning_end_token=tmvp_config.reasoning_end_token,
+                          solution_start_token=tmvp_config.solution_start_token,
+                          solution_end_token=tmvp_config.solution_end_token,
+                      ),
+                      question=question,
+                  ),
+              },
+          ],
+          tokenize=False,
+          add_generation_prompt=True,
+      ),
+      # passed to reward functions
+      "question": question,
+      # passed to reward functions
+      "answer": answer,
+  }