From 5231397dc2e7bf679d09944c86d60a3150a2662c Mon Sep 17 00:00:00 2001
From: ykarnati <ykarnati@nvidia.com>
Date: Fri, 29 May 2026 10:24:51 -0700
Subject: [PATCH 1/3] Move MIMO unit tests into models/mimo

---
 tests/unit_tests/models/{ => mimo}/test_mimo_1f1b_schedule.py | 2 +-
 .../models/{ => mimo}/test_mimo_audio_submodules.py           | 2 +-
 tests/unit_tests/models/{ => mimo}/test_mimo_checkpoint.py    | 4 ++--
 .../models/{ => mimo}/test_mimo_colocated_communicator.py     | 0
 .../models/{ => mimo}/test_mimo_colocated_correctness.py      | 4 ++--
 .../models/{ => mimo}/test_mimo_embedding_alignment.py        | 2 +-
 tests/unit_tests/models/{ => mimo}/test_mimo_model.py         | 2 +-
 tests/unit_tests/models/{ => mimo}/test_mimo_partition.py     | 2 +-
 tests/unit_tests/models/{ => mimo}/test_mimo_role.py          | 0
 tests/unit_tests/models/{ => mimo}/test_mimo_submodules.py    | 2 +-
 10 files changed, 10 insertions(+), 10 deletions(-)
 rename tests/unit_tests/models/{ => mimo}/test_mimo_1f1b_schedule.py (99%)
 rename tests/unit_tests/models/{ => mimo}/test_mimo_audio_submodules.py (99%)
 rename tests/unit_tests/models/{ => mimo}/test_mimo_checkpoint.py (98%)
 rename tests/unit_tests/models/{ => mimo}/test_mimo_colocated_communicator.py (100%)
 rename tests/unit_tests/models/{ => mimo}/test_mimo_colocated_correctness.py (99%)
 rename tests/unit_tests/models/{ => mimo}/test_mimo_embedding_alignment.py (99%)
 rename tests/unit_tests/models/{ => mimo}/test_mimo_model.py (99%)
 rename tests/unit_tests/models/{ => mimo}/test_mimo_partition.py (99%)
 rename tests/unit_tests/models/{ => mimo}/test_mimo_role.py (100%)
 rename tests/unit_tests/models/{ => mimo}/test_mimo_submodules.py (99%)

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 5cb28313cfd..de6c030eeff 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 accee3faf7f..4051b08486c 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

From d12ab934010dd9cbc70f3c58aa2fb81e4b53450f Mon Sep 17 00:00:00 2001
From: ykarnati <ykarnati@nvidia.com>
Date: Fri, 29 May 2026 10:35:59 -0700
Subject: [PATCH 2/3] Add MIMO test run README

---
 tests/unit_tests/models/mimo/README.md | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)
 create mode 100644 tests/unit_tests/models/mimo/README.md

diff --git a/tests/unit_tests/models/mimo/README.md b/tests/unit_tests/models/mimo/README.md
new file mode 100644
index 00000000000..29c3f5ecf29
--- /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=2 \
+  -m pytest tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py::TestMimo1F1BSchedule::test_baseline_2gpu -v -s
+```
+
+Larger non-colocated 4-GPU and 8-GPU examples are in `test_mimo_1f1b_schedule.py`.

From 142d733efb5a0f57e26c3b7be4dfc85b13eaebbe Mon Sep 17 00:00:00 2001
From: ykarnati <ykarnati@nvidia.com>
Date: Fri, 29 May 2026 10:43:06 -0700
Subject: [PATCH 3/3] Use 8 GPU non-colocated MIMO example

---
 tests/unit_tests/models/mimo/README.md | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/unit_tests/models/mimo/README.md b/tests/unit_tests/models/mimo/README.md
index 29c3f5ecf29..327cc68fcce 100644
--- a/tests/unit_tests/models/mimo/README.md
+++ b/tests/unit_tests/models/mimo/README.md
@@ -12,8 +12,8 @@ uv run python -m torch.distributed.run --nproc_per_node=8 \
 ## Non-colocated encoder + LLM
 
 ```bash
-uv run python -m torch.distributed.run --nproc-per-node=2 \
-  -m pytest tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py::TestMimo1F1BSchedule::test_baseline_2gpu -v -s
+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
 ```
 
-Larger non-colocated 4-GPU and 8-GPU examples are in `test_mimo_1f1b_schedule.py`.
+More non-colocated 8-GPU examples are in `test_mimo_1f1b_schedule.py`.