Ran linter fixes

Author: root

examples/nemo_gym/run_grpo_nemo_gym.py

@@ -23,7 +23,6 @@
 
 import ray
 from omegaconf import OmegaConf
-from wandb import Table
 
 from nemo_rl.algorithms.grpo import (
     ColocatablePolicyInterface,
@@ -54,6 +53,7 @@
     register_omegaconf_resolvers,
 )
 from nemo_rl.utils.logger import get_next_experiment_dir
+from wandb import Table
 
 
 def parse_args() -> tuple[argparse.Namespace, list[str]]:

nemo_rl/experience/rollouts.py

@@ -27,7 +27,6 @@
 import ray
 import torch
 from transformers import PreTrainedTokenizerBase
-from wandb import Histogram, Table
 
 from nemo_rl.data.interfaces import (
     DatumSpec,
@@ -50,6 +49,7 @@
     GenerationOutputSpec,
 )
 from nemo_rl.utils.timer import Timer
+from wandb import Histogram, Table
 
 TokenizerType = PreTrainedTokenizerBase
 

nemo_rl/models/megatron/setup.py

@@ -705,9 +705,13 @@ def freeze_moe_router(megatron_model):
     if use_peft:
         peft_cfg = policy_cfg["megatron_cfg"].get("peft", {})
         if "dim" not in peft_cfg or peft_cfg["dim"] is None:
-            raise ValueError("If megtatron_cfg.peft.enabled is True, dim must be set in peft_cfg")
+            raise ValueError(
+                "If megtatron_cfg.peft.enabled is True, dim must be set in peft_cfg"
+            )
         if "alpha" not in peft_cfg or peft_cfg["alpha"] is None:
-            raise ValueError("If megtatron_cfg.peft.enabled is True, alpha must be set in peft_cfg")
+            raise ValueError(
+                "If megtatron_cfg.peft.enabled is True, alpha must be set in peft_cfg"
+            )
         peft = LoRA(
             target_modules=peft_cfg.get("target_modules", []),
             exclude_modules=peft_cfg.get("exclude_modules", []),
@@ -875,13 +879,17 @@ def setup_reference_model_state(
 
     ref_pre_wrap_hooks = []
     use_peft = config["megatron_cfg"].get("peft", {}).get("enabled", False)
-    
+
     if use_peft:
         peft_cfg = config["megatron_cfg"].get("peft", {})
         if "dim" not in peft_cfg or peft_cfg["dim"] is None:
-            raise ValueError("If megtatron_cfg.peft.enabled is True, dim must be set in peft_cfg")
+            raise ValueError(
+                "If megtatron_cfg.peft.enabled is True, dim must be set in peft_cfg"
+            )
         if "alpha" not in peft_cfg or peft_cfg["alpha"] is None:
-            raise ValueError("If megtatron_cfg.peft.enabled is True, alpha must be set in peft_cfg")
+            raise ValueError(
+                "If megtatron_cfg.peft.enabled is True, alpha must be set in peft_cfg"
+            )
         peft = LoRA(
             target_modules=peft_cfg.get("target_modules", []),
             exclude_modules=peft_cfg.get("exclude_modules", []),
@@ -931,7 +939,7 @@ def composed_peft_hook(model: list[MegatronModule]) -> list[MegatronModule]:
         ref_megatron_cfg.checkpoint.finetune = False
 
     print("Loading the Reference Model")
-    
+
     if should_load_checkpoint:
         load_checkpoint(
             ref_state,

nemo_rl/utils/logger.py

@@ -30,7 +30,6 @@
 import requests
 import swanlab
 import torch
-import wandb
 from matplotlib import pyplot as plt
 from prometheus_client.parser import text_string_to_metric_families
 from prometheus_client.samples import Sample
@@ -40,6 +39,7 @@
 from rich.panel import Panel
 from torch.utils.tensorboard import SummaryWriter
 
+import wandb
 from nemo_rl.data.interfaces import LLMMessageLogType
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict