RLLM

src/fairseq2/assets/cards/models/llama.yaml

@@ -168,4 +168,34 @@ model_arch: llama3_1_8b
 checkpoint: "hg://deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
 tokenizer: "hg://deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
 tokenizer_family: llama
-use_v2_tokenizer: true
+use_v2_tokenizer: true
+
+---
+
+name: octothinker_8b_hybrid
+model_family: llama
+model_arch: llama3_1_8b
+checkpoint: /datasets/pretrained-llms/OctoThinker-8B-Hybrid-Base/
+tokenizer: /datasets/pretrained-llms/OctoThinker-8B-Hybrid-Base/
+tokenizer_family: llama
+use_v2_tokenizer: true
+
+---
+
+name: octothinker_8b_long
+model_family: llama
+model_arch: llama3_1_8b
+checkpoint: /datasets/pretrained-llms/OctoThinker-8B-Long-Base/
+tokenizer: /datasets/pretrained-llms/OctoThinker-8B-Long-Base/
+tokenizer_family: llama
+use_v2_tokenizer: true
+
+---
+
+name: octothinker_8b_short
+model_family: llama
+model_arch: llama3_1_8b
+checkpoint: /datasets/pretrained-llms/OctoThinker-8B-Short-Base/
+tokenizer: /datasets/pretrained-llms/OctoThinker-8B-Short-Base/
+tokenizer_family: llama
+use_v2_tokenizer: true

src/fairseq2/recipes/lm/__init__.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -48,6 +48,12 @@
 from fairseq2.recipes.lm._online_finetune._generative_judge import (
     J1PairwiseScoreExtractorHandler as J1PairwiseScoreExtractorHandler,
 )
+from fairseq2.recipes.lm._online_finetune._generative_judge import (
+    J1KwiseScoreExtractor as J1KwiseScoreExtractor,
+)
+from fairseq2.recipes.lm._online_finetune._generative_judge import (
+    J1KwiseScoreExtractorHandler as J1KwiseScoreExtractorHandler,
+)
 from fairseq2.recipes.lm._online_finetune._generative_judge import (
     J1PointwiseExtractor as J1PointwiseExtractor,
 )
@@ -84,6 +90,12 @@
 from fairseq2.recipes.lm._online_finetune._remote_model import (
     NoEnvGeneralVerifierPipeline as NoEnvGeneralVerifierPipeline,
 )
+from fairseq2.recipes.lm._online_finetune._remote_model import (
+    NoEnvAceMathRMPipeline as NoEnvAceMathRMPipeline,
+)
+from fairseq2.recipes.lm._online_finetune._remote_model import (
+    NoEnvSkyworkRMPipeline as NoEnvSkyworkRMPipeline,
+)
 from fairseq2.recipes.lm._online_finetune._remote_model import (
     RemoteModelHandler as RemoteModelHandler,
 )
@@ -93,6 +105,18 @@
 from fairseq2.recipes.lm._online_finetune._rewards import (
     AtheneVerifierHandler as AtheneVerifierHandler,
 )
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    SkyworkVerifier as SkyworkVerifier,
+)
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    SkyworkVerifierHandler as SkyworkVerifierHandler,
+)
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    AceMathVerifier as AceMathVerifier,
+)
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    AceMathVerifierHandler as AceMathVerifierHandler,
+)
 from fairseq2.recipes.lm._online_finetune._rewards import (
     GenerativePairwiseVerifier as GenerativePairwiseVerifier,
 )
@@ -105,6 +129,12 @@
 from fairseq2.recipes.lm._online_finetune._rewards import (
     GenerativePointwiseVerifierHandler as GenerativePointwiseVerifierHandler,
 )
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    GenerativeKwiseVerifier as GenerativeKwiseVerifier,
+)
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    GenerativeKwiseVerifierHandler as GenerativeKwiseVerifierHandler,
+)
 from fairseq2.recipes.lm._online_finetune._rewards import GSM8kVerifier as GSM8kVerifier
 from fairseq2.recipes.lm._online_finetune._rewards import (
     GSM8kVerifierHandler as GSM8kVerifierHandler,

src/fairseq2/recipes/lm/_online_finetune/_common.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -8,6 +8,7 @@
 
 import contextlib
 import io
+import re
 from dataclasses import dataclass
 from typing import List, cast
 
@@ -17,14 +18,8 @@
 from torch import Tensor
 from vllm import RequestOutput
 
-from fairseq2.data import (
-    CollateOptionsOverride,
-    Collater,
-    SequenceData,
-)
-from fairseq2.datasets import (
-    SequenceBatch,
-)
+from fairseq2.data import CollateOptionsOverride, Collater, SequenceData
+from fairseq2.datasets import SequenceBatch
 from fairseq2.datasets.preference import PreferenceBatch
 from fairseq2.datasets.prompt import PromptBatch
 from fairseq2.gang import Gang, Gangs
@@ -93,9 +88,13 @@
 
     seq_data = cast(SequenceData, collater(to_collate))
 
+    seq_lens = seq_data["seqs"]["seq_lens"]
+    assert isinstance(seq_lens, Tensor) or isinstance(seq_lens, list)
+    if isinstance(seq_lens, Tensor):
+        seq_lens = seq_lens.tolist()
     batch = SequenceBatch(
         seq_data["seqs"]["seqs"],
-        seq_data["seqs"]["seq_lens"],
+        seq_lens,
         target_mask=seq_data["target_loss_mask"]["seqs"],
     )
     batch.to(device)
@@ -395,6 +394,8 @@
         prompt = prompt_batch.meta_info.get("prompt_raw")[0]
     elif "raw_prompt" in prompt_batch.meta_info:
         prompt = prompt_batch.meta_info.get("raw_prompt")[0]
+    elif "problem" in prompt_batch.meta_info:
+        prompt = prompt_batch.meta_info.get("problem")[0]
     else:
         # raw text prompt doesn't exist for this dataset
         prompt = "DUMMY PROMPT"
@@ -416,6 +417,48 @@
     return rollout_lengths
 
 
+def strip_think_tokens(rollouts: List[SequenceData]):
+    count_stripped, count_not_stripped, total_count, think_present = 0, 0, 0, 0
+    for sample in rollouts:
+        for rollout in sample.outputs:
+            rollout_text = rollout.text
+            if "<think>" in rollout_text:
+                think_present += 1
+            if rollout.finish_reason == "length":
+                count_not_stripped += 1
+            if rollout.finish_reason == "stop":
+                count_stripped += 1
+            total_count += 1
+            rollout.text = re.sub(
+                r"<think>.*?</think>", "", rollout_text, flags=re.DOTALL
+            ).strip()
+
+    log.info(f"Total count: {total_count}")
+    log.info(f"Think present: {think_present}")
+    log.info(f"Count stripped: {count_stripped/total_count}")
+    log.info(f"Count not stripped: {count_not_stripped/total_count}")
+
+    return rollouts
+
+def get_failed_to_parse_answers(reward_output: dict, batch_size: int):
+    if "answers" in reward_output:
+        log.info(f"Answers: {reward_output['answers']}")
+        failed_to_parse = sum(answer is None for rollouts in reward_output["answers"] for answer in rollouts)
+        return failed_to_parse/batch_size
+    else:
+        return 0.0
+
+def strip_for_octothinker(rollouts: List[SequenceData]):
+    for sample in rollouts:
+        for rollout in sample.outputs:
+            rollout_text = rollout.text
+            if "\nUser:" in rollout_text:
+                rollout_text = rollout_text[:rollout_text.find("\nUser:")]
+            rollout.text = rollout_text
+
+    return rollouts
+
+
 class StatefulRolloutBag:
     """A stateful container for managing and reusing model rollouts across multiple micro-batches.
 
@@ -504,11 +547,23 @@
 @torch.inference_mode()
 def update_avg_reward(metric_bag: MetricBag, avg_reward):
     metric_bag.get(Mean, "avg_reward").update(avg_reward, weight=1)
+
+@torch.inference_mode()
+def update_avg_second_reward(metric_bag: MetricBag, avg_reward):
+    metric_bag.get(Mean, "avg_second_reward").update(avg_reward, weight=1)
+
+@torch.inference_mode()
+def update_reward_matches(metric_bag: MetricBag, reward_matches):
+    metric_bag.get(Mean, "reward_matches").update(reward_matches, weight=1)
 
 
 @torch.inference_mode()
 def update_std_reward(metric_bag: MetricBag, std_reward):
     metric_bag.get(Mean, "std_reward").update(std_reward, weight=1)
+
+@torch.inference_mode()
+def update_failed_to_parse_answers(metric_bag: MetricBag, failed_to_parse_answers):
+    metric_bag.get(Mean, "failed_to_parse_answers").update(failed_to_parse_answers, weight=1)
 
 
 @torch.inference_mode()