Repository: pytorch/torchtitan (language: python)

Issue:
Add test for toml-based pp split points
Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #664

Adds 2 test cases:
- first sets split points in new toml file
- second keeps split points in toml, but also overrides them with different values on cmdline.

both tests work as expected locally (logs from model init show the expected layers are used per-stage in both cases)



Discussion / hints:

Fix patch:
diff --git a/.ci/docker/dev-requirements.txt b/.ci/docker/dev-requirements.txt
index 770301a0..bd611222 100644
--- a/.ci/docker/dev-requirements.txt
+++ b/.ci/docker/dev-requirements.txt
@@ -2,3 +2,4 @@ expecttest==0.1.6
 pytest==7.3.2
 pytest-cov
 pre-commit
+tomli-w >= 1.1.0


Test patch:
diff --git a/test/test_job_config.py b/test/test_job_config.py
index e4ef04ba..aed007ad 100644
--- a/test/test_job_config.py
+++ b/test/test_job_config.py
@@ -7,6 +7,7 @@
 import tempfile
 
 import pytest
+import tomli_w
 from torchtitan.config_manager import JobConfig
 
 
@@ -44,6 +45,77 @@ def test_job_config_file_cmd_overrides(self):
         )
         assert config.job.dump_folder == "/tmp/test_tt/"
 
+    def test_parse_pp_split_points(self):
+
+        toml_splits = ["layers.2", "layers.4", "layers.6"]
+        toml_split_str = ",".join(toml_splits)
+        cmdline_splits = ["layers.1", "layers.3", "layers.5"]
+        cmdline_split_str = ",".join(cmdline_splits)
+        # no split points specified
+        config = JobConfig()
+        config.parse_args(
+            [
+                "--job.config_file",
+                "./train_configs/debug_model.toml",
+            ]
+        )
+        assert config.experimental.pipeline_parallel_split_points == []
+
+        # toml has no split points, but cmdline splits are specified
+        config = JobConfig()
+        config.parse_args(
+            [
+                "--job.config_file",
+                "./train_configs/debug_model.toml",
+                "--experimental.pipeline_parallel_split_points",
+                f"{cmdline_split_str}",
+            ]
+        )
+        assert (
+            config.experimental.pipeline_parallel_split_points == cmdline_splits
+        ), config.experimental.pipeline_parallel_split_points
+
+        # toml has split points, cmdline does not
+        with tempfile.NamedTemporaryFile() as fp:
+            with open(fp.name, "wb") as f:
+                tomli_w.dump(
+                    {
+                        "experimental": {
+                            "pipeline_parallel_split_points": toml_split_str,
+                        }
+                    },
+                    f,
+                )
+            config = JobConfig()
+            config.parse_args(["--job.config_file", fp.name])
+            assert (
+                config.experimental.pipeline_parallel_split_points == toml_splits
+            ), config.experimental.pipeline_parallel_split_points
+
+        # toml has split points, cmdline overrides them
+        with tempfile.NamedTemporaryFile() as fp:
+            with open(fp.name, "wb") as f:
+                tomli_w.dump(
+                    {
+                        "experimental": {
+                            "pipeline_parallel_split_points": toml_split_str,
+                        }
+                    },
+                    f,
+                )
+            config = JobConfig()
+            config.parse_args(
+                [
+                    "--job.config_file",
+                    fp.name,
+                    "--experimental.pipeline_parallel_split_points",
+                    f"{cmdline_split_str}",
+                ]
+            )
+            assert (
+                config.experimental.pipeline_parallel_split_points == cmdline_splits
+            ), config.experimental.pipeline_parallel_split_points
+
     def test_print_help(self):
         config = JobConfig()
         parser = config.parser