More lint

Author: Jack-Khuu

apps/grpo/main.py

@@ -12,17 +12,17 @@
 
 import torch
 import torch.nn.functional as F
-from torchtitan.config.job_config import Model as TitanJobModelConfig
 from datasets import load_dataset
 from forge.actors.policy import Policy, PolicyConfig, SamplingOverrides, WorkerConfig
-# from forge.actors.reference_model import ReferenceModel
+from forge.actors.reference_model import ReferenceModel  # noqa: F401
 from forge.actors.replay_buffer import ReplayBuffer
 from forge.controller.actor import ForgeActor
 from forge.controller.service import ServiceConfig, shutdown_service, spawn_service
 from forge.data.rewards import MathReward, ThinkingReward
 from forge.util.metric_logging import get_metric_logger
 from monarch.actor import endpoint
 from torch import nn
+from torchtitan.config.job_config import Model as TitanJobModelConfig
 from transformers import AutoModelForCausalLM
 from vllm.transformers_utils.tokenizer import get_tokenizer
 

src/forge/actors/reference_model.py

@@ -15,7 +15,7 @@
 import torch
 from monarch.actor import current_rank, current_size, endpoint
 
-from torchtitan.config.job_config import Compile, Checkpoint, Model, Parallelism
+from torchtitan.config.job_config import Checkpoint, Compile, Model, Parallelism
 from torchtitan.distributed import utils as dist_utils
 from torchtitan.experiments.forge.engine import ForgeEngine
 from torchtitan.experiments.forge.job_config import ForgeJobConfig
@@ -84,7 +84,7 @@ async def setup(self):
         self.engine = ForgeEngine(ForgeJobConfig(**engine_config))
 
     @endpoint
-    async def forward(self, episode: 'Episode') -> torch.Tensor:
+    async def forward(self, episode: "Episode") -> torch.Tensor:
         """
         Given an episode, return the log_probability of the
         token_ids, shape (completion_len, )