lints

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

Author: ananthsub

examples/llm/llama3_8b/conf/llama3_8b_pretrain_override_example.yaml

@@ -56,4 +56,4 @@ profiling_config:
   profile_step_start: 1000
   profile_step_end: 10000
   use_pytorch_profiler: true
-  profile_ranks: [0, 1]
+  profile_ranks: [0, 1]

examples/llm/llama3_8b/pretrain_llama3_8b.py

@@ -15,12 +15,12 @@
 
 import argparse
 import dataclasses
+import logging
 import os
 import sys
-import logging
 from typing import Any, Dict, TypeVar
 
-from omegaconf import OmegaConf, DictConfig
+from omegaconf import DictConfig, OmegaConf
 
 from nemo_lm.models.utils import forward_step
 from nemo_lm.recipes.llm.llama3_8b import pretrain_config
@@ -60,6 +60,7 @@ def apply_overrides_recursively(config_obj: DataclassInstance, overrides_dict: D
 
 
 def main() -> None:
+    """Entry point for the script."""
     parser = argparse.ArgumentParser(description="Pretrain Llama3 8B model using NeMo-LM")
     parser.add_argument(
         "--config-file",
@@ -74,7 +75,6 @@ def main() -> None:
     logger.info("------------------------------------")
 
     # 1. Load base configuration from the recipe - this is our primary Python object
-    logger.info(f"Loading base configuration using 'pretrain_config' from nemo_lm.recipes.llm.llama3_8b...")
     final_cfg_container: ConfigContainer = pretrain_config()
 
     # 2. Load YAML overrides into an OmegaConf dictionary

nemo_lm/recipes/llm/llama3_8b.py

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 from typing import Optional
 
-import os
 import torch
 
 from megatron.core.distributed import DistributedDataParallelConfig
@@ -24,7 +24,6 @@
 from nemo_lm.training.config import (
     CheckpointConfig,
     ConfigContainer,
-    FinetuningDatasetConfig,
     GPTDatasetConfig,
     LoggerConfig,
     RNGConfig,