diff --git a/tests/unit_tests/models/mimo/README.md b/tests/unit_tests/models/mimo/README.md new file mode 100644 index 00000000000..327cc68fcce --- /dev/null +++ b/tests/unit_tests/models/mimo/README.md @@ -0,0 +1,19 @@ +# MIMO Training Examples + +Run from the repository root. + +## Colocated encoder + LLM + +```bash +uv run python -m torch.distributed.run --nproc_per_node=8 \ + -m pytest tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py -v -s +``` + +## Non-colocated encoder + LLM + +```bash +uv run python -m torch.distributed.run --nproc-per-node=8 \ + -m pytest tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py::TestMimo1F1BSchedule::test_encoder_tp2_llm_tp2_pp3_8gpu -v -s +``` + +More non-colocated 8-GPU examples are in `test_mimo_1f1b_schedule.py`. diff --git a/tests/unit_tests/models/test_mimo_1f1b_schedule.py b/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py similarity index 99% rename from tests/unit_tests/models/test_mimo_1f1b_schedule.py rename to tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py index 836382b21cc..bdd565dae36 100644 --- a/tests/unit_tests/models/test_mimo_1f1b_schedule.py +++ b/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py @@ -3,7 +3,7 @@ """Integration tests for MIMO model with 1F1B pipeline schedule. Run with: - uv run python -m torch.distributed.run --nproc-per-node=2 -m pytest tests/unit_tests/models/test_mimo_1f1b_schedule.py -v + uv run python -m torch.distributed.run --nproc-per-node=2 -m pytest tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py -v """ import logging diff --git a/tests/unit_tests/models/test_mimo_audio_submodules.py b/tests/unit_tests/models/mimo/test_mimo_audio_submodules.py similarity index 99% rename from tests/unit_tests/models/test_mimo_audio_submodules.py rename to tests/unit_tests/models/mimo/test_mimo_audio_submodules.py index f9a18838f60..c8ed8980bc2 100644 --- a/tests/unit_tests/models/test_mimo_audio_submodules.py +++ b/tests/unit_tests/models/mimo/test_mimo_audio_submodules.py @@ -1,7 +1,7 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. ''' -WORLD_SIZE=1 LOCAL_RANK=0 python -m pytest tests/unit_tests/models/test_mimo_audio_submodules.py +WORLD_SIZE=1 LOCAL_RANK=0 python -m pytest tests/unit_tests/models/mimo/test_mimo_audio_submodules.py ''' import math import random diff --git a/tests/unit_tests/models/test_mimo_checkpoint.py b/tests/unit_tests/models/mimo/test_mimo_checkpoint.py similarity index 98% rename from tests/unit_tests/models/test_mimo_checkpoint.py rename to tests/unit_tests/models/mimo/test_mimo_checkpoint.py index 730d01e1c5b..dc12e53571a 100644 --- a/tests/unit_tests/models/test_mimo_checkpoint.py +++ b/tests/unit_tests/models/mimo/test_mimo_checkpoint.py @@ -4,7 +4,7 @@ Run with 8 GPUs: uv run python -m torch.distributed.run --nproc-per-node=8 \ - -m pytest tests/unit_tests/models/test_mimo_checkpoint.py -v -s + -m pytest tests/unit_tests/models/mimo/test_mimo_checkpoint.py -v -s """ import os @@ -20,7 +20,7 @@ from megatron.core.dist_checkpointing.validation import StrictHandling from megatron.core.models.mimo.optimizer import get_mimo_optimizer from megatron.core.optimizer.optimizer_config import OptimizerConfig -from tests.unit_tests.models.test_mimo_1f1b_schedule import ( +from tests.unit_tests.models.mimo.test_mimo_1f1b_schedule import ( create_all_embedding_groups, create_hypercomm_grid, destroy_all_grids, diff --git a/tests/unit_tests/models/test_mimo_colocated_communicator.py b/tests/unit_tests/models/mimo/test_mimo_colocated_communicator.py similarity index 100% rename from tests/unit_tests/models/test_mimo_colocated_communicator.py rename to tests/unit_tests/models/mimo/test_mimo_colocated_communicator.py diff --git a/tests/unit_tests/models/test_mimo_colocated_correctness.py b/tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py similarity index 99% rename from tests/unit_tests/models/test_mimo_colocated_correctness.py rename to tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py index e2d91bdf83e..71ff13ec557 100644 --- a/tests/unit_tests/models/test_mimo_colocated_correctness.py +++ b/tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py @@ -47,7 +47,7 @@ Run with:: uv run python -m torch.distributed.run --nproc_per_node=8 \\ - -m pytest tests/unit_tests/models/test_mimo_colocated_correctness.py -v -s + -m pytest tests/unit_tests/models/mimo/test_mimo_colocated_correctness.py -v -s """ import os @@ -65,7 +65,7 @@ from megatron.core.optimizer.optimizer_config import OptimizerConfig from megatron.core.transformer.enums import ModelType from megatron.core.utils import unwrap_model -from tests.unit_tests.models.test_mimo_1f1b_schedule import ( +from tests.unit_tests.models.mimo.test_mimo_1f1b_schedule import ( build_no_sync_func, create_all_embedding_groups, create_hypercomm_grid, diff --git a/tests/unit_tests/models/test_mimo_embedding_alignment.py b/tests/unit_tests/models/mimo/test_mimo_embedding_alignment.py similarity index 99% rename from tests/unit_tests/models/test_mimo_embedding_alignment.py rename to tests/unit_tests/models/mimo/test_mimo_embedding_alignment.py index 688ebe4832b..141093d7769 100644 --- a/tests/unit_tests/models/test_mimo_embedding_alignment.py +++ b/tests/unit_tests/models/mimo/test_mimo_embedding_alignment.py @@ -1,7 +1,7 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. ''' -WORLD_SIZE=1 LOCAL_RANK=0 python -m pytest tests/unit_tests/models/test_mimo_embedding_alignment.py +WORLD_SIZE=1 LOCAL_RANK=0 python -m pytest tests/unit_tests/models/mimo/test_mimo_embedding_alignment.py ''' from unittest.mock import MagicMock diff --git a/tests/unit_tests/models/test_mimo_model.py b/tests/unit_tests/models/mimo/test_mimo_model.py similarity index 99% rename from tests/unit_tests/models/test_mimo_model.py rename to tests/unit_tests/models/mimo/test_mimo_model.py index 88babea4678..85583d495e8 100644 --- a/tests/unit_tests/models/test_mimo_model.py +++ b/tests/unit_tests/models/mimo/test_mimo_model.py @@ -1,7 +1,7 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. ''' -WORLD_SIZE=1 LOCAL_RANK=0 python -m pytest tests/unit_tests/models/test_mimo_model.py +WORLD_SIZE=1 LOCAL_RANK=0 python -m pytest tests/unit_tests/models/mimo/test_mimo_model.py ''' import math diff --git a/tests/unit_tests/models/test_mimo_partition.py b/tests/unit_tests/models/mimo/test_mimo_partition.py similarity index 99% rename from tests/unit_tests/models/test_mimo_partition.py rename to tests/unit_tests/models/mimo/test_mimo_partition.py index 3db10e81f51..6f1dc587f89 100644 --- a/tests/unit_tests/models/test_mimo_partition.py +++ b/tests/unit_tests/models/mimo/test_mimo_partition.py @@ -3,7 +3,7 @@ ''' WORLD_SIZE=1 LOCAL_RANK=0 python -m torch.distributed.run \ --nproc_per_node=1 -m pytest \ - tests/unit_tests/models/test_mimo_partition.py -v + tests/unit_tests/models/mimo/test_mimo_partition.py -v ''' from unittest.mock import MagicMock, patch diff --git a/tests/unit_tests/models/test_mimo_role.py b/tests/unit_tests/models/mimo/test_mimo_role.py similarity index 100% rename from tests/unit_tests/models/test_mimo_role.py rename to tests/unit_tests/models/mimo/test_mimo_role.py diff --git a/tests/unit_tests/models/test_mimo_submodules.py b/tests/unit_tests/models/mimo/test_mimo_submodules.py similarity index 99% rename from tests/unit_tests/models/test_mimo_submodules.py rename to tests/unit_tests/models/mimo/test_mimo_submodules.py index 60730486d1c..0e8ca869ce3 100644 --- a/tests/unit_tests/models/test_mimo_submodules.py +++ b/tests/unit_tests/models/mimo/test_mimo_submodules.py @@ -3,7 +3,7 @@ ''' WORLD_SIZE=1 LOCAL_RANK=0 python -m torch.distributed.run \ --nproc_per_node=1 -m pytest \ - tests/unit_tests/models/test_mimo_submodules.py -v + tests/unit_tests/models/mimo/test_mimo_submodules.py -v ''' from typing import Any, Dict, List, Optional