Change online training verifier

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

@@ -66,15 +66,17 @@ num_shards: 8
 
 name: llama3
 model_family: llama
-checkpoint: "https://ai.meta.com/llama/;gated=true"
-tokenizer: "https://ai.meta.com/llama/;gated=true"
+checkpoint: "/datasets/pretrained-llms/Llama-3.1-8B/"
+tokenizer: "/datasets/pretrained-llms/Llama-3.1-8B/"
 tokenizer_family: llama
 use_v2_tokenizer: true
 
 ---
 
 name: llama3_instruct
 base: llama3
+checkpoint: "/datasets/pretrained-llms/Llama-3.1-8B-Instruct/"
+tokenizer: "/datasets/pretrained-llms/Llama-3.1-8B-Instruct/"
 use_eot: true  # instruct tokenizer to use EOT instead of EOS
 
 ---
@@ -168,4 +170,4 @@ 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

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

@@ -139,8 +139,8 @@ use_im_end: true
 name: qwen3_8b_base
 model_family: qwen
 model_arch: qwen3_8b
-checkpoint: "hg://qwen/qwen3-8b-base"
-tokenizer: "hg://qwen/qwen3-8b-base"
+checkpoint: "/checkpoint/data/jacklanchantin/pretrained-llms/Qwen3-8B-Base/"
+tokenizer: "/checkpoint/data/jacklanchantin/pretrained-llms/Qwen3-8B-Base/"
 tokenizer_family: qwen
 
 ---

src/fairseq2/datasets/prompt.py

@@ -232,8 +232,13 @@ def create_reader(
         def copy_prompt_text(example: dict[str, Any]) -> dict[str, Any]:
             example[f"{src_key}_text"] = example[f"{src_key}"]
             return example
+
+        def copy_prompt_raw_text(example: dict[str, Any]) -> dict[str, Any]:
+            example[f"raw_prompt_text"] = example[f"{src_key}"]
+            return example
 
         builder.map(copy_prompt_text, num_parallel_calls=npc)
+        builder.map(copy_prompt_raw_text, num_parallel_calls=npc)
 
         # Encode source and target texts.
         source_encoder = tokenizer.create_encoder(mode=options.source_encode_mode)

src/fairseq2/metrics/recorders/_wandb.py

@@ -150,6 +150,9 @@ def create(
         run_id = self._get_run_id(output_dir, config)
 
         try:
+            # Increase timeout for fairwandb.org which can be slow
+            settings = wandb.Settings(init_timeout=1200)  # 20 minutes
+
             run = wandb.init(
                 entity=config.entity,
                 project=config.project,
@@ -160,6 +163,7 @@ def create(
                 group=config.group,
                 job_type=config.job_type,
                 resume=config.resume_mode,
+                settings=settings,
             )
         except (RuntimeError, ValueError) as ex:
             raise MetricRecordError(

src/fairseq2/recipes/_trainer.py

@@ -189,7 +189,7 @@ def __init__(
         device_stat_tracker: DeviceStatTracker,
         wall_watch: Stopwatch,
         progress_reporter: ProgressReporter,
-        fp16_loss_scale: tuple[float, float] = (128.0, 0.0001),
+        fp16_loss_scale: tuple[float, float] = (16, 0.0001),
         no_sync_grad_accumulation: bool = False,
         max_grad_norm: float | None = None,
         grad_check: bool = False,

src/fairseq2/recipes/config.py

@@ -165,7 +165,7 @@ class TrainerSection:
     max_grad_norm: float | None = None
     """The maximum gradient norm. If ``None``, no clipping will be applied."""
 
-    fp16_loss_scale: tuple[float, float] = (128.0, 0.0001)
+    fp16_loss_scale: tuple[float, float] = (16, 0.0001)
     """The initial and minimum loss scale for fp16 training."""
 
     gc_every_n_steps: int | None = None

src/fairseq2/recipes/lm/__init__.py

@@ -54,6 +54,18 @@
 from fairseq2.recipes.lm._online_finetune._generative_judge import (
     J1PointwiseExtractorHandler as J1PointwiseExtractorHandler,
 )
+from fairseq2.recipes.lm._online_finetune._generative_judge import (
+    SelfAugmentingExtractor as SelfAugmentingExtractor,
+)
+from fairseq2.recipes.lm._online_finetune._generative_judge import (
+    SelfAugmentingExtractorHandler as SelfAugmentingExtractorHandler,
+)
+from fairseq2.recipes.lm._online_finetune._generative_judge import (
+    SelfAugmentingExtractor2 as SelfAugmentingExtractor2,
+)
+from fairseq2.recipes.lm._online_finetune._generative_judge import (
+    SelfAugmentingExtractorHandler2 as SelfAugmentingExtractorHandler2,
+)
 from fairseq2.recipes.lm._online_finetune._generative_judge import (
     JudgmentExtractorHandler as JudgmentExtractorHandler,
 )
@@ -116,6 +128,12 @@
 from fairseq2.recipes.lm._online_finetune._rewards import (
     MathVerifyHandler as MathVerifyHandler,
 )
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    PplDerivedVerifier as PplDerivedVerifier,
+)
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    PplDerivedVerifierHandler as PplDerivedVerifierHandler,
+)
 from fairseq2.recipes.lm._online_finetune._rewards import (
     VLLMOutputRewardHandler as VLLMOutputRewardHandler,
 )

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

@@ -129,7 +129,7 @@ def generate_rollouts(
             flat_request_list.extend(rank_prompts)
 
         rollouts = vllm_model.rollout_from_model(
-            flat_request_list, sampling_params=sampling_params
+            prompt_list=flat_request_list, sampling_params=sampling_params
         )
 
         rollouts_to_scatter = []
@@ -361,6 +361,7 @@ def combine_prompts_responses_for_scoring(
 
 def get_vllm_logprobs(
     vllm_outputs: List[RequestOutput],
+    model_logps: Tensor,
     gangs,
     rollout_start_end: tuple[int, int] | None = None,
 ):
@@ -405,6 +406,9 @@ def get_vllm_logprobs(
     for i, t in enumerate(sequences):
         padded[i, : t.size(0)] = t
 
+    # clip outputs to be same size as model_logps
+    if padded.size() != model_logps.size():
+        padded = padded[:, : model_logps.size(1)]
     return padded
 
 
@@ -421,11 +425,47 @@ def convert_vllm_output_to_ref_score(vllm_outputs: List[RequestOutput], gangs):
 
 
 def compute_token_level_entropy(logits: torch.Tensor, target_mask: torch.Tensor):
-    """Calculate entropy from logits. Returns sum of entropies averages for each sequence."""
-    pd = torch.nn.functional.softmax(logits, dim=-1)
-    entropy = torch.logsumexp(logits, dim=-1) - torch.sum(pd * logits, dim=-1)
+    """Calculate entropy from logits. Returns sum of entropies averages for each sequence.
+
+    Uses numerically stable entropy calculation: H = log(Z) - E[logits]
+    where Z = sum(exp(logits)) and E is the expectation over softmax distribution.
+
+    Note: Entropy is computed without gradients to avoid numerical instability.
+    Memory-efficient implementation processes sequences in chunks.
+    """
+    # Compute entropy without gradients to avoid backprop instability
+    with torch.no_grad():
+        logits_clamped = torch.clamp(logits, min=-100, max=100)
+
+        # Compute logsumexp for normalization constant
+        log_z = torch.logsumexp(logits_clamped, dim=-1)  # [batch, seq_len]
+
+        # Compute expected logits: E[logits] = sum(softmax(logits) * logits)
+        # Process in smaller chunks to avoid OOM
+        batch_size, seq_len, vocab_size = logits_clamped.shape
+        chunk_size = 128  # Process 128 positions at a time
+
+        expected_logits = torch.zeros(batch_size, seq_len, device=logits.device)
+        for i in range(0, seq_len, chunk_size):
+            end_i = min(i + chunk_size, seq_len)
+            chunk_logits = logits_clamped[:, i:end_i, :]
+            chunk_probs = torch.softmax(chunk_logits, dim=-1)
+            expected_logits[:, i:end_i] = torch.sum(chunk_probs * chunk_logits, dim=-1)
+            del chunk_probs, chunk_logits  # Free memory immediately
+
+        # Entropy = log(Z) - E[logits]
+        entropy = log_z - expected_logits
+
+        # Clamp entropy to prevent NaN values
+        entropy = torch.clamp(entropy, min=0.0, max=100.0)
+
     entropy_target_only = entropy * target_mask
-    entropy_per_seq = entropy_target_only.sum(dim=-1) / target_mask.sum(dim=-1)
+
+    # Compute target token counts and avoid division by zero
+    target_counts = target_mask.sum(dim=-1)
+    # Add small epsilon to avoid division by zero
+    target_counts = torch.clamp(target_counts, min=1e-8)
+    entropy_per_seq = entropy_target_only.sum(dim=-1) / target_counts
 
     return entropy_per_seq
 
@@ -438,6 +478,10 @@ def log_rollouts(prompt_batch: PromptBatch, rollouts, split_name, num_rollouts=1
         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 "raw_prompt_text" in prompt_batch.meta_info:
+        prompt = prompt_batch.meta_info.get("raw_prompt_text")[0]
+    elif "prefix_text" in prompt_batch.meta_info:
+        prompt = prompt_batch.meta_info.get("prefix_text")[0]
     else:
         # raw text prompt doesn't exist for this dataset
         prompt = "DUMMY PROMPT"
@@ -459,6 +503,46 @@ def get_rollout_lengths(rollouts: List[SequenceData]):
     return rollout_lengths
 
 
+def get_think_rollout_lengths(rollouts: List[SequenceData]):
+    """Get the lengths of tokens before the </think> tag in rollouts.
+
+    This function calculates the approximate number of tokens generated before
+    the </think> closing tag in each rollout. It uses a proportional approximation
+    based on character positions to estimate token counts.
+
+    Args:
+        rollouts: List of SequenceData containing rollout outputs
+
+    Returns:
+        List of token lengths before </think> tag for rollouts that contain the tag
+    """
+    think_rollout_lengths = []
+    think_tag = "</think>"
+
+    for rollout in rollouts:
+        for sample in rollout.outputs:
+            rollout_text = sample.text
+            if think_tag in rollout_text:
+                # Find the position of </think> in the text
+                think_end_pos = rollout_text.find(think_tag) + len(think_tag)
+                # Count tokens up to and including </think>
+                # We need to find how many tokens correspond to the text before </think>
+                # Since we have token_ids, we'll approximate by finding the proportion
+                text_before_think = rollout_text[:think_end_pos]
+                total_text = rollout_text
+                total_tokens = len(sample.token_ids)
+                # Approximate token count proportionally (rough estimate)
+                # A better approach would be to tokenize text_before_think, but we use approximation
+                think_token_length = (
+                    int((len(text_before_think) / len(total_text)) * total_tokens)
+                    if len(total_text) > 0
+                    else 0
+                )
+                think_rollout_lengths.append(think_token_length)
+
+    return think_rollout_lengths
+
+
 class StatefulRolloutBag:
     """A stateful container for managing and reusing model rollouts across multiple micro-batches.
 
@@ -559,6 +643,13 @@ def update_avg_rollout_length(metric_bag: MetricBag, avg_rollout_length):
     metric_bag.get(Mean, "avg_rollout_length").update(avg_rollout_length, weight=1)
 
 
+@torch.inference_mode()
+def update_avg_think_rollout_length(metric_bag: MetricBag, avg_think_rollout_length):
+    metric_bag.get(Mean, "avg_think_rollout_length").update(
+        avg_think_rollout_length, weight=1
+    )
+
+
 @torch.inference_mode()
 def update_avg_reward_len_norm(metric_bag: MetricBag, avg_reward_len_norm):
     metric_bag.get(Mean, "avg_reward_len_norm").update(avg_reward_len_norm, weight=1)
@@ -599,7 +690,9 @@ def update_grpo_batch_metrics(
 
 
 @torch.inference_mode()
-def update_grpo_loss(metric_bag: MetricBag, batch: PromptBatch, loss: Tensor) -> None:
+def update_grpo_loss(
+    metric_bag: MetricBag, batch: PromptBatch, loss: Tensor, tis_imp_ratio: Tensor
+) -> None:
     """Update the GRPO loss metric.
 
     :param batch:
@@ -611,6 +704,8 @@ def update_grpo_loss(metric_bag: MetricBag, batch: PromptBatch, loss: Tensor) ->
         loss / batch.batch_size, weight=batch.batch_size
     )
 
+    metric_bag.get(Mean, "tis_imp_ratio").update(tis_imp_ratio)
+
 
 def compute_reference_logps(
     gangs: Gangs,
@@ -651,3 +746,10 @@ def get_parameter_converter(model_config):
         raise RuntimeError(f"{model_config} not supported in online recipe")
 
     return _convert_parameter
+
+
+def masked_mean(x, mask, eps=1e-8):
+    # mask: boolean tensor with same shape as x
+    masked_x = x * mask
+    denom = mask.sum().clamp_min(eps)
+    return masked_x.sum() / denom