[refactor] split JobConfig and ConfigManager into two files (#1442)

This PR creates a new folder `torchtitan/config` to host `job_config.py`
and `manager.py`, for the reasons below:
- Both are complicated enough to worth their own files.
- The convention in torchtitan to extend custom `JobConfig` is to create
a file under model folder called `job_config.py` (see
https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/flux/job_config.py).
This PR makes the origin `JobConfig` consistent with that convention.
- (minor) Creating a more succinct `torchtitan.config` namespace is more
readable than importing from `torchtitan.config_manager`.

Author: tianyu-l

.github/workflows/integration_test_8gpu.yaml

@@ -46,4 +46,4 @@ jobs:
         USE_CPP=0 python -m pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu126
 
         mkdir artifacts-to-be-uploaded
-        python ./tests/integration_tests.py artifacts-to-be-uploaded --ngpu 8
+        python -m tests.integration_tests artifacts-to-be-uploaded --ngpu 8

.github/workflows/integration_test_8gpu_h100.yaml

@@ -47,4 +47,4 @@ jobs:
         USE_CPP=0 python -m pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu126
 
         mkdir artifacts-to-be-uploaded
-        python ./tests/integration_tests_h100.py artifacts-to-be-uploaded --ngpu 8
+        python -m tests.integration_tests_h100 artifacts-to-be-uploaded --ngpu 8

.github/workflows/integration_test_8gpu_torchft.yaml

@@ -49,5 +49,5 @@ jobs:
         RUST_BACKTRACE=1 torchft_lighthouse --min_replicas 1 --quorum_tick_ms 100 --join_timeout_ms 10000 > /dev/null 2>&1 &
         echo "ft_integration_test"
         # Getting error - Cuda failure 217 'peer access is not supported between these two devices'
-        python ./tests/integration_tests_ft.py artifacts-to-be-uploaded --ngpu 8
+        python -m tests.integration_tests_ft artifacts-to-be-uploaded --ngpu 8
         # pkill -9 torchft_lighthouse

scripts/estimate/estimation.py

@@ -17,7 +17,7 @@
 
 from torchtitan.components.lr_scheduler import build_lr_schedulers
 from torchtitan.components.optimizer import build_optimizers
-from torchtitan.config_manager import ConfigManager, JobConfig
+from torchtitan.config import ConfigManager, JobConfig
 from torchtitan.distributed import ParallelDims, utils as dist_utils
 from torchtitan.protocols.model_converter import build_model_converters
 from torchtitan.protocols.train_spec import get_train_spec

scripts/generate/test_generate.py

@@ -26,7 +26,7 @@
 )
 from torchtitan.components.checkpoint import excluded_parameters_for_model_only
 from torchtitan.components.metrics import build_device_memory_monitor
-from torchtitan.config_manager import ConfigManager
+from torchtitan.config import ConfigManager
 from torchtitan.distributed import ParallelDims, utils as dist_utils
 from torchtitan.protocols.train_spec import get_train_spec
 from torchtitan.tools import utils

tests/integration_tests_ft.py

@@ -10,8 +10,8 @@
 import os
 import subprocess
 from collections import defaultdict
-from dataclasses import dataclass
-from typing import Sequence
+
+from tests.integration_tests import OverrideDefinitions
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -22,22 +22,6 @@
     import tomli as tomllib
 
 
-@dataclass
-class OverrideDefinitions:
-    """
-    This class is used to define the override definitions for the integration tests.
-    """
-
-    override_args: Sequence[Sequence[str]] = tuple(tuple(" "))
-    test_descr: str = "default"
-    test_name: str = "default"
-    ngpu: int = 4
-    model_flavor: str = "debugmodel"
-
-    def __repr__(self):
-        return self.test_descr
-
-
 def build_test_list():
     """
     key is the config file name and value is a list of OverrideDefinitions
@@ -52,6 +36,7 @@ def build_test_list():
             ],
             "Default TorchFT integration test",
             "default_torchft",
+            ngpu=8,
         )
     ]
     return integration_tests_flavors
@@ -65,7 +50,6 @@ def run_test(test_flavor: OverrideDefinitions, full_path: str, output_dir: str):
     # run_test supports sequence of tests.
     test_name = test_flavor.test_name
     dump_folder_arg = f"--job.dump_folder {output_dir}/{test_name}"
-    model_flavor_arg = f"--model.flavor {test_flavor.model_flavor}"
 
     # Use all 8 GPUs in a single replica
     # TODO: Use two replica groups
@@ -79,14 +63,13 @@ def run_test(test_flavor: OverrideDefinitions, full_path: str, output_dir: str):
         for replica_id, ranks in enumerate(all_ranks):
             cmd = (
                 f'TORCH_TRACE="{output_dir}/{test_name}/compile_trace" '
-                + f"CUDA_VISIBLE_DEVICES={ranks}"
-                + f"CONFIG_FILE={full_path} NGPU={len(ranks)} ./run_train.sh "
+                + f"CUDA_VISIBLE_DEVICES={ranks} "
+                + f"CONFIG_FILE={full_path} NGPU={test_flavor.ngpu} ./run_train.sh "
                 + "--fault_tolerance.enable "
-                + f"--fault_tolerance.replica_id={replica_id} --fault_tolerance.group_size={len(all_ranks)}"
+                + f"--fault_tolerance.replica_id={replica_id} --fault_tolerance.group_size={test_flavor.ngpu}"
             )
 
             cmd += " " + dump_folder_arg
-            cmd += " " + model_flavor_arg
             if override_arg:
                 cmd += " " + " ".join(override_arg)
 

tests/integration_tests_h100.py

@@ -9,8 +9,8 @@
 import os
 import subprocess
 from collections import defaultdict
-from dataclasses import dataclass
-from typing import Sequence
+
+from .integration_tests import OverrideDefinitions
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -21,21 +21,6 @@
     import tomli as tomllib
 
 
-@dataclass
-class OverrideDefinitions:
-    """
-    This class is used to define the override definitions for the integration tests.
-    """
-
-    override_args: Sequence[Sequence[str]] = tuple(tuple(" "))
-    test_descr: str = "default"
-    test_name: str = "default"
-    ngpu: int = 4
-
-    def __repr__(self):
-        return self.test_descr
-
-
 def build_test_list():
     """
     key is the config file name and value is a list of OverrideDefinitions

tests/unit_tests/test_activation_checkpoint.py

@@ -10,11 +10,11 @@
 import torch.nn as nn
 from torch.utils.flop_counter import FlopCounterMode
 
-from torchtitan.config_manager import ActivationCheckpoint as ACConfig
+from torchtitan.config.job_config import ActivationCheckpoint as ACConfig
 from torchtitan.models.llama3.infra.parallelize import apply_ac
 
 
-class TestModule(nn.Module):
+class ToyModule(nn.Module):
     def __init__(self):
         super().__init__()
         self.layers = nn.ModuleDict({"0": TransformerBlock()})
@@ -56,12 +56,12 @@ def get_bw_flops(model_fn):
             return mode.get_total_flops() / (512**3 * 2)
 
         # 1. No AC
-        model_no_ac = TestModule()
+        model_no_ac = ToyModule()
         flops_no_ac = get_bw_flops(model_no_ac)
 
         # 2. SAC
         # Per-op SAC's policy is to save every other mm
-        model_selective_ac = TestModule()
+        model_selective_ac = ToyModule()
         ac_config_no_force = ACConfig(
             mode="selective",
             selective_ac_option="op",
@@ -72,7 +72,7 @@ def get_bw_flops(model_fn):
 
         # 3. Per-op SAC with force recompute "moe.router.gate"
         # This leads to two mms being recomputed since they share the same shape!
-        model_with_force_first = TestModule()
+        model_with_force_first = ToyModule()
         ac_config_with_force_first = ACConfig(
             mode="selective",
             selective_ac_option="op",
@@ -82,7 +82,7 @@ def get_bw_flops(model_fn):
         flops_with_force_first = get_bw_flops(model_with_force_first)
 
         # 4. Per-op SAC with force recompute "output"
-        model_with_force_last = TestModule()
+        model_with_force_last = ToyModule()
         ac_config_with_force_last = ACConfig(
             mode="selective",
             selective_ac_option="op",
@@ -92,7 +92,7 @@ def get_bw_flops(model_fn):
         flops_with_force_last = get_bw_flops(model_with_force_last)
 
         # 5. Full AC
-        model_with_full_ac = TestModule()
+        model_with_full_ac = ToyModule()
         ac_config_full_ac = ACConfig(
             mode="full",
         )
@@ -122,12 +122,12 @@ def get_act_mem(model_fn):
             return act_mem
 
         # 1. No AC
-        model_no_ac = TestModule().cuda()
+        model_no_ac = ToyModule().cuda()
         mem_no_ac = get_act_mem(model_no_ac)
 
         # 2. SAC
         # Per-op SAC's policy is to save every other mm
-        model_selective_ac = TestModule().cuda()
+        model_selective_ac = ToyModule().cuda()
         ac_config_no_force = ACConfig(
             mode="selective",
             selective_ac_option="op",
@@ -138,7 +138,7 @@ def get_act_mem(model_fn):
 
         # 3. Per-op SAC with force recompute "moe.router.gate"
         # This leads to two mms being recomputed since they share the same shape!
-        model_with_force_first = TestModule().cuda()
+        model_with_force_first = ToyModule().cuda()
         ac_config_with_force_first = ACConfig(
             mode="selective",
             selective_ac_option="op",
@@ -148,7 +148,7 @@ def get_act_mem(model_fn):
         mem_with_force_first = get_act_mem(model_with_force_first)
 
         # 4. Per-op SAC with force recompute "output"
-        model_with_force_last = TestModule().cuda()
+        model_with_force_last = ToyModule().cuda()
         ac_config_with_force_last = ACConfig(
             mode="selective",
             selective_ac_option="op",
@@ -158,7 +158,7 @@ def get_act_mem(model_fn):
         mem_with_force_last = get_act_mem(model_with_force_last)
 
         # 5. Full AC
-        model_with_full_ac = TestModule().cuda()
+        model_with_full_ac = ToyModule().cuda()
         ac_config_full_ac = ACConfig(
             mode="full",
         )
@@ -175,9 +175,9 @@ def get_act_mem(model_fn):
         # the size of the other two mms.
 
     def test_correctness(self):
-        model_no_ac = TestModule()
+        model_no_ac = ToyModule()
 
-        model_selective_ac = TestModule()
+        model_selective_ac = ToyModule()
         model_selective_ac.load_state_dict(model_no_ac.state_dict())
         apply_ac(
             model_selective_ac,
@@ -187,7 +187,7 @@ def test_correctness(self):
                 per_op_sac_force_recompute_mm_shapes_by_fqns=[],
             ),
         )
-        model_force_first = TestModule()
+        model_force_first = ToyModule()
         model_force_first.load_state_dict(model_no_ac.state_dict())
         apply_ac(
             model_force_first,
@@ -198,7 +198,7 @@ def test_correctness(self):
             ),
         )
 
-        model_force_last = TestModule()
+        model_force_last = ToyModule()
         model_force_last.load_state_dict(model_no_ac.state_dict())
         apply_ac(
             model_force_last,

tests/unit_tests/test_checkpoint.py

@@ -16,7 +16,7 @@
 import torch.nn as nn
 from torch.utils.data import DataLoader
 from torchtitan.components.checkpoint import CheckpointManager
-from torchtitan.config_manager import Checkpoint as CheckpointConfig
+from torchtitan.config.job_config import Checkpoint as CheckpointConfig
 
 
 class FakeOptimizersContainer:
@@ -176,6 +176,7 @@ def test_save_load_restores_state(self, mock_load, mock_save, mock_rank):
             lr_schedulers=self.lr_schedulers,
             states=self.states,
             checkpoint_config=self.job_config.checkpoint,
+            sd_adapter=None,
             base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
@@ -209,6 +210,7 @@ def test_save_and_purge_keeps_last_k_checkpoints(
             lr_schedulers=self.lr_schedulers,
             states=self.states,
             checkpoint_config=self.job_config.checkpoint,
+            sd_adapter=None,
             base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
@@ -250,6 +252,7 @@ def test_nonzero_rank_does_not_purge_or_save(self, mock_load, mock_save, mock_ra
             lr_schedulers=self.lr_schedulers,
             states=self.states,
             checkpoint_config=self.job_config.checkpoint,
+            sd_adapter=None,
             base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
@@ -273,6 +276,7 @@ def test_load_returns_false_when_no_checkpoint_folder(self):
             lr_schedulers=self.lr_schedulers,
             states=self.states,
             checkpoint_config=self.job_config.checkpoint,
+            sd_adapter=None,
             base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
@@ -297,6 +301,7 @@ def test_load_finds_latest_and_calls_dcp_load(self, mock_load, mock_rank):
             lr_schedulers=self.lr_schedulers,
             states=self.states,
             checkpoint_config=self.job_config.checkpoint,
+            sd_adapter=None,
             base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
@@ -327,6 +332,7 @@ def test_interval_respects_interval(self, mock_load, mock_save, mock_rank):
             lr_schedulers=self.lr_schedulers,
             states=self.states,
             checkpoint_config=self.job_config.checkpoint,
+            sd_adapter=None,
             base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
@@ -361,6 +367,7 @@ def test_last_save_model_only_and_initial_load_model_only(
             lr_schedulers=self.lr_schedulers,
             states=self.states,
             checkpoint_config=self.job_config.checkpoint,
+            sd_adapter=None,
             base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
@@ -381,6 +388,7 @@ def test_last_save_model_only_and_initial_load_model_only(
             lr_schedulers=self.lr_schedulers,
             states=self.states,
             checkpoint_config=self.job_config.checkpoint,
+            sd_adapter=None,
             base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
@@ -423,6 +431,7 @@ def test_async_save_calls_async_wait(self, mock_async_save, mock_new_group):
             lr_schedulers=self.lr_schedulers,
             states=states,
             checkpoint_config=checkpoint_config,
+            sd_adapter=None,
             base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
@@ -468,6 +477,7 @@ def test_ft_async_save_calls_async_wait(
             lr_schedulers=self.lr_schedulers,
             states=self.states,
             checkpoint_config=checkpoint_config,
+            sd_adapter=None,
             base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
@@ -504,6 +514,7 @@ def test_enable_first_step_checkpoint(self, mock_save, mock_rank):
             lr_schedulers=self.lr_schedulers,
             states=self.states,
             checkpoint_config=self.job_config.checkpoint,
+            sd_adapter=None,
             base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
@@ -530,6 +541,7 @@ def test_enable_first_step_checkpoint(self, mock_save, mock_rank):
             lr_schedulers=self.lr_schedulers,
             states=self.states,
             checkpoint_config=self.job_config.checkpoint,
+            sd_adapter=None,
             base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
@@ -576,6 +588,7 @@ def __init__(self):
             lr_schedulers=self.lr_schedulers,
             states=self.states,
             checkpoint_config=self.job_config.checkpoint,
+            sd_adapter=None,
             base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
@@ -626,6 +639,7 @@ def fake_load(state_dict: dict, checkpoint_id=None):
             lr_schedulers=self.lr_schedulers,
             states=self.states,
             checkpoint_config=self.job_config.checkpoint,
+            sd_adapter=None,
             base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )

tests/unit_tests/test_dataset_checkpointing.py

@@ -9,7 +9,7 @@
 import torch
 from datasets import load_dataset
 from torchtitan.components.tokenizer import HuggingFaceTokenizer
-from torchtitan.config_manager import ConfigManager
+from torchtitan.config import ConfigManager
 from torchtitan.datasets.hf_datasets import build_hf_dataloader, DatasetConfig, DATASETS