Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .github/workflows/verl_ci/run_megatron_deepseek.sh
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ trap cleanup EXIT
pip3 install git+https://github.com/NVIDIA-NeMo/Megatron-Bridge.git@main --no-deps --no-build-isolation
pip3 install git+https://github.com/NVIDIA/Megatron-LM.git@main --no-deps --no-build-isolation
pip3 install "nvidia-modelopt[torch]>=0.37.0"
pip3 install "transformers==5.12.1"
## cd to verl checkout root
pip3 install -r requirements-test.txt
pip3 install -r requirements.txt
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ trap cleanup EXIT
pip3 install git+https://github.com/NVIDIA-NeMo/Megatron-Bridge.git@main --no-deps --no-build-isolation
pip3 install git+https://github.com/NVIDIA/Megatron-LM.git@main --no-deps --no-build-isolation
pip3 install "nvidia-modelopt[torch]>=0.37.0"
pip3 install "transformers==5.12.1"
## cd to verl checkout root
pip3 install -r requirements-test.txt
pip3 install -r requirements.txt
Expand Down
1 change: 1 addition & 0 deletions .github/workflows/verl_ci/run_megatron_qwen3.sh
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ trap cleanup EXIT
pip3 install git+https://github.com/NVIDIA-NeMo/Megatron-Bridge.git@main --no-deps --no-build-isolation
pip3 install git+https://github.com/NVIDIA/Megatron-LM.git@main --no-deps --no-build-isolation
pip3 install "nvidia-modelopt[torch]>=0.37.0"
pip3 install "transformers==5.12.1"
## cd to verl checkout root
pip3 install -r requirements-test.txt
pip3 install -r requirements.txt
Expand Down
3 changes: 2 additions & 1 deletion src/megatron/bridge/data/builders/energon.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
from transformers import AutoProcessor, AutoTokenizer, Qwen3VLProcessor

from megatron.bridge.data.base import DataloaderConfig, DatasetBuildContext, validate_declarative_mapping
from megatron.bridge.data.energon.base_energon_datamodule import EnergonMultiModalDataModule
from megatron.bridge.models.hf_pretrained.utils import is_safe_repo


Expand Down Expand Up @@ -289,6 +288,8 @@ def __init__(self, config: EnergonDatasetConfig) -> None:

def build(self, context: DatasetBuildContext) -> tuple[Any | None, Any | None, None]:
"""Build requested Energon train and validation iterators."""
from megatron.bridge.data.energon.base_energon_datamodule import EnergonMultiModalDataModule

assert self.config.path is not None
build_train = context.train_samples > 0
build_validation = self.config.do_validation and context.valid_samples > 0
Expand Down
6 changes: 4 additions & 2 deletions tests/unit_tests/data/builders/test_energon_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
QwenVLEnergonTaskEncoderConfig,
build_energon_task_encoder,
)
from megatron.bridge.data.energon import base_energon_datamodule
from megatron.bridge.training.config import ConfigContainer
from megatron.bridge.training.utils.omegaconf_utils import process_config_with_overrides

Expand Down Expand Up @@ -226,7 +227,7 @@ def test_builder_honors_requested_splits_and_reuses_runtime_encoder(monkeypatch:
datamodule.val_dataloader.return_value = ["validation"]
datamodule_cls = MagicMock(return_value=datamodule)
monkeypatch.setattr("megatron.bridge.data.builders.energon.build_energon_task_encoder", build_encoder)
monkeypatch.setattr("megatron.bridge.data.builders.energon.EnergonMultiModalDataModule", datamodule_cls)
monkeypatch.setattr(base_energon_datamodule, "EnergonMultiModalDataModule", datamodule_cls)

train, validation, test = EnergonDatasetBuilder(config).build(
DatasetBuildContext(train_samples=10, valid_samples=5, test_samples=3)
Expand All @@ -248,7 +249,8 @@ def test_builder_skips_unrequested_validation(monkeypatch: pytest.MonkeyPatch):
datamodule.train_dataloader.return_value = []
monkeypatch.setattr("megatron.bridge.data.builders.energon.build_energon_task_encoder", lambda _: object())
monkeypatch.setattr(
"megatron.bridge.data.builders.energon.EnergonMultiModalDataModule",
base_energon_datamodule,
"EnergonMultiModalDataModule",
MagicMock(return_value=datamodule),
)

Expand Down
15 changes: 15 additions & 0 deletions tests/unit_tests/data/test_module_layout.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.

import importlib
import subprocess
import sys

import pytest

Expand All @@ -23,3 +25,16 @@
def test_removed_internal_data_modules_have_no_compatibility_shims(module_name):
with pytest.raises(ModuleNotFoundError):
importlib.import_module(module_name)


def test_training_config_import_does_not_require_energon():
code = "import sys; sys.modules['megatron.energon'] = None; import megatron.bridge.training.config"
result = subprocess.run(
[sys.executable, "-c", code],
capture_output=True,
text=True,
timeout=180,
check=False,
)

assert result.returncode == 0, result.stderr or result.stdout
Loading