build: update PyTorch base image to 26.09 (#6258)

Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>
Signed-off-by: Ajay Balasa <abalasa@nvidia.com>
Signed-off-by: Ajay <abalasa@nvidia.com>
Co-authored-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>

Author: balasaajay

.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

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" \

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",

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",
     }

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

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)