Skip to content

Commit 6408e9b

Browse files
liding-nvliding-nv
committed
tests
Signed-off-by: Li Ding <liding@nvidia.com>
1 parent 617c404 commit 6408e9b

2 files changed

Lines changed: 42 additions & 46 deletions

File tree

‎tests/unit_tests/training/megatron_mimo/test_megatron_mimo_checkpointing.py‎

Lines changed: 27 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -214,7 +214,6 @@ class TestPretrainMegatronMIMOSetup:
214214
@patch("megatron.bridge.training.setup_megatron_mimo.create_checkpoint_manager")
215215
@patch("megatron.bridge.training.setup_megatron_mimo.MultiModulePipelineCommunicator")
216216
@patch("megatron.bridge.training.setup_megatron_mimo.get_model_config")
217-
@patch("megatron.bridge.training.setup_megatron_mimo.validate_no_stub_ranks")
218217
@patch("megatron.bridge.training.setup_megatron_mimo.build_pg_collection_for_schedule")
219218
@patch("megatron.bridge.training.setup_megatron_mimo.get_module_to_grid_tuple")
220219
@patch("megatron.bridge.training.setup_megatron_mimo._update_megatron_mimo_model_config_funcs")
@@ -227,7 +226,6 @@ def test_setup_megatron_mimo_initializes_checkpoint_manager(
227226
mock_update_config_funcs,
228227
mock_get_grid,
229228
mock_build_pg,
230-
mock_validate,
231229
mock_get_config,
232230
mock_communicator,
233231
mock_create_ckpt_mgr,
@@ -278,20 +276,22 @@ def test_setup_megatron_mimo_initializes_checkpoint_manager(
278276
infra.topology = Mock()
279277
infra.module_output_ndim = {"language": 3}
280278
infra.pg_collections = {"language": Mock()}
281-
cfg.model.build_infra.return_value = infra
282-
cfg.model.provide_distributed_model.return_value = [Mock()]
279+
model = Mock()
283280

284281
mock_optimizer = MagicMock()
285282
mock_optimizer.module_infos = {}
283+
start_time_tensor = Mock()
284+
start_time_tensor.item.return_value = global_state.start_time
286285

287286
with (
288-
patch("megatron.bridge.training.setup_megatron_mimo._set_megatron_mimo_random_seeds"),
287+
patch("megatron.bridge.models.megatron_mimo.build_megatron_mimo_model", return_value=(model, infra)),
289288
patch("megatron.core.models.mimo.optimizer.get_mimo_optimizer", return_value=mock_optimizer),
290289
patch("megatron.core.num_microbatches_calculator._GLOBAL_NUM_MICROBATCHES_CALCULATOR", None),
291290
patch("megatron.core.num_microbatches_calculator.init_num_microbatches_calculator"),
292291
patch("megatron.core.parallel_state._TENSOR_MODEL_PARALLEL_GROUP", None),
293292
patch("megatron.core.parallel_state._DATA_PARALLEL_GROUP", None),
294293
patch("megatron.core.parallel_state._DATA_PARALLEL_GROUP_WITH_CP", None),
294+
patch("torch.tensor", return_value=start_time_tensor),
295295
):
296296
result = setup_megatron_mimo(state=global_state)
297297

@@ -990,13 +990,11 @@ class TestSetupMegatronMIMOCheckpointLoading:
990990
"megatron.bridge.training.setup_megatron_mimo.create_checkpoint_manager",
991991
"megatron.bridge.training.setup_megatron_mimo.MultiModulePipelineCommunicator",
992992
"megatron.bridge.training.setup_megatron_mimo.get_model_config",
993-
"megatron.bridge.training.setup_megatron_mimo.validate_no_stub_ranks",
994993
"megatron.bridge.training.setup_megatron_mimo.build_pg_collection_for_schedule",
995994
"megatron.bridge.training.setup_megatron_mimo.get_module_to_grid_tuple",
996995
"megatron.bridge.training.setup_megatron_mimo._update_megatron_mimo_model_config_funcs",
997996
"megatron.bridge.training.setup_megatron_mimo.unwrap_megatron_mimo_model",
998997
"megatron.bridge.training.setup_megatron_mimo.dist",
999-
"megatron.bridge.training.setup_megatron_mimo._set_megatron_mimo_random_seeds",
1000998
"megatron.core.num_microbatches_calculator._GLOBAL_NUM_MICROBATCHES_CALCULATOR",
1001999
"megatron.core.num_microbatches_calculator.init_num_microbatches_calculator",
10021000
"megatron.core.parallel_state._TENSOR_MODEL_PARALLEL_GROUP",
@@ -1039,8 +1037,7 @@ def _run_setup(self, *, load_path=None, checkpoint_exists_return=False):
10391037
infra.topology = Mock()
10401038
infra.module_output_ndim = {"language": 3}
10411039
infra.pg_collections = {"language": Mock()}
1042-
cfg.model.build_infra.return_value = infra
1043-
cfg.model.provide_distributed_model.return_value = [Mock()]
1040+
model = Mock()
10441041

10451042
local_pg = MagicMock()
10461043
mock_optimizer = MagicMock()
@@ -1087,12 +1084,18 @@ def _run_setup(self, *, load_path=None, checkpoint_exists_return=False):
10871084
stack.enter_context(
10881085
patch("megatron.core.models.mimo.optimizer.get_mimo_optimizer", return_value=mock_optimizer)
10891086
)
1087+
stack.enter_context(
1088+
patch("megatron.bridge.models.megatron_mimo.build_megatron_mimo_model", return_value=(model, infra))
1089+
)
10901090

10911091
model_config = Mock(pipeline_dtype=None, bf16=True)
10921092
stack.enter_context(
10931093
patch("megatron.bridge.training.setup_megatron_mimo.get_model_config", return_value=model_config)
10941094
)
10951095
stack.enter_context(patch("megatron.bridge.training.setup_megatron_mimo.unwrap_megatron_mimo_model"))
1096+
start_time_tensor = Mock()
1097+
start_time_tensor.item.return_value = state.start_time
1098+
stack.enter_context(patch("torch.tensor", return_value=start_time_tensor))
10961099

10971100
result = setup_megatron_mimo(state=state)
10981101

@@ -1171,8 +1174,7 @@ def _sig_fn(cfg, megatron_mimo_infra, *, train_state=None):
11711174
infra.topology = Mock()
11721175
infra.module_output_ndim = {"language": 3}
11731176
infra.pg_collections = {"language": Mock()}
1174-
cfg.model.build_infra.return_value = infra
1175-
cfg.model.provide_distributed_model.return_value = [Mock()]
1177+
model = Mock()
11761178

11771179
state = Mock()
11781180
state.cfg = cfg
@@ -1182,11 +1184,12 @@ def _sig_fn(cfg, megatron_mimo_infra, *, train_state=None):
11821184

11831185
mock_optimizer = MagicMock()
11841186
mock_optimizer.module_infos = {}
1187+
start_time_tensor = Mock()
1188+
start_time_tensor.item.return_value = state.start_time
11851189

11861190
with (
11871191
patch("megatron.bridge.training.setup_megatron_mimo.dist") as m_dist,
1188-
patch("megatron.bridge.training.setup_megatron_mimo._set_megatron_mimo_random_seeds"),
1189-
patch("megatron.bridge.training.setup_megatron_mimo.validate_no_stub_ranks"),
1192+
patch("megatron.bridge.models.megatron_mimo.build_megatron_mimo_model", return_value=(model, infra)),
11901193
patch("megatron.bridge.training.setup_megatron_mimo.build_pg_collection_for_schedule"),
11911194
patch("megatron.bridge.training.setup_megatron_mimo.get_module_to_grid_tuple"),
11921195
patch("megatron.bridge.training.setup_megatron_mimo.MultiModulePipelineCommunicator"),
@@ -1211,6 +1214,7 @@ def _sig_fn(cfg, megatron_mimo_infra, *, train_state=None):
12111214
patch("megatron.core.parallel_state._TENSOR_MODEL_PARALLEL_GROUP", None),
12121215
patch("megatron.core.parallel_state._DATA_PARALLEL_GROUP", None),
12131216
patch("megatron.core.parallel_state._DATA_PARALLEL_GROUP_WITH_CP", None),
1217+
patch("torch.tensor", return_value=start_time_tensor),
12141218
):
12151219
m_dist.get_rank.return_value = 0
12161220
m_dist.get_world_size.return_value = 2
@@ -1253,8 +1257,7 @@ def test_no_train_state_when_not_resuming(self):
12531257
infra.topology = Mock()
12541258
infra.module_output_ndim = {"language": 3}
12551259
infra.pg_collections = {"language": Mock()}
1256-
cfg.model.build_infra.return_value = infra
1257-
cfg.model.provide_distributed_model.return_value = [Mock()]
1260+
model = Mock()
12581261

12591262
state = Mock()
12601263
state.cfg = cfg
@@ -1263,11 +1266,12 @@ def test_no_train_state_when_not_resuming(self):
12631266

12641267
mock_optimizer = MagicMock()
12651268
mock_optimizer.module_infos = {}
1269+
start_time_tensor = Mock()
1270+
start_time_tensor.item.return_value = state.start_time
12661271

12671272
with (
12681273
patch("megatron.bridge.training.setup_megatron_mimo.dist") as m_dist,
1269-
patch("megatron.bridge.training.setup_megatron_mimo._set_megatron_mimo_random_seeds"),
1270-
patch("megatron.bridge.training.setup_megatron_mimo.validate_no_stub_ranks"),
1274+
patch("megatron.bridge.models.megatron_mimo.build_megatron_mimo_model", return_value=(model, infra)),
12711275
patch("megatron.bridge.training.setup_megatron_mimo.build_pg_collection_for_schedule"),
12721276
patch("megatron.bridge.training.setup_megatron_mimo.get_module_to_grid_tuple"),
12731277
patch("megatron.bridge.training.setup_megatron_mimo.MultiModulePipelineCommunicator"),
@@ -1292,6 +1296,7 @@ def test_no_train_state_when_not_resuming(self):
12921296
patch("megatron.core.parallel_state._TENSOR_MODEL_PARALLEL_GROUP", None),
12931297
patch("megatron.core.parallel_state._DATA_PARALLEL_GROUP", None),
12941298
patch("megatron.core.parallel_state._DATA_PARALLEL_GROUP_WITH_CP", None),
1299+
patch("torch.tensor", return_value=start_time_tensor),
12951300
):
12961301
m_dist.get_rank.return_value = 0
12971302
m_dist.get_world_size.return_value = 2
@@ -1337,8 +1342,7 @@ def legacy_builder(cfg, megatron_mimo_infra):
13371342
infra.topology = Mock()
13381343
infra.module_output_ndim = {"language": 3}
13391344
infra.pg_collections = {"language": Mock()}
1340-
cfg.model.build_infra.return_value = infra
1341-
cfg.model.provide_distributed_model.return_value = [Mock()]
1345+
model = Mock()
13421346

13431347
state = Mock()
13441348
state.cfg = cfg
@@ -1347,11 +1351,12 @@ def legacy_builder(cfg, megatron_mimo_infra):
13471351

13481352
mock_optimizer = MagicMock()
13491353
mock_optimizer.module_infos = {}
1354+
start_time_tensor = Mock()
1355+
start_time_tensor.item.return_value = state.start_time
13501356

13511357
with (
13521358
patch("megatron.bridge.training.setup_megatron_mimo.dist") as m_dist,
1353-
patch("megatron.bridge.training.setup_megatron_mimo._set_megatron_mimo_random_seeds"),
1354-
patch("megatron.bridge.training.setup_megatron_mimo.validate_no_stub_ranks"),
1359+
patch("megatron.bridge.models.megatron_mimo.build_megatron_mimo_model", return_value=(model, infra)),
13551360
patch("megatron.bridge.training.setup_megatron_mimo.build_pg_collection_for_schedule"),
13561361
patch("megatron.bridge.training.setup_megatron_mimo.get_module_to_grid_tuple"),
13571362
patch("megatron.bridge.training.setup_megatron_mimo.MultiModulePipelineCommunicator"),
@@ -1376,6 +1381,7 @@ def legacy_builder(cfg, megatron_mimo_infra):
13761381
patch("megatron.core.parallel_state._TENSOR_MODEL_PARALLEL_GROUP", None),
13771382
patch("megatron.core.parallel_state._DATA_PARALLEL_GROUP", None),
13781383
patch("megatron.core.parallel_state._DATA_PARALLEL_GROUP_WITH_CP", None),
1384+
patch("torch.tensor", return_value=start_time_tensor),
13791385
):
13801386
m_dist.get_rank.return_value = 0
13811387
m_dist.get_world_size.return_value = 2

‎tests/unit_tests/training/megatron_mimo/test_pretrain_megatron_mimo.py‎

Lines changed: 15 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -44,14 +44,10 @@ def _make_setup_output(module_to_grid_map):
4444
)
4545

4646

47-
@patch(
48-
"megatron.bridge.training.setup_megatron_mimo.is_current_rank_in_grid",
49-
side_effect=lambda grid: grid.rank_offset <= 4 < (grid.rank_offset + grid.size),
50-
)
51-
@patch("megatron.bridge.training.setup_megatron_mimo.dist")
52-
def test_set_megatron_mimo_random_seeds_calls_model_parallel_cuda_manual_seed(mock_dist, _mock_in_grid):
53-
"""_set_megatron_mimo_random_seeds should derive TP/PP ranks from grids and call model_parallel_cuda_manual_seed."""
54-
from megatron.bridge.training.setup_megatron_mimo import _set_megatron_mimo_random_seeds
47+
@patch("megatron.bridge.models.megatron_mimo.build_model.dist")
48+
def test_set_megatron_mimo_random_seeds_calls_model_parallel_cuda_manual_seed(mock_dist):
49+
"""_set_per_module_random_seeds should derive TP/PP ranks from grids and call model_parallel_cuda_manual_seed."""
50+
from megatron.bridge.models.megatron_mimo.build_model import _set_per_module_random_seeds
5551

5652
mock_dist.get_rank.return_value = 4 # e.g. first rank of vision encoder
5753

@@ -63,29 +59,25 @@ def test_set_megatron_mimo_random_seeds_calls_model_parallel_cuda_manual_seed(mo
6359
grid = MagicMock()
6460
grid.rank_offset = 4
6561
grid.size = 4
62+
grid.is_current_rank_in_grid.return_value = True
6663
grid.get_pg.side_effect = lambda dims: {"tp": tp_pg, "pp": pp_pg}[dims[0]]
6764

6865
megatron_mimo_infra = SimpleNamespace(module_to_grid_map={"vision": grid})
69-
cfg = SimpleNamespace(rng=SimpleNamespace(seed=42))
7066

71-
with patch("megatron.core.tensor_parallel.model_parallel_cuda_manual_seed") as mock_seed:
67+
with patch("megatron.bridge.models.megatron_mimo.build_model.tensor_parallel.model_parallel_cuda_manual_seed") as mock_seed:
7268
import torch
7369

7470
with patch.object(torch.cuda, "device_count", return_value=1):
75-
_set_megatron_mimo_random_seeds(cfg, megatron_mimo_infra)
71+
_set_per_module_random_seeds(megatron_mimo_infra, seed=42)
7672

7773
# pp_rank=0, so seed stays 42. tp_rank=0 passed explicitly.
7874
mock_seed.assert_called_once_with(42, tp_rank=0, ep_rank=0, etp_rank=0)
7975

8076

81-
@patch(
82-
"megatron.bridge.training.setup_megatron_mimo.is_current_rank_in_grid",
83-
side_effect=lambda grid: grid.rank_offset <= 2 < (grid.rank_offset + grid.size),
84-
)
85-
@patch("megatron.bridge.training.setup_megatron_mimo.dist")
86-
def test_set_megatron_mimo_random_seeds_offsets_by_pp_rank(mock_dist, _mock_in_grid):
77+
@patch("megatron.bridge.models.megatron_mimo.build_model.dist")
78+
def test_set_megatron_mimo_random_seeds_offsets_by_pp_rank(mock_dist):
8779
"""PP rank > 0 should offset the seed by 100 * pp_rank."""
88-
from megatron.bridge.training.setup_megatron_mimo import _set_megatron_mimo_random_seeds
80+
from megatron.bridge.models.megatron_mimo.build_model import _set_per_module_random_seeds
8981

9082
mock_dist.get_rank.return_value = 2
9183

@@ -97,16 +89,16 @@ def test_set_megatron_mimo_random_seeds_offsets_by_pp_rank(mock_dist, _mock_in_g
9789
grid = MagicMock()
9890
grid.rank_offset = 0
9991
grid.size = 4
92+
grid.is_current_rank_in_grid.return_value = True
10093
grid.get_pg.side_effect = lambda dims: {"tp": tp_pg, "pp": pp_pg}[dims[0]]
10194

10295
megatron_mimo_infra = SimpleNamespace(module_to_grid_map={"llm": grid})
103-
cfg = SimpleNamespace(rng=SimpleNamespace(seed=42))
10496

105-
with patch("megatron.core.tensor_parallel.model_parallel_cuda_manual_seed") as mock_seed:
97+
with patch("megatron.bridge.models.megatron_mimo.build_model.tensor_parallel.model_parallel_cuda_manual_seed") as mock_seed:
10698
import torch
10799

108100
with patch.object(torch.cuda, "device_count", return_value=1):
109-
_set_megatron_mimo_random_seeds(cfg, megatron_mimo_infra)
101+
_set_per_module_random_seeds(megatron_mimo_infra, seed=42)
110102

111103
# seed = 42 + 100 * 1 = 142, tp_rank=1
112104
mock_seed.assert_called_once_with(142, tp_rank=1, ep_rank=0, etp_rank=0)
@@ -246,12 +238,10 @@ def test_setup_megatron_mimo_asserts_when_constructor_fields_missing(
246238
mock_infra.module_to_grid_map = {"language": MagicMock()}
247239
mock_infra.topology = {"language": []}
248240
mock_infra.module_output_ndim = {"language": 3}
249-
cfg.model.build_infra.return_value = mock_infra
250-
cfg.model.provide_distributed_model.return_value = [MagicMock()]
241+
model = MagicMock()
251242

252243
with (
253-
patch("megatron.bridge.training.setup_megatron_mimo.validate_no_stub_ranks"),
254-
patch("megatron.bridge.training.setup_megatron_mimo._set_megatron_mimo_random_seeds"),
244+
patch("megatron.bridge.models.megatron_mimo.build_megatron_mimo_model", return_value=(model, mock_infra)),
255245
patch("megatron.bridge.training.setup_megatron_mimo.build_pg_collection_for_schedule"),
256246
patch("megatron.bridge.training.setup_megatron_mimo.get_module_to_grid_tuple"),
257247
patch("megatron.bridge.training.setup_megatron_mimo.MultiModulePipelineCommunicator"),

0 commit comments

Comments
 (0)