Skip to content

Commit 938967f

Browse files
committed
fix test
Signed-off-by: HuiyingLi <[email protected]>
1 parent 7abd810 commit 938967f

File tree

1 file changed

+11
-1
lines changed

1 file changed

+11
-1
lines changed

tests/unit_tests/recipes/test_train_ft.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,12 @@
1414

1515
import logging
1616
from unittest.mock import MagicMock, patch
17-
from nemo_automodel.recipes.llm.train_ft import build_validation_dataloader, build_dataloader, build_model_and_optimizer
17+
from nemo_automodel.recipes.llm.train_ft import (
18+
TrainFinetuneRecipeForNextTokenPrediction,
19+
build_validation_dataloader,
20+
build_dataloader,
21+
build_model_and_optimizer,
22+
)
1823
from nemo_automodel.components.config.loader import ConfigNode
1924
from unittest.mock import patch
2025
import importlib
@@ -468,6 +473,11 @@ def _patch_setup_minimals(monkeypatch, patch_mock: MagicMock):
468473
"nemo_automodel.recipes.llm.train_ft.build_checkpoint_config",
469474
lambda *a, **k: SimpleNamespace(checkpoint_dir="ckpts", model_state_dict_keys=None),
470475
)
476+
# Avoid requiring a distributed _target_
477+
monkeypatch.setattr(
478+
"nemo_automodel.components.config.loader.ConfigNode.instantiate",
479+
lambda self, *a, **k: SimpleNamespace(pp_size=0, device_mesh=None, moe_mesh=None),
480+
)
471481

472482
# Stub Checkpointer
473483
monkeypatch.setattr(

0 commit comments

Comments
 (0)