pytorch · tianyu-l · Dec 16, 2024 · Dec 16, 2024 · Dec 17, 2024 · Dec 17, 2024
diff --git a/.github/workflows/integration_test_4gpu.yaml b/.github/workflows/integration_test_4gpu.yaml
@@ -43,4 +43,4 @@ jobs:
         python -m pip install -e .
 
         mkdir artifacts-to-be-uploaded
-        python ./test_runner.py artifacts-to-be-uploaded --ngpu 4
+        python ./tests/integration_tests.py artifacts-to-be-uploaded --ngpu 4
diff --git a/.github/workflows/integration_test_8gpu.yaml b/.github/workflows/integration_test_8gpu.yaml
@@ -38,4 +38,4 @@ jobs:
 
         python -m pip install --force-reinstall --pre torch --index-url https://download.pytorch.org/whl/nightly/cu124
         mkdir artifacts-to-be-uploaded
-        python ./test_runner.py artifacts-to-be-uploaded --ngpu 8
+        python ./tests/integration_tests.py artifacts-to-be-uploaded --ngpu 8
diff --git a/.github/workflows/unit_test_cpu.yaml b/.github/workflows/unit_test_cpu.yaml
@@ -25,4 +25,4 @@ jobs:
         pip config --user set global.progress_bar off
 
         pip install --force-reinstall --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu
-        pytest test --cov=. --cov-report=xml --durations=20 -vv
+        pytest tests/unit_tests --cov=. --cov-report=xml --durations=20 -vv
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -24,7 +24,7 @@ repos:
         files: \.py$
         args:
         - --license-filepath
-        - docs/license_header.txt
+        - assets/license_header.txt
 
 -   repo: https://github.com/pycqa/flake8
     rev: 34cbf8ef3950f43d09b85e2e45c15ae5717dc37b

diff --git a/assets/images/readme.md b/assets/images/readme.md
diff --git a/docs/license_header.txt → assets/license_header.txt b/docs/license_header.txt → assets/license_header.txt
diff --git a/version.txt → assets/version.txt b/version.txt → assets/version.txt
diff --git a/pyproject.toml b/pyproject.toml
@@ -36,7 +36,7 @@ dev = [
 ]
 
 [tool.setuptools.dynamic]
-version = {file = "version.txt"}
+version = {file = "assets/version.txt"}
 
 
 # ---- Explicit project build information ---- #

diff --git a/test/__init__.py → tests/__init__.py b/test/__init__.py → tests/__init__.py
diff --git a/test/assets/c4_test/data.json → tests/assets/c4_test/data.json b/test/assets/c4_test/data.json → tests/assets/c4_test/data.json
diff --git a/test/assets/custom_schedule.csv → tests/assets/custom_schedule.csv b/test/assets/custom_schedule.csv → tests/assets/custom_schedule.csv
diff --git a/test/assets/test_tiktoken.model → tests/assets/test_tiktoken.model b/test/assets/test_tiktoken.model → tests/assets/test_tiktoken.model
diff --git a/test_runner.py → tests/integration_tests.py b/test_runner.py → tests/integration_tests.py
@@ -260,7 +260,7 @@ def build_test_list():
                     "--checkpoint.enable_checkpoint",
                     "--experimental.pipeline_parallel_degree 2",
                     "--experimental.pipeline_parallel_schedule PipelineScheduleMulti",
-                    "--experimental.pipeline_parallel_schedule_csv ./test/assets/custom_schedule.csv",
+                    "--experimental.pipeline_parallel_schedule_csv ./tests/assets/custom_schedule.csv",
                     "--experimental.pipeline_parallel_microbatches 8",
                 ],
             ],

diff --git a/test/datasets/__init__.py → tests/unit_tests/__init__.py b/test/datasets/__init__.py → tests/unit_tests/__init__.py
diff --git a/test/multimodal_model/__init__.py → ...s/unit_tests/multimodal_model/__init__.py b/test/multimodal_model/__init__.py → ...s/unit_tests/multimodal_model/__init__.py
diff --git a/...multimodal_model/test_multimodal_model.py → ...multimodal_model/test_multimodal_model.py b/...multimodal_model/test_multimodal_model.py → ...multimodal_model/test_multimodal_model.py
@@ -12,7 +12,10 @@
     VisionEncoder,
 )
 
-from test.multimodal_model.test_utils import fixed_init_model, fixed_init_tensor
+from tests.unit_tests.multimodal_model.test_utils import (
+    fixed_init_model,
+    fixed_init_tensor,
+)
 
 
 @pytest.fixture

diff --git a/test/multimodal_model/test_utils.py → ...unit_tests/multimodal_model/test_utils.py b/test/multimodal_model/test_utils.py → ...unit_tests/multimodal_model/test_utils.py
diff --git a/test/datasets/test_checkpoint.py → .../unit_tests/test_dataset_checkpointing.py b/test/datasets/test_checkpoint.py → .../unit_tests/test_dataset_checkpointing.py
@@ -9,10 +9,10 @@
 from torchtitan.datasets.tokenizer import build_tokenizer
 
 
-class TestCheckpoint:
+class TestDatasetCheckpointing:
     def test_c4_resumption(self):
         dataset_name = "c4_test"
-        dataset_path = "./test/assets/c4_test"
+        dataset_path = "./tests/assets/c4_test"
         batch_size = 1
         seq_len = 1024
         world_size = 4
@@ -41,8 +41,7 @@ def test_c4_resumption(self):
     def _build_dataloader(
         self, dataset_name, dataset_path, batch_size, seq_len, world_size, rank
     ):
-        tokenizer_type = "tiktoken"
-        tokenizer = build_tokenizer("tiktoken", "./test/assets/test_tiktoken.model")
+        tokenizer = build_tokenizer("tiktoken", "./tests/assets/test_tiktoken.model")
         return build_hf_data_loader(
             dataset_name=dataset_name,
             dataset_path=dataset_path,

diff --git a/test/test_fused_rms_norm.py → ...unit_tests/test_fused_rms_norm_dtensor.py b/test/test_fused_rms_norm.py → ...unit_tests/test_fused_rms_norm_dtensor.py
diff --git a/test/test_job_config.py → tests/unit_tests/test_job_config.py b/test/test_job_config.py → tests/unit_tests/test_job_config.py
diff --git a/torchtitan/datasets/hf_datasets.py b/torchtitan/datasets/hf_datasets.py
@@ -45,7 +45,7 @@ class DatasetConfig:
         text_processor=_process_c4_text,
     ),
     "c4_test": DatasetConfig(
-        path="test/assets/c4_test",
+        path="tests/assets/c4_test",
         loader=lambda path: load_dataset(path, split="train"),
         text_processor=_process_c4_text,
     ),

diff --git a/train_configs/debug_model.toml b/train_configs/debug_model.toml
@@ -24,7 +24,7 @@ name = "llama3"
 flavor = "debugmodel"
 norm_type = "rmsnorm"  # layernorm / np_layernorm / rmsnorm / fused_rmsnorm
 # test tokenizer.model, for debug purpose only
-tokenizer_path = "./test/assets/test_tiktoken.model"
+tokenizer_path = "./tests/assets/test_tiktoken.model"
 
 [optimizer]
 name = "AdamW"