Move MIMO unit tests into models/mimo

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`.

tests/unit_tests/models/test_mimo_1f1b_schedule.py → 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

tests/unit_tests/models/test_mimo_audio_submodules.py → 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

tests/unit_tests/models/test_mimo_checkpoint.py → 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,

tests/unit_tests/models/test_mimo_colocated_correctness.py → 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,

tests/unit_tests/models/test_mimo_embedding_alignment.py → 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

tests/unit_tests/models/test_mimo_model.py → 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

tests/unit_tests/models/test_mimo_partition.py → 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

tests/unit_tests/models/test_mimo_submodules.py → 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