diff --git a/.github/actions/test-template/action.yml b/.github/actions/test-template/action.yml index 57281bd0f9..2d7779c624 100644 --- a/.github/actions/test-template/action.yml +++ b/.github/actions/test-template/action.yml @@ -177,6 +177,9 @@ runs: set -e export GH_TOKEN=${{ inputs.github-token }} + # Keep the FlashAttention generation policy scoped to CI tests. + export NVTE_FLASH_ATTN_V2=1 NVTE_FLASH_ATTN_V3=0 NVTE_FLASH_ATTN_V4=0 + if [[ "${{ contains(inputs.runner, 'gcp') }}" == "true" ]]; then export NCCL_ENV_PLUGIN=none export NCCL_NET_PLUGIN=none diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index f746f9e291..81f6b9234c 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:26.08-py3 +ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:26.09-py3 FROM ${BASE_IMAGE} AS base ENV PATH="/opt/venv/bin:/usr/local/bin:$PATH" \ diff --git a/docker/common/fw_pyproject.toml b/docker/common/fw_pyproject.toml index 01a8efd965..aeb4ee5728 100644 --- a/docker/common/fw_pyproject.toml +++ b/docker/common/fw_pyproject.toml @@ -61,7 +61,9 @@ override-dependencies = [ "pyzmq==25.1.2", # @URL: https://github.com/NVIDIA/NeMo/issues/13625 "tiktoken>=0.9.0", # let megatron-hub win over nemo-toolkit "datasets==4.8.5", - "pyarrow>=21.0.0", # datasets 4.8.5 needs >=21; override export-deploy's <21 cap (not exercised by training) + # datasets 4.8.5 needs >=21, overriding export-deploy's <21 cap. + # PyArrow 26 requires NumPy 2; preserve MBridge's NumPy 1 runtime. + "pyarrow>=21.0.0,<26", "fsspec[http]>=2023.1.0,<=2024.9.0", "transformers; sys_platform == 'never'", # Defer to transformers in MBridge install "packaging>=25.0", diff --git a/tests/unit_tests/ci/test_test_template.py b/tests/unit_tests/ci/test_test_template.py index 383c22f2b3..4a6847926a 100644 --- a/tests/unit_tests/ci/test_test_template.py +++ b/tests/unit_tests/ci/test_test_template.py @@ -54,7 +54,8 @@ def test_container_startup_cannot_override_gcp_test_environment( "python3 - <<'PROBE'\n" "import json, os\n" "from pathlib import Path\n" - "names = ('NCCL_ENV_PLUGIN', 'NCCL_NET_PLUGIN', 'NCCL_PROFILER_PLUGIN', 'NCCL_NET', 'STARTUP_WAS_RUN')\n" + "names = ('NCCL_ENV_PLUGIN', 'NCCL_NET_PLUGIN', 'NCCL_PROFILER_PLUGIN', 'NCCL_NET', 'STARTUP_WAS_RUN', " + "'NVTE_FLASH_ATTN_V2', 'NVTE_FLASH_ATTN_V3', 'NVTE_FLASH_ATTN_V4')\n" "Path('observed.json').write_text(json.dumps({key: os.environ.get(key) for key in names}))\n" "PROBE\n" f"exit {test_exit_code}\n" @@ -65,6 +66,9 @@ def test_container_startup_cannot_override_gcp_test_environment( "BASH_ENV": str(startup), "GITHUB_OUTPUT": str(tmp_path / "github-output"), "NCCL_NET": "inherited-network", + "NVTE_FLASH_ATTN_V2": "0", + "NVTE_FLASH_ATTN_V3": "1", + "NVTE_FLASH_ATTN_V4": "1", } subprocess.run(["bash", "-c", create], cwd=tmp_path, env=env, check=True, capture_output=True, text=True) result = subprocess.run(["bash", "job.sh"], cwd=tmp_path, env=env, capture_output=True, text=True) @@ -77,4 +81,7 @@ def test_container_startup_cannot_override_gcp_test_environment( "NCCL_PROFILER_PLUGIN": plugin, "NCCL_NET": "Socket" if "gcp" in runner else "inherited-network", "STARTUP_WAS_RUN": "yes", + "NVTE_FLASH_ATTN_V2": "1", + "NVTE_FLASH_ATTN_V3": "0", + "NVTE_FLASH_ATTN_V4": "0", } diff --git a/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_model.py b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_model.py index 83fecd96d9..c0362a5dc0 100644 --- a/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_model.py +++ b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_model.py @@ -1261,6 +1261,8 @@ def test_real_radio_image_forward_with_collator_owned_cp1_packing( assert torch.isfinite(output).all() +# Dump all Python thread stacks before the CI wrapper terminates a stalled backward pass. +@pytest.mark.timeout(300, method="thread") @pytest.mark.run_only_on("GPU") def test_real_packed_multimodal_optimizer_step(single_rank_model_parallel): del single_rank_model_parallel diff --git a/tests/unit_tests/models/test_attention_backend_config.py b/tests/unit_tests/models/test_attention_backend_config.py index b067bd4aba..903324aadb 100644 --- a/tests/unit_tests/models/test_attention_backend_config.py +++ b/tests/unit_tests/models/test_attention_backend_config.py @@ -21,6 +21,7 @@ import pytest from megatron.core.models.common.language_module.language_module import LanguageModule +from megatron.core.transformer import attention as attention_module from megatron.core.transformer import utils as transformer_utils from megatron.core.transformer.enums import AttnBackend @@ -33,19 +34,21 @@ CONFIG_TYPES = (TransformerConfig, MLATransformerConfig, HeterogeneousTransformerConfig) BACKEND_ENV = ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN") +FLASH_VERSION_ENV = ("NVTE_FLASH_ATTN_V2", "NVTE_FLASH_ATTN_V3", "NVTE_FLASH_ATTN_V4") pytestmark = pytest.mark.unit @pytest.fixture(autouse=True) def isolated_attention_environment(): with patch.dict(os.environ): - for name in BACKEND_ENV: + for name in (*BACKEND_ENV, *FLASH_VERSION_ENV): os.environ.pop(name, None) yield -def make_config(config_type, backend): +def make_config(config_type, backend, **overrides): kwargs = dict(num_layers=2, hidden_size=64, num_attention_heads=4, attention_backend=backend) + kwargs.update(overrides) if config_type is HeterogeneousTransformerConfig: block = { "attention": {"no_op": False, "replace_with_linear": False, "num_query_groups": 4}, @@ -63,6 +66,63 @@ def configure_mcore_attention(config): LanguageModule._set_attention_backend(SimpleNamespace(config=config)) +@pytest.mark.parametrize("config_type", CONFIG_TYPES) +@pytest.mark.parametrize("backend", list(AttnBackend)) +@pytest.mark.skipif( + not hasattr(attention_module.Attention, "_resolve_flash_version"), + reason="The frozen MCore dev pin predates FlashAttention version controls", +) +def test_explicit_fa2_disables_fa4_without_changing_backend(config_type, backend, monkeypatch): + monkeypatch.setenv("NVTE_FLASH_ATTN_V4", "0") + config = make_config(config_type, backend, flash_attention_version=2) + + config.finalize() + configure_mcore_attention(config) + + assert config.flash_attention_version == 2 + assert config.attention_backend is backend + assert tuple(os.environ[name] for name in FLASH_VERSION_ENV) == ("1", "0", "0") + + +@pytest.mark.parametrize("config_type", CONFIG_TYPES) +@pytest.mark.skipif( + not hasattr(attention_module.Attention, "_resolve_flash_version"), + reason="The frozen MCore dev pin predates FlashAttention version controls", +) +def test_direct_attention_uses_explicit_fa2_even_when_fa4_is_installed(config_type, monkeypatch): + monkeypatch.setattr(attention_module, "HAVE_FA4", True) + monkeypatch.setattr(attention_module, "HAVE_FA3", True) + config = make_config(config_type, AttnBackend.auto, flash_attention_version=2) + config.finalize() + + attention = SimpleNamespace(flash_attention_version=config.flash_attention_version) + assert attention_module.Attention._resolve_flash_version(attention) == (False, False) + + +@pytest.mark.parametrize("config_type", CONFIG_TYPES) +@pytest.mark.parametrize("version", [None, 2, 3, 4]) +@pytest.mark.skipif( + not hasattr(attention_module.Attention, "_resolve_flash_version"), + reason="The frozen MCore dev pin predates FlashAttention version controls", +) +def test_explicit_flash_version_is_preserved(config_type, version): + config = make_config(config_type, AttnBackend.auto, flash_attention_version=version) + + config.finalize() + + assert config.flash_attention_version == version + + +@pytest.mark.parametrize("config_type", CONFIG_TYPES) +def test_default_flash_version_is_inherited_from_mcore(config_type): + config = make_config(config_type, AttnBackend.auto) + + config.finalize() + + assert getattr(config, "flash_attention_version", None) is None + assert all(name not in os.environ for name in FLASH_VERSION_ENV) + + @pytest.mark.parametrize("config_type", CONFIG_TYPES) def test_unset_backend_finalizes_to_mcore_auto(config_type): config = make_config(config_type, None)