Renames for consistency (#14)

ananthsub · web-flow · commit 117a22a24d4a · 2025-05-29T13:15:40.000-05:00
* Renames for consistency

Signed-off-by: Ananth Subramaniam &lt;ansubramania@nvidia.com&gt;

* cleanup

Signed-off-by: Ananth Subramaniam &lt;ansubramania@nvidia.com&gt;

* cleanup

Signed-off-by: Ananth Subramaniam &lt;ansubramania@nvidia.com&gt;

---------

Signed-off-by: Ananth Subramaniam &lt;ansubramania@nvidia.com&gt;
diff --git a/nemo_lm/converters/__init__.py b/nemo_lm/converters/__init__.py
diff --git a/nemo_lm/converters/common.py b/nemo_lm/converters/common.py
@@ -28,12 +28,12 @@
 from megatron.core.optimizer import OptimizerConfig
 from megatron.core.transformer.module import MegatronModule
 
-from nemo_lm.checkpointing import save_checkpoint
-from nemo_lm.config import CheckpointConfig, ConfigContainer, LoggerConfig, TokenizerConfig
-from nemo_lm.model.gpt import GPTConfig, torch_dtype_from_mcore_config
-from nemo_lm.model.t5 import T5Config
-from nemo_lm.state import GlobalState
+from nemo_lm.models.gpt import GPTConfig, torch_dtype_from_mcore_config
+from nemo_lm.models.t5 import T5Config
 from nemo_lm.tokenizers.tokenizer import _HuggingFaceTokenizer
+from nemo_lm.training.checkpointing import save_checkpoint
+from nemo_lm.training.config import CheckpointConfig, ConfigContainer, LoggerConfig, TokenizerConfig
+from nemo_lm.training.state import GlobalState
 from nemo_lm.utils.instantiate_utils import instantiate
 
 if TYPE_CHECKING:
diff --git a/nemo_lm/converters/llama.py b/nemo_lm/converters/llama.py
@@ -18,9 +18,9 @@
 
 import torch
 
-from nemo_lm.converter.common import BaseExporter, BaseImporter, dtype_from_hf
-from nemo_lm.converter.state_transform import TransformFns, apply_transforms, state_transform
-from nemo_lm.model.llama import Llama4Config, Llama31Config, LlamaConfig
+from nemo_lm.converters.common import BaseExporter, BaseImporter, dtype_from_hf
+from nemo_lm.converters.state_transform import TransformFns, apply_transforms, state_transform
+from nemo_lm.models.llama import Llama4Config, Llama31Config, LlamaConfig
 
 if TYPE_CHECKING:
     from transformers import LlamaConfig as HFLlamaConfig
diff --git a/nemo_lm/converters/state_transform.py b/nemo_lm/converters/state_transform.py
diff --git a/nemo_lm/data/builders/hf_dataset.py b/nemo_lm/data/builders/hf_dataset.py
@@ -24,10 +24,10 @@
 from datasets import Dataset, DatasetDict, load_dataset
 from tqdm import tqdm
 
-from nemo_lm.config import FinetuningDatasetConfig
 from nemo_lm.data.builders.finetuning_dataset import FinetuningDatasetBuilder
 from nemo_lm.data.datasets.sft import get_dataset_root
 from nemo_lm.tokenizers.tokenizer import MegatronTokenizer
+from nemo_lm.training.config import FinetuningDatasetConfig
 from nemo_lm.utils.common_utils import print_rank_0
 
 logger = logging.getLogger(__name__)
diff --git a/nemo_lm/data/loaders.py b/nemo_lm/data/loaders.py
@@ -21,9 +21,9 @@
 from megatron.core.rerun_state_machine import RerunDataIterator
 from torch.utils.data import DataLoader
 
-from nemo_lm.config import ConfigContainer
 from nemo_lm.data.samplers import build_pretraining_data_loader
-from nemo_lm.state import TrainState
+from nemo_lm.training.config import ConfigContainer
+from nemo_lm.training.state import TrainState
 from nemo_lm.utils.common_utils import print_rank_0
 from nemo_lm.utils.sig_utils import DistributedSignalHandler
 
diff --git a/nemo_lm/data/utils.py b/nemo_lm/data/utils.py
@@ -20,10 +20,10 @@
 from megatron.core.datasets.blended_megatron_dataset_config import BlendedMegatronDatasetConfig
 from megatron.core.datasets.gpt_dataset import GPTDataset, MockGPTDataset
 
-from nemo_lm.config import DataloaderConfig, FinetuningDatasetConfig, GPTDatasetConfig
-from nemo_lm.data.builders.finetuning_dataset import FinetuningDatasetBuilder
+from nemo_lm.data.builders import FinetuningDatasetBuilder
 from nemo_lm.data.builders.hf_dataset import HFDatasetBuilder, HFDatasetConfig
 from nemo_lm.tokenizers.tokenizer import MegatronTokenizer
+from nemo_lm.training.config import DataloaderConfig, FinetuningDatasetConfig, GPTDatasetConfig
 from nemo_lm.utils.common_utils import print_rank_0
 
 
diff --git a/nemo_lm/models/__init__.py b/nemo_lm/models/__init__.py
@@ -23,8 +23,8 @@
 from megatron.core.fp8_utils import is_float8tensor
 from megatron.core.transformer.module import Float16Module, MegatronModule
 
-from nemo_lm.model.gpt import GPTConfig
-from nemo_lm.model.t5 import T5Config
+from nemo_lm.models.gpt import GPTConfig
+from nemo_lm.models.t5 import T5Config
 
 
 def get_model_from_config(
diff --git a/nemo_lm/models/gpt.py b/nemo_lm/models/gpt.py
diff --git a/nemo_lm/models/llama.py b/nemo_lm/models/llama.py
@@ -20,8 +20,8 @@
 import torch
 import torch.nn.functional as F
 
-from nemo_lm.model.gpt import GPTConfig
-from nemo_lm.model.llama4_utils import get_llama4_layer_spec
+from nemo_lm.models.gpt import GPTConfig
+from nemo_lm.models.llama4_utils import get_llama4_layer_spec
 
 try:
     from megatron.core.transformer.spec_utils import ModuleSpec
diff --git a/nemo_lm/models/llama4_utils.py b/nemo_lm/models/llama4_utils.py
diff --git a/nemo_lm/models/t5.py b/nemo_lm/models/t5.py
@@ -21,7 +21,7 @@
 from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_config import TransformerConfig
 
-from nemo_lm.model.gpt import get_vocab_size
+from nemo_lm.models.gpt import get_vocab_size
 from nemo_lm.utils.import_utils import safe_import
 
 _, HAVE_TE = safe_import("transformer_engine")
diff --git a/nemo_lm/models/utils.py b/nemo_lm/models/utils.py
@@ -21,9 +21,9 @@
 from megatron.core.models.gpt import GPTModel
 from megatron.core.utils import get_batch_on_this_cp_rank
 
-from nemo_lm.config import ConfigContainer, FinetuningDatasetConfig
-from nemo_lm.losses import masked_next_token_loss
-from nemo_lm.state import GlobalState
+from nemo_lm.training.config import ConfigContainer, FinetuningDatasetConfig
+from nemo_lm.training.losses import masked_next_token_loss
+from nemo_lm.training.state import GlobalState
 
 logger = logging.getLogger(__name__)
 
diff --git a/nemo_lm/tokenizers/tokenizer.py b/nemo_lm/tokenizers/tokenizer.py
@@ -10,10 +10,10 @@
 
 from megatron.core.datasets.megatron_tokenizer import MegatronTokenizer as MegatronTokenizerCore
 
-from nemo_lm.config import TokenizerConfig
 from nemo_lm.tokenizers.bert_tokenization import FullTokenizer as FullBertTokenizer
 from nemo_lm.tokenizers.gpt2_tokenization import GPT2Tokenizer
 from nemo_lm.tokenizers.multimodal_tokenizer import MultimodalTokenizer
+from nemo_lm.training.config import TokenizerConfig
 from nemo_lm.utils.common_utils import get_rank_safe, print_rank_0
 
 
@@ -684,7 +684,6 @@ class _GPTSentencePieceTokenizer(_SentencePieceTokenizer):
     token handling. It primarily uses the BOS, EOS, and PAD IDs defined by the
     SentencePiece model itself, without adding extra tokens like <CLS>, <SEP>, etc.
     The `eod` (end-of-document) token is mapped to the `eos_id`.
-
     Args:
         model_file (str): Path to the SentencePiece model file.
     """
@@ -735,11 +734,9 @@ def additional_special_tokens_ids(self):
 
 class _Llama2Tokenizer(_SentencePieceTokenizer):
     """A tokenizer specifically for Llama-2 style models, using SentencePiece.
-
     This class inherits from `_SentencePieceTokenizer` and is configured for Llama-2's
     specific use of BOS and EOS tokens. It uses the BOS/EOS/PAD IDs directly from
     the SentencePiece model.
-
     Args:
         model_file (str): Path to the SentencePiece model file for Llama-2.
     """
@@ -755,12 +752,10 @@ def __init__(self, model_file):
 
     def tokenize(self, s: str, bos=True, eos=False):
         """Tokenizes a string, with options to add BOS and EOS tokens.
-
         Args:
             s (str): The input string to tokenize.
             bos (bool, optional): Whether to prepend the BOS token. Defaults to True.
             eos (bool, optional): Whether to append the EOS token. Defaults to False.
-
         Returns:
             list[int]: A list of token IDs.
         """
@@ -806,15 +801,12 @@ def reload_mergeable_ranks(path: str, max_vocab: Optional[int] = None) -> Dict[b
     """
     Reloads a tokenizer vocabulary from a JSON file (NeMo format) and converts it
     into the mergeable ranks format required by Tiktoken.
-
     The input JSON file is expected to be a list of dictionaries, each with
     "rank", "token_bytes" (base64 encoded), and "token_str" keys.
-
     Args:
         path (str): Path to the JSON vocabulary file.
         max_vocab (Optional[int], optional): If provided, truncates the vocabulary
                                            to this maximum size. Defaults to None.
-
     Returns:
         Dict[bytes, int]: A dictionary mapping token bytes to their ranks.
     """
@@ -851,12 +843,10 @@ def reload_mergeable_ranks(path: str, max_vocab: Optional[int] = None) -> Dict[b
 
 class CustomTikTokenizer(MegatronTokenizer):
     """A custom tokenizer using the Tiktoken library with a NeMo-style vocabulary file.
-
     This tokenizer loads a vocabulary from a JSON file (processed by
     `reload_mergeable_ranks`) and uses it with Tiktoken for encoding and decoding.
     It supports a configurable number of special tokens, which are placed at the
     beginning of the vocabulary ID space.
-
     Args:
         path (str): Path to the JSON vocabulary file (NeMo format).
         pattern (str): The regex pattern string for Tiktoken.
@@ -962,12 +952,10 @@ def inv_vocab(self):
 
     def tokenize(self, s: str, bos: bool = False, eos: bool = False) -> List[int]:
         """Tokenizes a string, with options to add BOS and EOS tokens.
-
         Args:
             s (str): The input string to tokenize.
             bos (bool, optional): Whether to prepend the BOS token. Defaults to False.
             eos (bool, optional): Whether to append the EOS token. Defaults to False.
-
         Returns:
             List[int]: A list of token IDs.
         """
@@ -1005,11 +993,9 @@ def decoder(self):
 
 class _NullTokenizer(MegatronTokenizer):
     """A simple tokenizer that splits text by spaces and converts tokens to integers.
-
     This tokenizer is primarily for testing or placeholder purposes where actual
     linguistic tokenization is not required. It assumes tokens are space-separated
     integers.
-
     Args:
         vocab_size (int): The vocabulary size, excluding the EOD token.
                           The EOD token will be assigned `vocab_size` as its ID.
diff --git a/nemo_lm/training/checkpointing.py b/nemo_lm/training/checkpointing.py
@@ -44,9 +44,9 @@
 from megatron.core.num_microbatches_calculator import update_num_microbatches
 from megatron.core.rerun_state_machine import get_rerun_state_machine
 
-from nemo_lm import fault_tolerance
-from nemo_lm.config import ConfigContainer
-from nemo_lm.state import GlobalState, TrainState
+from nemo_lm.training import fault_tolerance
+from nemo_lm.training.config import ConfigContainer
+from nemo_lm.training.state import GlobalState, TrainState
 from nemo_lm.utils import wandb_utils
 from nemo_lm.utils.async_utils import is_empty_async_queue, schedule_async_save
 from nemo_lm.utils.common_utils import (
diff --git a/nemo_lm/training/config.py b/nemo_lm/training/config.py
@@ -24,8 +24,8 @@
 from megatron.core.distributed import DistributedDataParallelConfig
 from megatron.core.optimizer import OptimizerConfig
 
-from nemo_lm.model.gpt import GPTConfig
-from nemo_lm.model.t5 import T5Config
+from nemo_lm.models.gpt import GPTConfig
+from nemo_lm.models.t5 import T5Config
 from nemo_lm.utils.common_utils import get_world_size_safe
 from nemo_lm.utils.config_utils import ConfigContainer as Container
 
diff --git a/nemo_lm/training/eval.py b/nemo_lm/training/eval.py
@@ -23,9 +23,9 @@
 from megatron.core.rerun_state_machine import RerunDataIterator, RerunMode, get_rerun_state_machine
 from megatron.core.transformer import MegatronModule
 
-from nemo_lm import fault_tolerance
-from nemo_lm.config import ConfigContainer
-from nemo_lm.state import GlobalState
+from nemo_lm.training import fault_tolerance
+from nemo_lm.training.config import ConfigContainer
+from nemo_lm.training.state import GlobalState
 from nemo_lm.utils.common_utils import is_last_rank, print_rank_0, print_rank_last
 from nemo_lm.utils.train_utils import check_forward_step_func_num_args, maybe_inject_state
 
diff --git a/nemo_lm/training/fault_tolerance.py b/nemo_lm/training/fault_tolerance.py
@@ -56,8 +56,8 @@
 
 import torch
 
-from nemo_lm.config import ConfigContainer, FaultToleranceConfig
-from nemo_lm.state import GlobalState
+from nemo_lm.training.config import ConfigContainer, FaultToleranceConfig
+from nemo_lm.training.state import GlobalState
 from nemo_lm.utils.common_utils import get_rank_safe, print_rank_0
 
 _NUM_WARMUP_ITERS: int = 1
diff --git a/nemo_lm/training/initialize.py b/nemo_lm/training/initialize.py
@@ -30,9 +30,9 @@
 )
 from megatron.core.utils import get_te_version, is_te_min_version, is_torch_min_version
 
-from nemo_lm.config import ConfigContainer, DistributedInitConfig, RerunStateMachineConfig, RNGConfig
-from nemo_lm.model.gpt import GPTConfig
-from nemo_lm.model.t5 import T5Config
+from nemo_lm.models.gpt import GPTConfig
+from nemo_lm.models.t5 import T5Config
+from nemo_lm.training.config import ConfigContainer, DistributedInitConfig, RerunStateMachineConfig, RNGConfig
 from nemo_lm.utils.common_utils import get_local_rank_preinit, get_rank_safe, get_world_size_safe
 
 
diff --git a/nemo_lm/training/losses.py b/nemo_lm/training/losses.py
diff --git a/nemo_lm/training/optim.py b/nemo_lm/training/optim.py
@@ -19,7 +19,7 @@
 from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler
 from megatron.core.transformer.module import MegatronModule
 
-from nemo_lm.config import SchedulerConfig
+from nemo_lm.training.config import SchedulerConfig
 
 
 def setup_optimizer(
diff --git a/nemo_lm/training/pretrain.py b/nemo_lm/training/pretrain.py
@@ -14,12 +14,12 @@
 
 from typing import Callable, Optional
 
-from nemo_lm.checkpointing import save_checkpoint
-from nemo_lm.config import ConfigContainer
 from nemo_lm.data.utils import get_dataset_provider
-from nemo_lm.eval import evaluate_and_print_results
-from nemo_lm.setup import setup
-from nemo_lm.train import _finish_train, train
+from nemo_lm.training.checkpointing import save_checkpoint
+from nemo_lm.training.config import ConfigContainer
+from nemo_lm.training.eval import evaluate_and_print_results
+from nemo_lm.training.setup import setup
+from nemo_lm.training.train import _finish_train, train
 from nemo_lm.utils.common_utils import print_rank_0
 from nemo_lm.utils.decorators import experimental_fn
 from nemo_lm.utils.log_utils import barrier_and_log
diff --git a/nemo_lm/training/setup.py b/nemo_lm/training/setup.py
@@ -24,17 +24,17 @@
 from megatron.core.rerun_state_machine import RerunDataIterator
 from megatron.core.transformer import MegatronModule
 
-from nemo_lm import fault_tolerance
-from nemo_lm.checkpointing import checkpoint_exists, load_checkpoint
-from nemo_lm.config import CheckpointConfig, ConfigContainer
 from nemo_lm.data.loaders import setup_data_iterators
-from nemo_lm.init import initialize_megatron, set_jit_fusion_options
-from nemo_lm.model import get_model_from_config
-from nemo_lm.model.gpt import GPTConfig
-from nemo_lm.model.t5 import T5Config
-from nemo_lm.optim import setup_optimizer
-from nemo_lm.state import GlobalState
+from nemo_lm.models import get_model_from_config
+from nemo_lm.models.gpt import GPTConfig
+from nemo_lm.models.t5 import T5Config
 from nemo_lm.tokenizers.tokenizer import build_tokenizer
+from nemo_lm.training import fault_tolerance
+from nemo_lm.training.checkpointing import checkpoint_exists, load_checkpoint
+from nemo_lm.training.config import CheckpointConfig, ConfigContainer
+from nemo_lm.training.initialize import initialize_megatron, set_jit_fusion_options
+from nemo_lm.training.optim import setup_optimizer
+from nemo_lm.training.state import GlobalState
 from nemo_lm.utils.common_utils import print_rank_0
 from nemo_lm.utils.import_utils import safe_import
 from nemo_lm.utils.log_utils import append_to_progress_log, barrier_and_log, setup_logging
diff --git a/nemo_lm/training/state.py b/nemo_lm/training/state.py
@@ -25,8 +25,8 @@
 from torch.distributed.checkpoint.stateful import Stateful
 from torch.utils.tensorboard.writer import SummaryWriter
 
-from nemo_lm.config import ConfigContainer
 from nemo_lm.tokenizers.tokenizer import build_tokenizer
+from nemo_lm.training.config import ConfigContainer
 from nemo_lm.utils.common_utils import get_rank_safe, get_world_size_safe
 from nemo_lm.utils.sig_utils import DistributedSignalHandler
 from nemo_lm.utils.yaml_utils import dump_dataclass_to_yaml
diff --git a/nemo_lm/training/train.py b/nemo_lm/training/train.py
@@ -36,12 +36,12 @@
 from megatron.core.transformer import MegatronModule
 from megatron.core.utils import check_param_hashes_across_dp_replicas, get_model_config
 
-from nemo_lm import fault_tolerance
-from nemo_lm.checkpointing import save_checkpoint
-from nemo_lm.config import ConfigContainer
-from nemo_lm.eval import evaluate_and_print_results
-from nemo_lm.init import destroy_global_state
-from nemo_lm.state import GlobalState
+from nemo_lm.training import fault_tolerance
+from nemo_lm.training.checkpointing import save_checkpoint
+from nemo_lm.training.config import ConfigContainer
+from nemo_lm.training.eval import evaluate_and_print_results
+from nemo_lm.training.initialize import destroy_global_state
+from nemo_lm.training.state import GlobalState
 from nemo_lm.utils import flop_utils
 from nemo_lm.utils.async_utils import maybe_finalize_async_save
 from nemo_lm.utils.common_utils import get_world_size_safe, print_rank_0
diff --git a/nemo_lm/utils/async_utils.py b/nemo_lm/utils/async_utils.py
@@ -19,7 +19,7 @@
 
 from megatron.core.dist_checkpointing.strategies.async_utils import AsyncCallsQueue, AsyncRequest
 
-from nemo_lm.config import CheckpointConfig
+from nemo_lm.training.config import CheckpointConfig
 from nemo_lm.utils.common_utils import print_rank_0
 
 # Singleton manager of async calls
diff --git a/nemo_lm/utils/flop_utils.py b/nemo_lm/utils/flop_utils.py
@@ -14,7 +14,7 @@
 
 import torch.nn.functional as F
 
-from nemo_lm.config import ConfigContainer
+from nemo_lm.training.config import ConfigContainer
 
 
 def num_floating_point_operations(cfg: ConfigContainer, batch_size: int) -> float:
diff --git a/nemo_lm/utils/theoretical_memory_utils.py b/nemo_lm/utils/theoretical_memory_utils.py
@@ -19,7 +19,7 @@
 
 import torch.nn.functional as F
 
-from nemo_lm.config import ConfigContainer
+from nemo_lm.training.config import ConfigContainer
 
 NUM_BYTES_IN_MEGABYTE: int = 1024 * 1024
 
diff --git a/nemo_lm/utils/train_utils.py b/nemo_lm/utils/train_utils.py
@@ -25,8 +25,8 @@
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.utils import get_data_parallel_group_if_dtensor, to_local_if_dtensor
 
-from nemo_lm.config import ConfigContainer
-from nemo_lm.state import GlobalState
+from nemo_lm.training.config import ConfigContainer
+from nemo_lm.training.state import GlobalState
 from nemo_lm.utils.common_utils import get_world_size_safe, is_last_rank, print_rank_last
 from nemo_lm.utils.theoretical_memory_utils import report_theoretical_memory
 
