precommit

megatron/logging.py

@@ -24,6 +24,7 @@
 from megatron import mpu, print_rank_0
 from megatron.utils import report_memory
 
+
 class Tee:
     """Duplicate output to both stdout/err and file"""
 

megatron/model/transformer.py

@@ -230,7 +230,7 @@ def __init__(
                 #     skip_bias_add=False,
                 #     mup_rescale_parameters=is_last_layer,  # only called if neox_args.use_mup = True, despite it not being included here
                 # )
-            else: # Not using cross entropy loss for RMs
+            else:  # Not using cross entropy loss for RMs
                 self.rm_linear = mpu.RowParallelLinear(
                     neox_args=neox_args,
                     input_size=neox_args.hidden_size,

megatron/neox_arguments/arguments.py

@@ -51,15 +51,15 @@
 )
 
 ### Logging colors ###
-GREEN = '\033[92m'
-RED = '\033[91m'
-YELLOW = '\033[93m'
-END = '\033[0m'
+GREEN = "\033[92m"
+RED = "\033[91m"
+YELLOW = "\033[93m"
+END = "\033[0m"
 SUCCESS = f"{GREEN} [SUCCESS] {END}"
 OKAY = f"{GREEN}[OKAY]{END}"
 WARNING = f"{YELLOW}[WARNING]{END}"
-FAIL = f'{RED}[FAIL]{END}'
-INFO = '[INFO]'
+FAIL = f"{RED}[FAIL]{END}"
+INFO = "[INFO]"
 
 # ZERO defaults by deespeed
 # These values should not be changed unless defaults in deepspeed are changed
@@ -205,11 +205,13 @@ def initialize_comet(self):
 
                 logging.info("> setting up comet ...")
             except ImportError as e:
-                logging.error(f"{FAIL} importing comet. Comet can be installed with \"pip install comet_llm\". See https://github.com/comet-ml/comet-llm for more info. Full error is:")
+                logging.error(
+                    f'{FAIL} importing comet. Comet can be installed with "pip install comet_llm". See https://github.com/comet-ml/comet-llm for more info. Full error is:'
+                )
                 raise e
             except Exception as e:
                 logging.error(
-                        f"{FAIL} Error setting up Comet. Either set \"use_comet: False\" in your configuration file, or resolve the issue with Comet. Full error is:",
+                    f'{FAIL} Error setting up Comet. Either set "use_comet: False" in your configuration file, or resolve the issue with Comet. Full error is:',
                 )
                 raise e
 
@@ -1219,7 +1221,9 @@ def validate_values(self):
 
         # learning rate
         if self.lr is None:
-            error_message = f"{FAIL} " + self.__class__.__name__ + ".validate_values() lr is None"
+            error_message = (
+                f"{FAIL} " + self.__class__.__name__ + ".validate_values() lr is None"
+            )
             logging.error(error_message)
             raise ValueError(error_message)
             return False

tools/datasets/preprocess_data_with_chat_template.py

@@ -122,7 +122,7 @@ def build_chat(
         chat_tokens = tokenizer.apply_chat_template(
             chat[: i + 1], add_generation_prompt=add_gen
         )[len(tokens) :]
-  
+
         # remove previous stuff...
         tokens.extend(chat_tokens)
         if only_last_turn and (i != len(chat) - 1):