Ruff

Author: JacobHelwig

verl/experimental/agent_loop/agent_loop.py

@@ -478,7 +478,9 @@ async def generate_sequences(self, batch: DataProto) -> DataProto:
             )
         outputs = await asyncio.gather(*tasks)
 
-        output = self._postprocess(outputs, input_non_tensor_batch=batch.non_tensor_batch, validate=batch.meta_info.get("validate", False))
+        output = self._postprocess(
+            outputs, input_non_tensor_batch=batch.non_tensor_batch, validate=batch.meta_info.get("validate", False)
+        )
         return output
 
     async def _run_agent_loop(
@@ -736,10 +738,18 @@ async def _compute_score(self, output, prompts, responses, attention_mask, input
     async def _compute_teacher_logprobs(self, output, prompt_ids, response_ids, validate):
         """Compute teacher logprobs for single sample."""
         if self.distillation_enabled and not validate:
-            data = DataProto(batch=TensorDict({"prompt_ids": torch.tensor([prompt_ids]), "response_ids": torch.tensor([response_ids])}, batch_size=1))
+            data = DataProto(
+                batch=TensorDict(
+                    {"prompt_ids": torch.tensor([prompt_ids]), "response_ids": torch.tensor([response_ids])},
+                    batch_size=1,
+                )
+            )
             selected_teacher_loop_worker_handle = random.choice(self.teacher_loop_worker_handles)
             result = await selected_teacher_loop_worker_handle.compute_logprobs.remote(data)
-            response_ids, response_logprobs = result["response_ids"], result["response_logprobs"] # (1, S, K), S=sequence length, K=topk/1
+            response_ids, response_logprobs = (
+                result["response_ids"],
+                result["response_logprobs"],
+            )  # (1, S, K), S=sequence length, K=topk/1
 
             pad_size = self.config.actor_rollout_ref.rollout.response_length - response_ids.shape[1]
             padding = (0, 0, 0, pad_size)  # pad the sequence dimension
@@ -976,7 +986,9 @@ def _init_agent_loop_workers(self):
                     scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
                         node_id=node_id, soft=True
                     ),
-                ).remote(self.config, self.server_handles, self.reward_loop_worker_handles, self.teacher_loop_worker_handles)
+                ).remote(
+                    self.config, self.server_handles, self.reward_loop_worker_handles, self.teacher_loop_worker_handles
+                )
             )
 
     def generate_sequences(self, prompts: DataProto) -> DataProto:

verl/experimental/teacher_loop/__init__.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .teacher_loop import TeacherLoopManager, TeacherLoopWorker
-from .teacher_loop import TeacherModelManager
+from .teacher_loop import TeacherLoopManager, TeacherLoopWorker, TeacherModelManager
 
 __all__ = ["TeacherModelManager", "TeacherLoopWorker", "TeacherLoopManager"]

verl/experimental/teacher_loop/teacher_loop.py

@@ -20,21 +20,16 @@
 import numpy as np
 import ray
 import torch
-import torch.nn.functional as F
-from omegaconf import DictConfig, open_dict
+from omegaconf import DictConfig
 from tensordict import TensorDict
 
 from verl.protocol import DataProto
 from verl.single_controller.ray.base import RayResourcePool
-from verl.trainer.ppo.reward import load_reward_manager
-from verl.utils import hf_tokenizer
-from verl.utils.fs import copy_to_local
+from verl.trainer.distillation.losses import DistillationLossSettings, get_distillation_loss_settings
 from verl.utils.config import omega_conf_to_dataclass
-
-from .teacher_model import TeacherModelManager
 from verl.workers.config import DistillationConfig, DistillationLossConfig
 
-from verl.trainer.distillation.losses import get_distillation_loss_settings, DistillationLossSettings
+from .teacher_model import TeacherModelManager
 
 logger = logging.getLogger(__file__)
 logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN"))
@@ -52,7 +47,9 @@ def __init__(self, config: DictConfig, teacher_router_address: str = None):
         self.config = config
         self.distillation_config: DistillationConfig = self.config.distillation
         self.distillation_loss_config: DistillationLossConfig = self.distillation_config.distillation_loss
-        self.distillation_loss_settings: DistillationLossSettings = get_distillation_loss_settings(self.distillation_loss_config.loss_mode)
+        self.distillation_loss_settings: DistillationLossSettings = get_distillation_loss_settings(
+            self.distillation_loss_config.loss_mode
+        )
         self.teacher_router_address = teacher_router_address
         # # Serialize teacher requests per actor to reduce pressure on the teacher vLLM router/backend.
         # self._request_semaphore = asyncio.Semaphore(1)
@@ -112,16 +109,16 @@ async def _post_request(self, payload: dict, endpoint: str, max_retries: int = 1
             raise last_exception
 
     async def _compute_logprobs(self, data: DataProto) -> dict:
-        prompt_ids = data.batch['prompt_ids']
-        response_ids = data.batch['response_ids']
+        prompt_ids = data.batch["prompt_ids"]
+        response_ids = data.batch["response_ids"]
         input_ids = torch.cat([prompt_ids, response_ids], dim=1).squeeze(0).tolist()
         engine_name = self.config.distillation.teacher_model.inference.name
         model_name = self.config.distillation.teacher_model.model_path
         if engine_name == "vllm":
             if self.distillation_loss_settings.use_topk:
                 num_logprobs = topk = self.distillation_loss_config.topk
             else:
-                num_logprobs = 0 # only the sampled logprob    
+                num_logprobs = 0  # only the sampled logprob
             payloads = {
                 "model": model_name,
                 "prompt": input_ids,
@@ -140,7 +137,7 @@ async def _compute_logprobs(self, data: DataProto) -> dict:
             for logprobs_dict in response_logprob_dicts:
                 if num_logprobs == 0:
                     token_id_str = list(logprobs_dict.keys())[0]
-                    logprob = logprobs_dict[token_id_str]['logprob']
+                    logprob = logprobs_dict[token_id_str]["logprob"]
                     response_logprobs_ls.append([logprob])
                     response_ids_ls.append([int(token_id_str)])
                 else:
@@ -149,18 +146,22 @@ async def _compute_logprobs(self, data: DataProto) -> dict:
                     # We get either top-k logprobs or top-k plus the sampled logprob (if sampled token is not in top-k)
                     assert len(logprobs_dict) in [topk, topk + 1], len(logprobs_dict)
                     for token_id_str, token_dict in logprobs_dict.items():
-                        if token_dict['rank'] > topk:
-                            continue # the sampled token is not in the top-k
-                        rank = token_dict['rank']
-                        logprob = token_dict['logprob']
+                        if token_dict["rank"] > topk:
+                            continue  # the sampled token is not in the top-k
+                        rank = token_dict["rank"]
+                        logprob = token_dict["logprob"]
                         response_ids[rank - 1] = int(token_id_str)
                         response_logprobs[rank - 1] = logprob
                     response_logprobs_ls.append(response_logprobs)
                     response_ids_ls.append(response_ids)
-            logprobs_dtype = torch.bfloat16 if self.distillation_config.teacher_model.inference.dtype == "bfloat16" else torch.float32
+            logprobs_dtype = (
+                torch.bfloat16
+                if self.distillation_config.teacher_model.inference.dtype == "bfloat16"
+                else torch.float32
+            )
             response_logprobs = torch.tensor(response_logprobs_ls, dtype=logprobs_dtype).unsqueeze(0)
             response_ids = torch.tensor(response_ids_ls, dtype=torch.long).unsqueeze(0)
-            
+
         elif engine_name == "sglang":
             raise ValueError("SGLang backend does not support distillation currently.")
             payloads = {
@@ -200,7 +201,9 @@ class TeacherLoopManager:
 
     def __init__(self, config: DictConfig, teacher_resource_pool: RayResourcePool = None):
         self.config = config
-        self.distillation_config: DistillationConfig = omega_conf_to_dataclass(self.config.distillation) # to dataclass for the post init to handle top-k and engine kwargs
+        self.distillation_config: DistillationConfig = omega_conf_to_dataclass(
+            self.config.distillation
+        )  # to dataclass for the post init to handle top-k and engine kwargs
         self.teacher_model_manager = TeacherModelManager(self.distillation_config.teacher_model, teacher_resource_pool)
         self.teacher_router_address = self.teacher_model_manager.get_router_address()
 
@@ -266,6 +269,7 @@ def compute_teacher_logprobs(self, data: DataProto) -> DataProto:
 
     def _run_all(self, tasks: list[asyncio.Task]):
         raise NotImplementedError("TODO:RM")
+
         async def run_all():
             return await asyncio.gather(*tasks)
 

verl/experimental/teacher_loop/teacher_model.py

@@ -17,7 +17,7 @@
 import os
 
 from verl.single_controller.ray.base import RayResourcePool, split_resource_pool
-from verl.workers.config import HFModelConfig, DistillationTeacherModelConfig
+from verl.workers.config import DistillationTeacherModelConfig, HFModelConfig
 from verl.workers.rollout.replica import get_rollout_replica_class
 
 logger = logging.getLogger(__file__)
@@ -91,6 +91,7 @@ def _initialize_router(self):
         worker_urls = [f"http://{server_address}" for server_address in self.server_addresses]
 
         from ..reward_loop.router.naive_router import launch_router_process
+
         self.router_address, _ = launch_router_process(worker_urls=worker_urls)
 
     def get_router_address(self):

verl/trainer/distillation/fsdp/losses.py

@@ -27,6 +27,7 @@ def kl_divergence(log_q: torch.Tensor, log_p: torch.Tensor) -> torch.Tensor:
     kld = p * (log_p - log_q)
     return kld.sum(dim=-1)
 
+
 def compute_forward_kl_topk(
     student_logits: torch.Tensor,
     teacher_topk_log_probs: torch.Tensor,

verl/trainer/distillation/losses.py

@@ -22,7 +22,7 @@
 from verl.trainer.distillation.types import DistillationLossInputs
 from verl.trainer.ppo.core_algos import agg_loss, kl_penalty
 from verl.utils.metric import AggregationType, Metric
-from verl.workers.config import DistillationConfig, DistillationLossConfig, ActorConfig
+from verl.workers.config import ActorConfig, DistillationConfig, DistillationLossConfig
 
 DistillationLossFn = Callable[
     [
@@ -56,8 +56,7 @@ def __post_init__(self):
         self.names = [self.names] if isinstance(self.names, str) else self.names
         if sum([self.use_topk, self.use_estimator]) > 1:
             raise ValueError(
-                f"Expected only one of use_estimator, use_topk, but got "
-                f"{self.use_estimator=}, {self.use_topk=}."
+                f"Expected only one of use_estimator, use_topk, but got {self.use_estimator=}, {self.use_topk=}."
             )
 
 
@@ -206,7 +205,11 @@ def compute_forward_kl_topk(
         teacher_topk_ids=teacher_topk_ids,
         config=distillation_config,
     )
-    distillation_losses, student_mass, teacher_mass = outputs["distillation_losses"], outputs["student_mass"], outputs["teacher_mass"]
+    distillation_losses, student_mass, teacher_mass = (
+        outputs["distillation_losses"],
+        outputs["student_mass"],
+        outputs["teacher_mass"],
+    )
 
     # Log amount of mass in the top-k log probabilities for both student and teacher.
     student_mass = student_mass[response_mask]

verl/trainer/main_ppo.py

@@ -242,7 +242,9 @@ def init_resource_pool_mgr(self, config):
             if distillation_config.teacher_model.nnodes <= 0:
                 raise ValueError("config.distillation.teacher_model.nnodes must be greater than 0")
 
-            teacher_pool = [distillation_config.teacher_model.n_gpus_per_node] * distillation_config.teacher_model.nnodes
+            teacher_pool = [
+                distillation_config.teacher_model.n_gpus_per_node
+            ] * distillation_config.teacher_model.nnodes
             resource_pool_spec["teacher_pool"] = teacher_pool
 
         from verl.trainer.ppo.ray_trainer import ResourcePoolManager
@@ -274,7 +276,6 @@ def add_teacher_model_resource_pool(self, config):
             else:
                 self.mapping[Role.TeacherModel] = "global_pool"
 
-
     def add_ref_policy_worker(self, config, ref_policy_cls):
         """Add reference policy worker if KL loss or KL reward is used."""
         from verl.trainer.ppo.ray_trainer import Role
@@ -462,4 +463,4 @@ def create_rl_sampler(data_config, dataset):
 
 
 if __name__ == "__main__":
-    main()
+    main()

verl/trainer/ppo/ray_trainer.py

@@ -38,7 +38,7 @@
 from verl.experimental.dataset.sampler import AbstractCurriculumSampler
 from verl.protocol import pad_dataproto_to_divisor, unpad_dataproto
 from verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup, ResourcePoolManager
-from verl.single_controller.ray.base import create_colocated_worker_cls, split_resource_pool
+from verl.single_controller.ray.base import create_colocated_worker_cls
 from verl.trainer.config import AlgoConfig
 from verl.trainer.ppo import core_algos
 from verl.trainer.ppo.core_algos import AdvantageEstimator, agg_loss
@@ -54,9 +54,9 @@
     Role,
     WorkerType,
     need_critic,
-    need_teacher_policy,
     need_reference_policy,
     need_reward_model,
+    need_teacher_policy,
 )
 from verl.utils import tensordict_utils as tu
 from verl.utils.checkpoint.checkpoint_manager import find_latest_ckpt_path, should_save_ckpt_esi
@@ -744,7 +744,6 @@ def init_workers(self):
             )
             self.resource_pool_to_cls[resource_pool][str(Role.RefPolicy)] = ref_policy_cls
 
-
         # initialize WorkerGroup
         # NOTE: if you want to use a different resource pool for each role, which can support different parallel size,
         # you should not use `create_colocated_worker_cls`.
@@ -1150,7 +1149,6 @@ def _compute_ref_log_prob(self, batch: DataProto) -> DataProto:
 
         return ref_log_prob
 
-
     def _compute_old_log_prob(self, batch: DataProto):
         if self.use_legacy_worker_impl == "disable":
             # TODO: remove step 1, 2, 4 after we make the whole training tensordict and padding free

verl/utils/config.py

@@ -199,4 +199,4 @@ def check_mutually_exclusive(mbs, mbs_per_gpu, name: str):
 
         get_vllm_max_lora_rank(lora_rank)
 
-    print("[validate_config] All configuration checks passed successfully!")
+    print("[validate_config] All configuration checks passed successfully!")

verl/utils/stages.py

@@ -20,4 +20,4 @@ class Stage(Enum):
 
     OLD_LOG_PROB = "old_log_prob"
     REF_LOG_PROB = "ref_log_prob"
-    ACTOR_UPDATE = "actor_update"
+    ACTOR_UPDATE = "actor_update"