From 8f160254d16ba8da36e20da6e7308463c72e7799 Mon Sep 17 00:00:00 2001 From: svcnemo-autobot Date: Tue, 29 Sep 2026 05:08:39 +0000 Subject: [PATCH 01/14] build: update PyTorch base image to 26.09 Signed-off-by: svcnemo-autobot Signed-off-by: Ajay Balasa --- docker/Dockerfile.ci | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index e9dcc4d12a..f550d8c022 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 megatron_bridge WORKDIR /opt/Megatron-Bridge From 837960141620fc50090d3fe946faf16c0a936596 Mon Sep 17 00:00:00 2001 From: Ajay Balasa Date: Wed, 30 Sep 2026 15:47:53 +0000 Subject: [PATCH 02/14] fix(training): count full dense sequences in context-parallel FLOPs Signed-off-by: Ajay Balasa --- src/megatron/bridge/training/gpt_step.py | 13 +++++--- tests/unit_tests/training/test_gpt_step.py | 39 ++++++++++++++++++++++ 2 files changed, 47 insertions(+), 5 deletions(-) diff --git a/src/megatron/bridge/training/gpt_step.py b/src/megatron/bridge/training/gpt_step.py index 3ff54efc2f..4024db6125 100644 --- a/src/megatron/bridge/training/gpt_step.py +++ b/src/megatron/bridge/training/gpt_step.py @@ -524,11 +524,14 @@ def _forward_step_common( ) timers("batch-generator").stop() - # Packed CP partitions tokens but preserves full-sequence cu_seqlens. - # Restore the physical token count for linear FLOPS; attention uses the - # unchanged sub-sequence boundaries. train.py reduces over pure DP, not CP. + # Both dense and packed tokens are CP-sharded. Recover the actual full batch + # length instead of relying on seq_length, which the builder's transformer + # config omits and which may exceed a dynamic batch's length. Packed attention + # still uses full-sequence cu_seqlens. train.py reduces over pure DP, not CP. # A post-process-only PP stage has labels but no input token ids. flops_tokens = tokens if tokens is not None else labels + cp_size = pg_collection.cp.size() + full_seq_length = flops_tokens.shape[1] * cp_size if flops_tokens is not None else None cu_seqlens = None cu_seqlens_argmin = None cu_seqlens_unpadded = None @@ -548,8 +551,8 @@ def _forward_step_common( state, flops_tokens, vp_stage=vp_stage, - config_seq_len=getattr(config, "seq_length", None), - context_parallel_size=pg_collection.cp.size(), + config_seq_len=full_seq_length, + context_parallel_size=cp_size, cu_seqlens=cu_seqlens, cu_seqlens_argmin=cu_seqlens_argmin, cu_seqlens_unpadded=cu_seqlens_unpadded, diff --git a/tests/unit_tests/training/test_gpt_step.py b/tests/unit_tests/training/test_gpt_step.py index 2ae725bf1a..23d981e931 100644 --- a/tests/unit_tests/training/test_gpt_step.py +++ b/tests/unit_tests/training/test_gpt_step.py @@ -13,6 +13,7 @@ # limitations under the License. from functools import partial +from types import SimpleNamespace from unittest.mock import MagicMock, Mock, patch import modelopt.torch.distill as mtd @@ -20,6 +21,7 @@ import torch from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.transformer.moe.router import TopKRouter +from megatron.core.transformer.transformer_config import TransformerConfig from megatron.bridge.training.gpt_step import ( _create_loss_function_modelopt, @@ -741,6 +743,43 @@ def test_forward_common_passes_loss_mask_to_model( data_iterator, state.cfg, mtp_num_layers > 0, pg_collection=pg_collection, vp_stage=None ) + @pytest.mark.unit + @pytest.mark.parametrize("cp_size", [1, 2, 8]) + @pytest.mark.parametrize("full_seq_length", [16, 32]) + @pytest.mark.parametrize("has_input_ids", [False, True]) + @pytest.mark.parametrize("vp_stage", [None, 1]) + def test_forward_common_counts_dense_flops_with_nested_model_config( + self, monkeypatch, cp_size, full_seq_length, has_input_ids, vp_stage + ): + """Builder configs omit seq_length; count actual dense CP shards, including shorter batches.""" + tokens = torch.arange(full_seq_length // cp_size).repeat(2, 1) + labels = tokens + 1 + loss_mask = torch.ones_like(tokens, dtype=torch.float32) + model = _RecordingModel(vp_stage=vp_stage) + model.config = TransformerConfig(num_layers=2, hidden_size=32, num_attention_heads=4) + assert not hasattr(model.config, "seq_length") + state = SimpleNamespace( + cfg=SimpleNamespace(model=SimpleNamespace(seq_length=128)), + timers=_NoopTimer(), + straggler_timer=_NoopTimer(), + _flops_seqlen_sum=0, + _flops_seqlen_sq_sum=0, + _flops_requires_global_reduce=False, + ) + monkeypatch.setattr( + "megatron.bridge.training.gpt_step.get_pg_collection", lambda model: _MockPGCollection(cp_size=cp_size) + ) + get_batch_mock = Mock(return_value=(tokens if has_input_ids else None, labels, loss_mask, None, None, None)) + + _forward_step_common(state, _Iterator({}), model, _get_batch_fn=get_batch_mock) + + expected_batches = 2 if vp_stage is None else 0 + assert state._flops_seqlen_sum == expected_batches * full_seq_length + assert state._flops_seqlen_sq_sum == expected_batches * full_seq_length**2 + assert not state._flops_requires_global_reduce + assert model.forward_kwargs["input_ids"] is (tokens if has_input_ids else None) + assert model.forward_kwargs["labels"] is labels + @pytest.mark.unit @pytest.mark.parametrize("cp_size", [1, 2, 8]) @pytest.mark.parametrize("metadata_format", ["current", "legacy"]) From efee69596cdaaa266c9f5bcb6ab709028e3559bb Mon Sep 17 00:00:00 2001 From: Ajay Balasa Date: Wed, 30 Sep 2026 15:58:21 +0000 Subject: [PATCH 03/14] docs: include the formal PR review rubric in navigation Signed-off-by: Ajay Balasa --- docs/skills-index.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/skills-index.md b/docs/skills-index.md index 49b5461b13..ec36d9ba3e 100644 --- a/docs/skills-index.md +++ b/docs/skills-index.md @@ -14,6 +14,7 @@ skills/bump-dependency/SKILL skills/testing/SKILL skills/cicd/SKILL skills/review-pr/SKILL +Formal PR Review Rubric skills/nemo-mbridge-mlm-bridge-training/SKILL skills/nemo-mbridge-recipe-recommender/SKILL ``` From 4d62eb8dbe28c46f5f0e3dfb4e8a5d1eecc0c5b5 Mon Sep 17 00:00:00 2001 From: Ajay Balasa Date: Fri, 9 Oct 2026 08:10:35 +0000 Subject: [PATCH 04/14] test(nemotron-omni): capture stalled backward thread stacks Signed-off-by: Ajay Balasa --- .../unit_tests/models/nemotron_omni/test_nemotron_omni_model.py | 2 ++ 1 file changed, 2 insertions(+) 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 From c1017130c09e909e3d01aa9b96e74be3edd25f1b Mon Sep 17 00:00:00 2001 From: Ajay Balasa Date: Fri, 9 Oct 2026 11:38:17 +0000 Subject: [PATCH 05/14] build: align FlashAttention with the 26.09 CUTLASS runtime Signed-off-by: Ajay Balasa --- docker/Dockerfile.ci | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 81f6b9234c..f266b24a63 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -244,7 +244,11 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \ # its CUTLASS-backed kernels. if [ "$MBRIDGE_PRESERVE_BASE_RUNTIME" != "True" ]; then \ pip uninstall -y nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base nvidia-cutlass-dsl-libs-cu13 && \ - pip install nvidia-cutlass-dsl[cu13]==4.8.0; \ + pip install nvidia-cutlass-dsl[cu13]==4.8.0 && \ + # FA4 includes the single-tile scheduler deadlock fix from Dao-AILab/flash-attention#2705. + # Keep FA4 and Quack compatible with the installed CUTLASS 4.8 runtime. + uv pip install --python /usr/bin/python3 --break-system-packages --no-deps \ + flash-attn-4==4.0.0b34 quack-kernels==0.6.5; \ fi && \ if [ "$MCORE_TRIGGERED_TESTING" = "true" ]; then \ echo "⚙️ MCore testing mode: syncing the dispatched dependency delta"; \ From 541b945ccda6953623c721b9af8dec405ff74e0b Mon Sep 17 00:00:00 2001 From: Ajay Balasa Date: Fri, 9 Oct 2026 13:30:09 +0000 Subject: [PATCH 06/14] build: pin MCore RADIO packed sequence length fix Signed-off-by: Ajay Balasa --- 3rdparty/Megatron-LM | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index 5052ce8c48..38a2399e6d 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit 5052ce8c48d999a9422ad5ef02954089639db832 +Subproject commit 38a2399e6d3811b4455e50ae199ffbf6ce60e8df From d9500bb6ce5dd16ca0a401fc5d99daecf8afd5c6 Mon Sep 17 00:00:00 2001 From: Ajay Balasa Date: Fri, 9 Oct 2026 15:02:59 +0000 Subject: [PATCH 07/14] build(mcore): synchronize the pinned main revision Signed-off-by: Ajay Balasa --- .main.commit | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.main.commit b/.main.commit index 2708f79afc..7b1b7092cd 100644 --- a/.main.commit +++ b/.main.commit @@ -1 +1 @@ -441a987409b8fdd2b90b463f7999bac970109a43 +38a2399e6d3811b4455e50ae199ffbf6ce60e8df From d8df76d823fc55226850ea641a194728c88129f1 Mon Sep 17 00:00:00 2001 From: Ajay Balasa Date: Fri, 9 Oct 2026 16:12:56 +0000 Subject: [PATCH 08/14] fix(training): keep Omni vision sequence maxima as integers Return host integer maxima from the training-step vision metadata builder. Non-temporal RADIO preserves this metadata through class-token insertion, and FA4 rejects scalar tensor maxima. Keep cumulative offsets on device. Cover square and ragged image sizes plus packed and unpacked first-stage batches with strict metadata type checks. Signed-off-by: Ajay Balasa --- .../bridge/training/nemotron_omni_step.py | 2 +- .../training/test_nemotron_omni_step.py | 44 +++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/src/megatron/bridge/training/nemotron_omni_step.py b/src/megatron/bridge/training/nemotron_omni_step.py index 2e0dfd4b1b..e0732fb680 100644 --- a/src/megatron/bridge/training/nemotron_omni_step.py +++ b/src/megatron/bridge/training/nemotron_omni_step.py @@ -158,7 +158,7 @@ def _build_vision_packed_seq_params( cu.append(cu[-1] + sl) device = imgs_sizes.device if torch.is_tensor(imgs_sizes) else torch.device("cpu") cu_tensor = torch.tensor(cu, dtype=torch.int32, device=device) - max_len = torch.tensor(max(seq_lens) if seq_lens else 0, dtype=torch.int32, device=device) + max_len = max(seq_lens, default=0) return PackedSeqParams( qkv_format="thd", cu_seqlens_q=cu_tensor, diff --git a/tests/unit_tests/training/test_nemotron_omni_step.py b/tests/unit_tests/training/test_nemotron_omni_step.py index c6911dbe55..ce32480228 100644 --- a/tests/unit_tests/training/test_nemotron_omni_step.py +++ b/tests/unit_tests/training/test_nemotron_omni_step.py @@ -22,6 +22,29 @@ from megatron.bridge.training.utils.packed_seq_utils import get_packed_seq_params +@pytest.mark.unit +@pytest.mark.parametrize( + "image_sizes,expected_boundaries,expected_max", + [ + ([[512, 512]], [0, 1024], 1024), + ([[32, 48], [64, 32]], [0, 6, 14], 8), + ], +) +def test_vision_packed_metadata_uses_integer_maxima(image_sizes, expected_boundaries, expected_max): + image_sizes = torch.tensor(image_sizes, dtype=torch.int64) + + metadata = nemotron_omni_step._build_vision_packed_seq_params(image_sizes) + + assert metadata.qkv_format == "thd" + assert type(metadata.max_seqlen_q) is int + assert type(metadata.max_seqlen_kv) is int + assert metadata.max_seqlen_q == metadata.max_seqlen_kv == expected_max + for boundaries in (metadata.cu_seqlens_q, metadata.cu_seqlens_kv): + assert boundaries.tolist() == expected_boundaries + assert boundaries.dtype == torch.int32 + assert boundaries.device == image_sizes.device + + def test_batch_moves_only_one_compatible_token_alias_to_cuda(monkeypatch): cuda_inputs = [] @@ -122,6 +145,27 @@ def _pipeline_cfg(*, packed=True, defer_packing=False, temporal_patch_dim=1): ) +@pytest.mark.unit +@pytest.mark.parametrize("packed", [False, True]) +def test_first_pipeline_batch_keeps_vision_maxima_as_integers(monkeypatch, packed): + monkeypatch.setattr(torch.Tensor, "cuda", lambda self, **kwargs: self) + monkeypatch.setattr(nemotron_omni_step, "is_pp_first_stage", lambda group: True) + monkeypatch.setattr(nemotron_omni_step, "is_pp_last_stage", lambda group: False) + batch = _packed_pipeline_batch() + if not packed: + for key in (*nemotron_omni_step._PACKED_SEQ_PARAM_KEYS, "padding_mask"): + batch.pop(key, None) + + result = get_batch(iter([batch]), _pipeline_cfg(packed=packed), pg_collection=SimpleNamespace(pp=object())) + + vision_metadata = result[12] + assert type(vision_metadata.max_seqlen_q) is int + assert type(vision_metadata.max_seqlen_kv) is int + assert vision_metadata.max_seqlen_q == vision_metadata.max_seqlen_kv == 4 + assert vision_metadata.cu_seqlens_q.tolist() == [0, 4, 8] + assert vision_metadata.cu_seqlens_kv.tolist() == [0, 4, 8] + + def test_forward_rejects_deferred_multimodal_packing(): state = SimpleNamespace( timers=None, From 973efa22089b7f3c5a33fa881eb3de1870f17419 Mon Sep 17 00:00:00 2001 From: Ajay Balasa Date: Fri, 9 Oct 2026 16:25:44 +0000 Subject: [PATCH 09/14] fix(vision): pass integer maxima to packed attention Normalize Qwen3-VL and ERNIE vision sequence maxima to Python integers at the PackedSeqParams boundary. Their scalar tensor maxima are rejected by the CUTLASS host shape operation used by FA4. Preserve cumulative offsets and the existing CUDA-graph path. Add CPU regression coverage for single images and ragged multi-frame grids. Signed-off-by: Ajay Balasa --- .../modeling_ernie45_vl/vision_model.py | 2 +- .../modelling_qwen3_vl/vision_model.py | 2 +- .../models/ernie_vl/test_ernie45_vl_bridge.py | 24 +++++++++++++++ .../modelling_qwen3_vl/test_vision_model.py | 30 +++++++++++++++++++ 4 files changed, 56 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_model.py b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_model.py index dce0f58bb6..7850593584 100644 --- a/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_model.py +++ b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_model.py @@ -267,7 +267,7 @@ def build_packed_seq_params( cu_seqlens = seqlens.cumsum(dim=0) cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0).int() - max_seqlen_q = seqlens.max() + max_seqlen_q = int(seqlens.max().item()) return PackedSeqParams( cu_seqlens_q=cu_seqlens, cu_seqlens_kv=cu_seqlens, diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/vision_model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/vision_model.py index ce9c027261..6afa4144e2 100644 --- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/vision_model.py +++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/vision_model.py @@ -383,7 +383,7 @@ def build_packed_seq_params( cu_seqlens = seqlens.cumsum(dim=0) cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0).int() - max_seqlen_q = seqlens.max() + max_seqlen_q = int(seqlens.max().item()) return PackedSeqParams( cu_seqlens_q=cu_seqlens, cu_seqlens_kv=cu_seqlens, diff --git a/tests/unit_tests/models/ernie_vl/test_ernie45_vl_bridge.py b/tests/unit_tests/models/ernie_vl/test_ernie45_vl_bridge.py index 984d7afeeb..8d1762fe54 100644 --- a/tests/unit_tests/models/ernie_vl/test_ernie45_vl_bridge.py +++ b/tests/unit_tests/models/ernie_vl/test_ernie45_vl_bridge.py @@ -34,6 +34,30 @@ MultiTypeMoeSubmodules, ) from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.model import ErnieMultimodalRotaryEmbedding +from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.vision_model import ErnieVLVisionModel + + +@pytest.mark.unit +@pytest.mark.parametrize( + "grid,expected_boundaries,expected_max", + [ + ([[1, 2, 3]], [0, 6], 6), + ([[2, 2, 3], [1, 2, 4]], [0, 6, 12, 20], 8), + ], +) +def test_vision_packed_metadata_uses_integer_maxima(grid, expected_boundaries, expected_max): + grid = torch.tensor(grid, dtype=torch.int64) + + packed = ErnieVLVisionModel.build_packed_seq_params(None, grid) + + assert packed.qkv_format == "thd" + assert type(packed.max_seqlen_q) is int + assert type(packed.max_seqlen_kv) is int + assert packed.max_seqlen_q == packed.max_seqlen_kv == expected_max + for boundaries in (packed.cu_seqlens_q, packed.cu_seqlens_kv): + assert boundaries.tolist() == expected_boundaries + assert boundaries.dtype == torch.int32 + assert boundaries.device == grid.device def _make_vision_config(): diff --git a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_vision_model.py b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_vision_model.py index 82043ac117..2265958860 100644 --- a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_vision_model.py +++ b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_vision_model.py @@ -95,6 +95,36 @@ def test_raises_when_seq_exceeds_max(self): class TestVisionForwardPackedAttentionSetup: """Tests for ``_vision_forward_packed_attention_setup``.""" + @pytest.mark.unit + @pytest.mark.parametrize( + "grid,expected_boundaries,expected_max", + [ + ([[1, 2, 3]], [0, 6], 6), + ([[2, 2, 3], [1, 2, 4]], [0, 6, 12, 20], 8), + ], + ) + def test_non_cuda_graph_path_uses_integer_maxima(self, grid, expected_boundaries, expected_max): + grid = torch.tensor(grid, dtype=torch.int64) + seq_len = expected_boundaries[-1] + packed, mask = _vision_forward_packed_attention_setup( + use_cuda_graph_padding=False, + hidden_states=torch.zeros(seq_len, 1, 8), + original_seq_len=seq_len, + seq_len=seq_len, + grid_thw=grid, + build_packed_seq_params=lambda value: Qwen3VLVisionModel.build_packed_seq_params(None, value), + ) + + assert mask is None + assert packed.qkv_format == "thd" + assert type(packed.max_seqlen_q) is int + assert type(packed.max_seqlen_kv) is int + assert packed.max_seqlen_q == packed.max_seqlen_kv == expected_max + for boundaries in (packed.cu_seqlens_q, packed.cu_seqlens_kv): + assert boundaries.tolist() == expected_boundaries + assert boundaries.dtype == torch.int32 + assert boundaries.device == grid.device + def test_non_cuda_graph_path_calls_builder_and_no_mask(self): grid = torch.tensor([[1, 4, 4]], dtype=torch.long) hidden = torch.zeros(5, 1, 8, dtype=torch.float32) From fa0b9cb3f67c42065a0e45ab3ee05f1c6bb7f93e Mon Sep 17 00:00:00 2001 From: Ajay Balasa Date: Fri, 9 Oct 2026 17:05:47 +0000 Subject: [PATCH 10/14] fix(container): keep PyArrow compatible with NumPy 1 Constrain the final framework image override below PyArrow 26, which requires NumPy 2 at import time. The Bridge runtime uses NumPy 1.26.4, and the unlocked final sync otherwise upgrades its compatible PyArrow. Retain the lower bound required by Datasets and the existing root lock. Signed-off-by: Ajay Balasa --- docker/common/fw_pyproject.toml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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", From 266efd6f6fad6cb723bf0522e4a792ff239a1aa2 Mon Sep 17 00:00:00 2001 From: Ajay Balasa Date: Fri, 9 Oct 2026 18:17:58 +0000 Subject: [PATCH 11/14] fix(attention): default to FA2 and disable FA4 in CI Signed-off-by: Ajay Balasa --- docker/Dockerfile.ci | 7 +-- .../bridge/models/transformer_config.py | 5 ++ .../models/test_attention_backend_config.py | 50 ++++++++++++++++++- 3 files changed, 55 insertions(+), 7 deletions(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index f266b24a63..62add59beb 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -72,6 +72,7 @@ ARG REINSTALL_NVSHMEM=True ARG PRESERVE_BASE_RUNTIME=False ENV MBRIDGE_PRESERVE_BASE_RUNTIME=${PRESERVE_BASE_RUNTIME} \ + NVTE_FLASH_ATTN_V4=0 \ NVTE_BUILD_NUM_PHILOX_ROUNDS=3 \ NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD=1 \ HYBRID_EP_MULTINODE=1 \ @@ -244,11 +245,7 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \ # its CUTLASS-backed kernels. if [ "$MBRIDGE_PRESERVE_BASE_RUNTIME" != "True" ]; then \ pip uninstall -y nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base nvidia-cutlass-dsl-libs-cu13 && \ - pip install nvidia-cutlass-dsl[cu13]==4.8.0 && \ - # FA4 includes the single-tile scheduler deadlock fix from Dao-AILab/flash-attention#2705. - # Keep FA4 and Quack compatible with the installed CUTLASS 4.8 runtime. - uv pip install --python /usr/bin/python3 --break-system-packages --no-deps \ - flash-attn-4==4.0.0b34 quack-kernels==0.6.5; \ + pip install nvidia-cutlass-dsl[cu13]==4.8.0; \ fi && \ if [ "$MCORE_TRIGGERED_TESTING" = "true" ]; then \ echo "⚙️ MCore testing mode: syncing the dispatched dependency delta"; \ diff --git a/src/megatron/bridge/models/transformer_config.py b/src/megatron/bridge/models/transformer_config.py index 2f9d030599..178a5464c3 100644 --- a/src/megatron/bridge/models/transformer_config.py +++ b/src/megatron/bridge/models/transformer_config.py @@ -20,6 +20,7 @@ import copy from dataclasses import dataclass, field, fields, is_dataclass +from typing import Literal from megatron.core.transformer.enums import AttnBackend from megatron.core.transformer.heterogeneous.heterogeneous_config import ( @@ -168,6 +169,10 @@ class TransformerConfig(MCoreTransformerConfig): _NO_COPY_KEYS = {"_pg_collection"} + # Match the FA2 validation policy in NVIDIA/Megatron-LM#7725. The TE environment + # flag alone does not control MCore's direct inference attention dispatch. + flash_attention_version: Literal[2, 3, 4] | None = 2 + # Generalized tensor-parallel metadata was added after the frozen MCore dev pin. # Keep it on the Bridge wrapper so one configuration remains usable with both pins. tensor_parallel_num_weight_shards: int | None = None diff --git a/tests/unit_tests/models/test_attention_backend_config.py b/tests/unit_tests/models/test_attention_backend_config.py index b067bd4aba..e8dd1551e1 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,49 @@ 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_default_flash_version_disables_fa4_without_changing_backend(config_type, backend, monkeypatch): + monkeypatch.setenv("NVTE_FLASH_ATTN_V4", "0") + config = make_config(config_type, backend) + + 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_defaults_to_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) + 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]) +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_unset_backend_finalizes_to_mcore_auto(config_type): config = make_config(config_type, None) From 6f07fbfb95e70f9e77730412c4f7db0dbb5b2974 Mon Sep 17 00:00:00 2001 From: Ajay Balasa Date: Fri, 9 Oct 2026 19:19:07 +0000 Subject: [PATCH 12/14] fix(ci): scope FlashAttention selection to tests Signed-off-by: Ajay Balasa --- .github/actions/test-template/action.yml | 3 +++ docker/Dockerfile.ci | 1 - .../bridge/models/transformer_config.py | 5 ----- tests/unit_tests/ci/test_test_template.py | 9 +++++++- .../models/test_attention_backend_config.py | 22 +++++++++++++++---- 5 files changed, 29 insertions(+), 11 deletions(-) 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 62add59beb..81f6b9234c 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -72,7 +72,6 @@ ARG REINSTALL_NVSHMEM=True ARG PRESERVE_BASE_RUNTIME=False ENV MBRIDGE_PRESERVE_BASE_RUNTIME=${PRESERVE_BASE_RUNTIME} \ - NVTE_FLASH_ATTN_V4=0 \ NVTE_BUILD_NUM_PHILOX_ROUNDS=3 \ NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD=1 \ HYBRID_EP_MULTINODE=1 \ diff --git a/src/megatron/bridge/models/transformer_config.py b/src/megatron/bridge/models/transformer_config.py index 178a5464c3..2f9d030599 100644 --- a/src/megatron/bridge/models/transformer_config.py +++ b/src/megatron/bridge/models/transformer_config.py @@ -20,7 +20,6 @@ import copy from dataclasses import dataclass, field, fields, is_dataclass -from typing import Literal from megatron.core.transformer.enums import AttnBackend from megatron.core.transformer.heterogeneous.heterogeneous_config import ( @@ -169,10 +168,6 @@ class TransformerConfig(MCoreTransformerConfig): _NO_COPY_KEYS = {"_pg_collection"} - # Match the FA2 validation policy in NVIDIA/Megatron-LM#7725. The TE environment - # flag alone does not control MCore's direct inference attention dispatch. - flash_attention_version: Literal[2, 3, 4] | None = 2 - # Generalized tensor-parallel metadata was added after the frozen MCore dev pin. # Keep it on the Bridge wrapper so one configuration remains usable with both pins. tensor_parallel_num_weight_shards: int | None = None 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/test_attention_backend_config.py b/tests/unit_tests/models/test_attention_backend_config.py index e8dd1551e1..903324aadb 100644 --- a/tests/unit_tests/models/test_attention_backend_config.py +++ b/tests/unit_tests/models/test_attention_backend_config.py @@ -72,9 +72,9 @@ def configure_mcore_attention(config): not hasattr(attention_module.Attention, "_resolve_flash_version"), reason="The frozen MCore dev pin predates FlashAttention version controls", ) -def test_default_flash_version_disables_fa4_without_changing_backend(config_type, backend, monkeypatch): +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) + config = make_config(config_type, backend, flash_attention_version=2) config.finalize() configure_mcore_attention(config) @@ -89,10 +89,10 @@ def test_default_flash_version_disables_fa4_without_changing_backend(config_type not hasattr(attention_module.Attention, "_resolve_flash_version"), reason="The frozen MCore dev pin predates FlashAttention version controls", ) -def test_direct_attention_defaults_to_fa2_even_when_fa4_is_installed(config_type, monkeypatch): +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) + config = make_config(config_type, AttnBackend.auto, flash_attention_version=2) config.finalize() attention = SimpleNamespace(flash_attention_version=config.flash_attention_version) @@ -101,6 +101,10 @@ def test_direct_attention_defaults_to_fa2_even_when_fa4_is_installed(config_type @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) @@ -109,6 +113,16 @@ def test_explicit_flash_version_is_preserved(config_type, version): 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) From 36d03cedc42437d2990716ebcca5677f89610a4d Mon Sep 17 00:00:00 2001 From: Ajay Balasa Date: Fri, 9 Oct 2026 21:17:03 +0000 Subject: [PATCH 13/14] build(mcore): use upstream revision for FA2 Signed-off-by: Ajay Balasa --- .main.commit | 2 +- 3rdparty/Megatron-LM | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.main.commit b/.main.commit index 7b1b7092cd..2708f79afc 100644 --- a/.main.commit +++ b/.main.commit @@ -1 +1 @@ -38a2399e6d3811b4455e50ae199ffbf6ce60e8df +441a987409b8fdd2b90b463f7999bac970109a43 diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index 38a2399e6d..441a987409 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit 38a2399e6d3811b4455e50ae199ffbf6ce60e8df +Subproject commit 441a987409b8fdd2b90b463f7999bac970109a43 From d28698ba7ae16880e318b49a696d21c3558623ab Mon Sep 17 00:00:00 2001 From: Ajay Balasa Date: Fri, 9 Oct 2026 22:04:34 +0000 Subject: [PATCH 14/14] revert: remove FA4-only vision metadata changes Signed-off-by: Ajay Balasa --- .../modeling_ernie45_vl/vision_model.py | 2 +- .../modelling_qwen3_vl/vision_model.py | 2 +- .../bridge/training/nemotron_omni_step.py | 2 +- .../models/ernie_vl/test_ernie45_vl_bridge.py | 24 ---------- .../modelling_qwen3_vl/test_vision_model.py | 30 ------------- .../training/test_nemotron_omni_step.py | 44 ------------------- 6 files changed, 3 insertions(+), 101 deletions(-) diff --git a/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_model.py b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_model.py index 7850593584..dce0f58bb6 100644 --- a/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_model.py +++ b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_model.py @@ -267,7 +267,7 @@ def build_packed_seq_params( cu_seqlens = seqlens.cumsum(dim=0) cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0).int() - max_seqlen_q = int(seqlens.max().item()) + max_seqlen_q = seqlens.max() return PackedSeqParams( cu_seqlens_q=cu_seqlens, cu_seqlens_kv=cu_seqlens, diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/vision_model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/vision_model.py index 6afa4144e2..ce9c027261 100644 --- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/vision_model.py +++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/vision_model.py @@ -383,7 +383,7 @@ def build_packed_seq_params( cu_seqlens = seqlens.cumsum(dim=0) cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0).int() - max_seqlen_q = int(seqlens.max().item()) + max_seqlen_q = seqlens.max() return PackedSeqParams( cu_seqlens_q=cu_seqlens, cu_seqlens_kv=cu_seqlens, diff --git a/src/megatron/bridge/training/nemotron_omni_step.py b/src/megatron/bridge/training/nemotron_omni_step.py index e0732fb680..2e0dfd4b1b 100644 --- a/src/megatron/bridge/training/nemotron_omni_step.py +++ b/src/megatron/bridge/training/nemotron_omni_step.py @@ -158,7 +158,7 @@ def _build_vision_packed_seq_params( cu.append(cu[-1] + sl) device = imgs_sizes.device if torch.is_tensor(imgs_sizes) else torch.device("cpu") cu_tensor = torch.tensor(cu, dtype=torch.int32, device=device) - max_len = max(seq_lens, default=0) + max_len = torch.tensor(max(seq_lens) if seq_lens else 0, dtype=torch.int32, device=device) return PackedSeqParams( qkv_format="thd", cu_seqlens_q=cu_tensor, diff --git a/tests/unit_tests/models/ernie_vl/test_ernie45_vl_bridge.py b/tests/unit_tests/models/ernie_vl/test_ernie45_vl_bridge.py index 8d1762fe54..984d7afeeb 100644 --- a/tests/unit_tests/models/ernie_vl/test_ernie45_vl_bridge.py +++ b/tests/unit_tests/models/ernie_vl/test_ernie45_vl_bridge.py @@ -34,30 +34,6 @@ MultiTypeMoeSubmodules, ) from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.model import ErnieMultimodalRotaryEmbedding -from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.vision_model import ErnieVLVisionModel - - -@pytest.mark.unit -@pytest.mark.parametrize( - "grid,expected_boundaries,expected_max", - [ - ([[1, 2, 3]], [0, 6], 6), - ([[2, 2, 3], [1, 2, 4]], [0, 6, 12, 20], 8), - ], -) -def test_vision_packed_metadata_uses_integer_maxima(grid, expected_boundaries, expected_max): - grid = torch.tensor(grid, dtype=torch.int64) - - packed = ErnieVLVisionModel.build_packed_seq_params(None, grid) - - assert packed.qkv_format == "thd" - assert type(packed.max_seqlen_q) is int - assert type(packed.max_seqlen_kv) is int - assert packed.max_seqlen_q == packed.max_seqlen_kv == expected_max - for boundaries in (packed.cu_seqlens_q, packed.cu_seqlens_kv): - assert boundaries.tolist() == expected_boundaries - assert boundaries.dtype == torch.int32 - assert boundaries.device == grid.device def _make_vision_config(): diff --git a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_vision_model.py b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_vision_model.py index 2265958860..82043ac117 100644 --- a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_vision_model.py +++ b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_vision_model.py @@ -95,36 +95,6 @@ def test_raises_when_seq_exceeds_max(self): class TestVisionForwardPackedAttentionSetup: """Tests for ``_vision_forward_packed_attention_setup``.""" - @pytest.mark.unit - @pytest.mark.parametrize( - "grid,expected_boundaries,expected_max", - [ - ([[1, 2, 3]], [0, 6], 6), - ([[2, 2, 3], [1, 2, 4]], [0, 6, 12, 20], 8), - ], - ) - def test_non_cuda_graph_path_uses_integer_maxima(self, grid, expected_boundaries, expected_max): - grid = torch.tensor(grid, dtype=torch.int64) - seq_len = expected_boundaries[-1] - packed, mask = _vision_forward_packed_attention_setup( - use_cuda_graph_padding=False, - hidden_states=torch.zeros(seq_len, 1, 8), - original_seq_len=seq_len, - seq_len=seq_len, - grid_thw=grid, - build_packed_seq_params=lambda value: Qwen3VLVisionModel.build_packed_seq_params(None, value), - ) - - assert mask is None - assert packed.qkv_format == "thd" - assert type(packed.max_seqlen_q) is int - assert type(packed.max_seqlen_kv) is int - assert packed.max_seqlen_q == packed.max_seqlen_kv == expected_max - for boundaries in (packed.cu_seqlens_q, packed.cu_seqlens_kv): - assert boundaries.tolist() == expected_boundaries - assert boundaries.dtype == torch.int32 - assert boundaries.device == grid.device - def test_non_cuda_graph_path_calls_builder_and_no_mask(self): grid = torch.tensor([[1, 4, 4]], dtype=torch.long) hidden = torch.zeros(5, 1, 8, dtype=torch.float32) diff --git a/tests/unit_tests/training/test_nemotron_omni_step.py b/tests/unit_tests/training/test_nemotron_omni_step.py index ce32480228..c6911dbe55 100644 --- a/tests/unit_tests/training/test_nemotron_omni_step.py +++ b/tests/unit_tests/training/test_nemotron_omni_step.py @@ -22,29 +22,6 @@ from megatron.bridge.training.utils.packed_seq_utils import get_packed_seq_params -@pytest.mark.unit -@pytest.mark.parametrize( - "image_sizes,expected_boundaries,expected_max", - [ - ([[512, 512]], [0, 1024], 1024), - ([[32, 48], [64, 32]], [0, 6, 14], 8), - ], -) -def test_vision_packed_metadata_uses_integer_maxima(image_sizes, expected_boundaries, expected_max): - image_sizes = torch.tensor(image_sizes, dtype=torch.int64) - - metadata = nemotron_omni_step._build_vision_packed_seq_params(image_sizes) - - assert metadata.qkv_format == "thd" - assert type(metadata.max_seqlen_q) is int - assert type(metadata.max_seqlen_kv) is int - assert metadata.max_seqlen_q == metadata.max_seqlen_kv == expected_max - for boundaries in (metadata.cu_seqlens_q, metadata.cu_seqlens_kv): - assert boundaries.tolist() == expected_boundaries - assert boundaries.dtype == torch.int32 - assert boundaries.device == image_sizes.device - - def test_batch_moves_only_one_compatible_token_alias_to_cuda(monkeypatch): cuda_inputs = [] @@ -145,27 +122,6 @@ def _pipeline_cfg(*, packed=True, defer_packing=False, temporal_patch_dim=1): ) -@pytest.mark.unit -@pytest.mark.parametrize("packed", [False, True]) -def test_first_pipeline_batch_keeps_vision_maxima_as_integers(monkeypatch, packed): - monkeypatch.setattr(torch.Tensor, "cuda", lambda self, **kwargs: self) - monkeypatch.setattr(nemotron_omni_step, "is_pp_first_stage", lambda group: True) - monkeypatch.setattr(nemotron_omni_step, "is_pp_last_stage", lambda group: False) - batch = _packed_pipeline_batch() - if not packed: - for key in (*nemotron_omni_step._PACKED_SEQ_PARAM_KEYS, "padding_mask"): - batch.pop(key, None) - - result = get_batch(iter([batch]), _pipeline_cfg(packed=packed), pg_collection=SimpleNamespace(pp=object())) - - vision_metadata = result[12] - assert type(vision_metadata.max_seqlen_q) is int - assert type(vision_metadata.max_seqlen_kv) is int - assert vision_metadata.max_seqlen_q == vision_metadata.max_seqlen_kv == 4 - assert vision_metadata.cu_seqlens_q.tolist() == [0, 4, 8] - assert vision_metadata.cu_seqlens_kv.tolist() == [0, 4, 8] - - def test_forward_rejects_deferred_multimodal_packing(): state = SimpleNamespace( timers=None,