Run formatter

JenniferWang · JenniferWang · commit 13f918c52506 · 2025-11-15T09:04:20.000-08:00
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
diff --git a/.meta/mast/hydrate_cache.py b/.meta/mast/hydrate_cache.py
@@ -14,6 +14,7 @@
 python .meta/mast/hydrate_cache.py --model-id Qwen/Qwen3-32B
 
 """
+
 import argparse
 import os
 import sys
diff --git a/apps/grpo/main.py b/apps/grpo/main.py
@@ -140,7 +140,6 @@ def simple_grpo_loss(
 
 @dataclass
 class RewardActor(ForgeActor):
-
     reward_functions: list[Callable]
 
     @endpoint
diff --git a/apps/sft/main.py b/apps/sft/main.py
@@ -287,7 +287,6 @@ def __repr__(self) -> str:
 
 
 async def run(cfg: DictConfig) -> None:
-
     logging.info("Spawning recipe...")
     process_cfg = cfg.pop("processes")
 
diff --git a/src/forge/actors/coder.py b/src/forge/actors/coder.py
@@ -10,10 +10,10 @@
 import tempfile
 from pathlib import Path
 
-from monarch.actor import endpoint
-
 from forge.controller import ForgeActor
 
+from monarch.actor import endpoint
+
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
diff --git a/src/forge/actors/generator.py b/src/forge/actors/generator.py
@@ -17,6 +17,28 @@
 
 import torch
 import torchstore as ts
+
+from forge.actors._torchstore_utils import (
+    extract_param_name,
+    get_dcp_whole_state_dict_key,
+    get_param_key,
+    get_param_prefix,
+    load_tensor_from_dcp,
+    rdma_available,
+)
+
+from forge.controller import (
+    ForgeActor,
+    get_proc_mesh,
+    host_mesh_from_proc,
+    stop_proc_mesh,
+)
+from forge.data_models.completion import Completion
+from forge.data_models.prompt import to_prompt
+from forge.observability.metrics import record_metric, Reduce
+from forge.observability.perf_tracker import Tracer
+from forge.types import ProcessConfig
+from forge.util._shared_tensor import SharedTensor, SharedTensorHandle
 from monarch.actor import current_rank, endpoint, ProcMesh, this_host
 
 from vllm.config import VllmConfig
@@ -42,28 +64,6 @@
 from vllm.v1.structured_output import StructuredOutputManager
 from vllm.worker.worker_base import WorkerWrapperBase
 
-from forge.actors._torchstore_utils import (
-    extract_param_name,
-    get_dcp_whole_state_dict_key,
-    get_param_key,
-    get_param_prefix,
-    load_tensor_from_dcp,
-    rdma_available,
-)
-
-from forge.controller import (
-    ForgeActor,
-    get_proc_mesh,
-    host_mesh_from_proc,
-    stop_proc_mesh,
-)
-from forge.data_models.completion import Completion
-from forge.data_models.prompt import to_prompt
-from forge.observability.metrics import record_metric, Reduce
-from forge.observability.perf_tracker import Tracer
-from forge.types import ProcessConfig
-from forge.util._shared_tensor import SharedTensor, SharedTensorHandle
-
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
diff --git a/src/forge/actors/reference_model.py b/src/forge/actors/reference_model.py
@@ -13,6 +13,11 @@
 from dataclasses import dataclass, field, fields
 
 import torch
+
+from forge.controller import ForgeActor
+from forge.observability.metrics import record_metric, Reduce
+from forge.observability.perf_tracker import Tracer
+from forge.util.ops import compute_logprobs
 from monarch.actor import current_rank, current_size, endpoint
 from torch.distributed.tensor import DTensor
 
@@ -27,11 +32,6 @@
 from torchtitan.experiments.forge.engine import ForgeEngine
 from torchtitan.experiments.forge.job_config import ForgeJobConfig
 
-from forge.controller import ForgeActor
-from forge.observability.metrics import record_metric, Reduce
-from forge.observability.perf_tracker import Tracer
-from forge.util.ops import compute_logprobs
-
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
diff --git a/src/forge/actors/replay_buffer.py b/src/forge/actors/replay_buffer.py
@@ -11,12 +11,12 @@
 from operator import itemgetter
 from typing import Any, Callable
 
-from monarch.actor import endpoint
-
 from forge.controller import ForgeActor
 from forge.observability.metrics import record_metric, Reduce
 from forge.observability.perf_tracker import trace
 
+from monarch.actor import endpoint
+
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
diff --git a/src/forge/actors/trainer/titan.py b/src/forge/actors/trainer/titan.py
@@ -16,6 +16,18 @@
 import torch.distributed.checkpoint as dcp
 import torchstore as ts
 
+from forge.actors._torchstore_utils import (
+    DcpHandle,
+    get_dcp_whole_state_dict_key,
+    get_param_key,
+    rdma_available,
+)
+
+from forge.controller import ForgeActor
+from forge.data.utils import batch_to_device
+from forge.observability.metrics import record_metric, Reduce
+from forge.observability.perf_tracker import Tracer
+
 from monarch.actor import endpoint
 from torch import Tensor
 from torch.distributed.checkpoint._nested_dict import flatten_state_dict
@@ -36,18 +48,6 @@
 from torchtitan.experiments.forge.engine import ForgeEngine
 from torchtitan.experiments.forge.job_config import ForgeJobConfig
 
-from forge.actors._torchstore_utils import (
-    DcpHandle,
-    get_dcp_whole_state_dict_key,
-    get_param_key,
-    rdma_available,
-)
-
-from forge.controller import ForgeActor
-from forge.data.utils import batch_to_device
-from forge.observability.metrics import record_metric, Reduce
-from forge.observability.perf_tracker import Tracer
-
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
diff --git a/src/forge/controller/launcher.py b/src/forge/controller/launcher.py
@@ -16,6 +16,8 @@
 
 import monarch
 import torchx.specs as specs
+
+from forge.types import Launcher, LauncherConfig
 from monarch._rust_bindings.monarch_hyperactor.alloc import AllocConstraints
 from monarch._rust_bindings.monarch_hyperactor.channel import ChannelTransport
 from monarch._rust_bindings.monarch_hyperactor.config import configure
@@ -25,8 +27,6 @@
 from monarch.tools.commands import create, info
 from monarch.tools.config import Config, Workspace
 
-from forge.types import Launcher, LauncherConfig
-
 _MAST_AVAILABLE = False
 
 try:
@@ -269,7 +269,6 @@ def add_additional_packages(self, packages: "Packages") -> "Packages":
         return packages
 
     def build_appdef(self) -> specs.AppDef:
-
         # create the app definition for the worker
         remote_end_python_path = ":".join(
             [
diff --git a/src/forge/controller/provisioner.py b/src/forge/controller/provisioner.py
@@ -5,6 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 """Remote and local resource manager for allocation and provisioning."""
+
 import asyncio
 import logging
 
@@ -14,6 +15,10 @@
 
 import torch
 
+from forge.controller.launcher import BaseLauncher, get_launcher
+from forge.env import all_env_vars, FORGE_DISABLE_METRICS
+from forge.types import ProcessConfig, ProvisionerConfig
+
 from monarch._src.actor.actor_mesh import ActorMesh
 from monarch._src.actor.shape import Extent
 
@@ -22,10 +27,6 @@
 from monarch.tools import commands
 from monarch.utils import setup_env_for_distributed
 
-from forge.controller.launcher import BaseLauncher, get_launcher
-from forge.env import all_env_vars, FORGE_DISABLE_METRICS
-from forge.types import ProcessConfig, ProvisionerConfig
-
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
@@ -586,7 +587,6 @@ async def shutdown_metric_logger():
 
 
 async def shutdown():
-
     await shutdown_metric_logger()
 
     logger.info("Shutting down provisioner..")
diff --git a/src/forge/controller/service/replica.py b/src/forge/controller/service/replica.py
@@ -12,11 +12,11 @@
 from dataclasses import dataclass, field
 from enum import Enum
 
-from monarch.actor import ActorError
-
 from forge.controller import ForgeActor
 from forge.types import ProcessConfig
 
+from monarch.actor import ActorError
+
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
diff --git a/src/forge/controller/service/service.py b/src/forge/controller/service/service.py
@@ -37,8 +37,6 @@
 import pprint
 import uuid
 
-from monarch.actor import Actor, endpoint
-
 from forge.controller.service.interface import _session_context, Session
 
 from forge.controller.service.metrics import ServiceMetrics
@@ -51,6 +49,8 @@
 )
 from forge.types import ServiceConfig
 
+from monarch.actor import Actor, endpoint
+
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
diff --git a/src/forge/controller/service/spawn.py b/src/forge/controller/service/spawn.py
@@ -8,13 +8,13 @@
 import logging
 from typing import Type
 
-from monarch.actor import proc_mesh
-
 from forge.controller import ForgeActor
 from forge.controller.service import ServiceActor, ServiceConfig
 
 from forge.controller.service.interface import ServiceInterfaceV2
 
+from monarch.actor import proc_mesh
+
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
diff --git a/src/forge/data/datasets/packed.py b/src/forge/data/datasets/packed.py
@@ -10,14 +10,14 @@
 from typing import Any, Generic, Iterable, Iterator, TypeVar
 
 import torch
+
+from forge.data import CROSS_ENTROPY_IGNORE_IDX
+from forge.observability.metrics import Metric, Reduce
 from torch.nn.attention.flex_attention import (
     create_block_mask as create_block_mask_flex,
 )
 from torchdata.stateful_dataloader import Stateful
 
-from forge.data import CROSS_ENTROPY_IGNORE_IDX
-from forge.observability.metrics import Metric, Reduce
-
 from .dataset import DatasetInfo, InfiniteTuneIterableDataset
 
 logger = logging.getLogger(__name__)
@@ -464,7 +464,6 @@ def __iter__(self) -> Iterator[SampleDict]:
 
         # Main packing loop
         while True:
-
             # Stop if the source is exhausted and there's no data left to pack
             if self._exhausted and not self._buffer and self._current_pack_size == 0:
                 break
diff --git a/src/forge/data/datasets/sft_dataset.py b/src/forge/data/datasets/sft_dataset.py
@@ -123,7 +123,6 @@ class SFTOutputTransform:
     """
 
     def __call__(self, sample: dict[str, Any]) -> dict[str, Any]:
-
         # Sanity checks
         if not isinstance(sample["tokens"], torch.Tensor):
             sample["tokens"] = torch.tensor(sample["tokens"])
diff --git a/src/forge/data/tokenizer.py b/src/forge/data/tokenizer.py
@@ -8,13 +8,13 @@
 from typing import Any, Optional
 
 import jinja2
-from jinja2 import StrictUndefined
-
-from tokenizers import Tokenizer
 
 from forge.data.utils import truncate
 from forge.interfaces import BaseTokenizer, ModelTokenizer
 from forge.types import Message
+from jinja2 import StrictUndefined
+
+from tokenizers import Tokenizer
 
 
 class HuggingFaceBaseTokenizer(BaseTokenizer):
diff --git a/src/forge/losses/reinforce_loss.py b/src/forge/losses/reinforce_loss.py
@@ -5,9 +5,9 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
-from torch import nn
 
 from forge.util.ops import compute_logprobs
+from torch import nn
 
 
 class ReinforceLoss(nn.Module):
diff --git a/src/forge/observability/metric_actors.py b/src/forge/observability/metric_actors.py
@@ -9,14 +9,6 @@
 import uuid
 from typing import Any, Union
 
-from monarch.actor import (
-    context,
-    endpoint,
-    get_or_spawn_controller,
-    ProcMesh,
-    this_proc,
-)
-
 from forge.controller.actor import ForgeActor
 
 from forge.env import FORGE_DISABLE_METRICS
@@ -29,6 +21,14 @@
     reduce_metrics_states,
 )
 
+from monarch.actor import (
+    context,
+    endpoint,
+    get_or_spawn_controller,
+    ProcMesh,
+    this_proc,
+)
+
 
 logger = logging.getLogger(__name__)
 
diff --git a/src/forge/observability/metrics.py b/src/forge/observability/metrics.py
@@ -13,11 +13,11 @@
 from typing import Any
 
 import pytz
-from monarch.actor import current_rank
 
 from forge.observability.utils import get_proc_name_with_rank
 
 from forge.util.logging import get_logger, log_once
+from monarch.actor import current_rank
 
 logger = get_logger("INFO")
 
@@ -654,7 +654,6 @@ class LoggerBackend(ABC):
     def __init__(
         self, *, logging_mode: LoggingMode, per_rank_share_run: bool = False, **kwargs
     ) -> None:
-
         self.logging_mode = logging_mode
         self.per_rank_share_run = per_rank_share_run
         self.backend_kwargs = kwargs
diff --git a/src/forge/observability/perf_tracker.py b/src/forge/observability/perf_tracker.py
diff --git a/tests/sandbox/toy_rl/toy_metrics/main.py b/tests/sandbox/toy_rl/toy_metrics/main.py
diff --git a/tests/unit_tests/datasets/test_hf.py b/tests/unit_tests/datasets/test_hf.py
diff --git a/tests/unit_tests/datasets/test_interleaved.py b/tests/unit_tests/datasets/test_interleaved.py
diff --git a/tests/unit_tests/observability/conftest.py b/tests/unit_tests/observability/conftest.py
diff --git a/tests/unit_tests/observability/test_metrics.py b/tests/unit_tests/observability/test_metrics.py
diff --git a/tests/unit_tests/observability/test_perf_tracker.py b/tests/unit_tests/observability/test_perf_tracker.py
diff --git a/tests/unit_tests/test_coder.py b/tests/unit_tests/test_coder.py
