Merge branch 'meta-pytorch:main' into main

Author: DNXie

.github/workflows/unit_test.yaml

@@ -23,9 +23,12 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Update pip
         run: python -m pip install --upgrade pip
+      - name: Install pytorch
+        run: python -m pip install torch==2.9.0.dev20250826 --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+      - name: Install monarch
+        run: python -m pip install monarch-no-torch==0.1.0.dev20250826 --find-links assets/wheels
       - name: Install dependencies
-        run: |
-          python -m pip install --no-build-isolation -e ".[dev,cpu]" --extra-index-url https://download.pytorch.org/whl/nightly/cpu --find-links assets/wheels
+        run: python -m pip install --no-build-isolation -e ".[dev]"
       - name: Run unit tests with coverage
         # TODO add all tests
         run: pytest tests/unit_tests --cov=. --cov-report=xml --durations=20 -vv

README.md

@@ -35,14 +35,14 @@ source .venv/bin/activate
 
 ```bash
 # feature install if you don't have /user/local/cuda-12.8
-feature install --persist cuda_12_8
+feature install --persist cuda_12_9
 
 # add env variables
-export CUDA_VERSION=12.8
-export NVCC=/usr/local/cuda-${CUDA_VERSION}/bin/nvcc
-export CUDA_NVCC_EXECUTABLE=/usr/local/cuda-${CUDA_VERSION}/bin/nvcc
-export CUDA_HOME=/usr/local/cuda-${CUDA_VERSION}
-export PATH="${CUDA_HOME}/bin:$PATH"
+export CUDA_VERSION=12.9
+export NVCC=/usr/local/cuda-$CUDA_VERSION/bin/nvcc
+export CUDA_NVCC_EXECUTABLE=/usr/local/cuda-$CUDA_VERSION/bin/nvcc
+export CUDA_HOME=/usr/local/cuda-$CUDA_VERSION
+export PATH="$CUDA_HOME/bin:$PATH"
 export CUDA_INCLUDE_DIRS=$CUDA_HOME/include
 export CUDA_CUDART_LIBRARY=$CUDA_HOME/lib64/libcudart.so
 export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH

apps/grpo/main.py

@@ -12,8 +12,11 @@
 import torch
 from datasets import load_dataset
 from forge.actors.policy import Policy, PolicyConfig, SamplingOverrides, WorkerConfig
+from forge.actors.replay_buffer import ReplayBuffer
 from forge.controller import ServiceConfig, spawn_service
 from forge.controller.actor import ForgeActor
+from forge.data.rewards import MathReward, ThinkingReward
+from forge.util.metric_logging import get_metric_logger
 from monarch.actor import endpoint
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
@@ -209,66 +212,18 @@ async def update_weights(self, policy_actor):
         self.logger.info(f"Updating weights took {end_time - start_time:.2f} seconds")
 
 
-def math_scoring_function(prompt: str, response: str, target: str) -> float:
-    """Function to score math correctness."""
-    import re
-
-    # Extract expected answer from target
-    expected_answer = (
-        float(target.strip())
-        if target.strip().replace(".", "").replace("-", "").isdigit()
-        else None
-    )
-
-    # Extract model answer from response
-    patterns = [
-        r"####\s*([+-]?\d+(?:\.\d+)?)",  # GSM8K style answer format
-        r"(?:the\s+)?answer\s+is\s*([+-]?\d+(?:\.\d+)?)",
-        r"(?:answer:|result:)\s*([+-]?\d+(?:\.\d+)?)",
-        r"=\s*([+-]?\d+(?:\.\d+)?)\s*(?:\.|$)",  # equals near end
-        r"\b([+-]?\d+(?:\.\d+)?)\s*(?:\.|$)",  # number at end
-        r"([+-]?\d+(?:\.\d+)?)",  # any number (fallback)
-    ]
-
-    model_answer = None
-    response_lower = response.lower().strip()
-    for pattern in patterns:
-        matches = re.findall(pattern, response_lower)
-        if matches:
-            model_answer = float(matches[-1])
-            break
-
-    if expected_answer is None or model_answer is None:
-        return 0.1  # Partial credit for attempting
-
-    # Check if answers match (with some tolerance for floating point)
-    if abs(expected_answer - model_answer) < 1e-6:
-        return 1.0  # Correct answer
-    else:
-        return 0.0  # Incorrect answer
-
-
-def thinking_scoring_function(prompt: str, response: str, target: str) -> float:
-    """Function to score thinking tag usage."""
-    # Check if response contains <think></think> tags
-    if "<think>" in response.lower() and "</think>" in response.lower():
-        return 0.5
-    else:
-        return 0.0
-
-
 class RewardActor(ForgeActor):
     """Reward actor that uses a list of scoring functions."""
 
-    def __init__(self, scoring_functions: list[Callable]):
+    def __init__(self, reward_functions: list[Callable]):
         super().__init__()
-        self.scoring_functions = scoring_functions
+        self.reward_functions = reward_functions
 
     @endpoint
     async def evaluate_response(self, prompt: str, response: str, target: str) -> float:
         total_reward = 0.0
-        for scoring_fn in self.scoring_functions:
-            reward = scoring_fn(prompt, response, target)
+        for reward_fn in self.reward_functions:
+            reward = reward_fn(prompt, response, target)
             total_reward += reward
         return total_reward
 
@@ -388,6 +343,13 @@ async def main():
     group_size = 1
     model = "Qwen/Qwen3-1.7B"
 
+    # ---- Setup WandB Logger ---- #
+    logger = get_metric_logger(
+        "wandb",
+        freq=1,
+        project="grpo-training",
+    )
+
     # ---- Setup services ---- #
     default_service_cfg = ServiceConfig(
         procs_per_replica=1,
@@ -447,7 +409,7 @@ async def main():
     reward_actor = await spawn_service(
         default_service_cfg,
         RewardActor,
-        scoring_functions=[math_scoring_function, thinking_scoring_function],
+        reward_functions=[MathReward(), ThinkingReward()],
     )
 
     print("All services initialized successfully!")
@@ -498,6 +460,7 @@ async def continuous_rollouts():
                 print(
                     f"Generated {rollout_count} rollouts w/ average reward {avg_reward}"
                 )
+                logger.log("reward/rollout", avg_reward, rollout_count)
 
     async def continuous_training():
         training_step = 0
@@ -511,7 +474,9 @@ async def continuous_training():
                 if training_step % 10 == 0:
                     print(f"Completed {training_step} training steps")
                     if training_result:
-                        print(f"Latest loss: {training_result.get('loss', 'N/A')}")
+                        loss_value = training_result.get("loss", 0.0)
+                        print(f"Latest loss: {loss_value}")
+                        logger.log("loss/training_step", loss_value, training_step)
                 # await trainer.update_weights(policy)
 
     print("Starting GRPO training loops...")

pyproject.toml

@@ -47,11 +47,6 @@ oss = [
     "torchmonarch-nightly==2025.8.1",
     "torchstore",
 ]
-cpu = [
-    "torch==2.9.0.dev20250826",
-    "monarch-no-torch==0.1.0.dev20250826",
-]
-
 
 [project.scripts]
 forge = "forge.cli.forge:main"
@@ -78,8 +73,8 @@ members = [
 # pytorch
 # TODO: get auto backend to work
 [[tool.uv.index]]
-name = "pytorch-nightly-cu128"
-url  = "https://download.pytorch.org/whl/nightly/cu128"
+name = "pytorch-nightly-cu129"
+url  = "https://download.pytorch.org/whl/nightly/cu129"
 #explicit = true
 
 # vllm
@@ -89,8 +84,8 @@ url  = "https://download.pytorch.org/whl/nightly/cu128"
 # explicit = true
 
 [tool.uv.sources]
-torchtitan = { index = "pytorch-nightly-cu128" }
-torch = { index = "pytorch-nightly-cu128" }
+torchtitan = { index = "pytorch-nightly-cu129" }
+torch = { index = "pytorch-nightly-cu129" }
 torchstore = { git = "ssh://[email protected]/meta-pytorch/torchstore.git" }
 #vllm = { index = "vllm-nightly" }
 

src/forge/cli/config.py

@@ -1,3 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
 import argparse
 import functools
 import sys

src/forge/data/rewards.py

@@ -1,3 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
 import re
 from typing import Optional