From ea1af52f39ecbd6056461ae81ae390679886f529 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Thu, 9 Apr 2026 15:52:49 +0800
Subject: [PATCH 01/51] fix: align THD MoE aux-loss token accounting with valid
 tokens

Pass a THD padding_mask derived from lm_input_ids into GPTModel so MoE router
aux/global-aux metrics exclude tail padding and match BSHD token-count semantics.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../models/qwen_vl/modelling_qwen3_vl/model.py  | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
index 702723fc06..5b0e8424e7 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
@@ -622,6 +622,22 @@ def forward(
         torch.cuda.nvtx.range_pop()
         torch.cuda.nvtx.range_push("Qwen3VLModel.forward.language_model")
 
+        # For THD packed path we intentionally keep attention_mask=None for model forward.
+        # MoE aux/global-aux accounting still needs to ignore tail padding tokens, so pass
+        # a dedicated padding_mask (True=padding) to GPTModel/Router.
+        padding_mask_for_moe = None
+        if packed_seq_params is not None and lm_input_ids is not None:
+            padding_mask_for_moe = lm_input_ids.eq(0)
+            if _thd_diag_enabled() and _rank0():
+                pad_cnt = int(padding_mask_for_moe.sum().item())
+                tok_cnt = int(padding_mask_for_moe.numel())
+                logger.info(
+                    "[THD_DIAG][model] padding_mask_for_moe: total_tokens=%d padding_tokens=%d valid_tokens=%d",
+                    tok_cnt,
+                    pad_cnt,
+                    tok_cnt - pad_cnt,
+                )
+
         output = self.language_model(
             input_ids=lm_input_ids,
             position_ids=position_ids,  # None in encoder
@@ -629,6 +645,7 @@ def forward(
             decoder_input=combined_embeddings,  # only not None in the first decoder PP stage
             labels=labels,  # only not None in the last decoder PP stage
             loss_mask=loss_mask,  # Added for THD training compatibility
+            padding_mask=padding_mask_for_moe,  # for MoE routing/aux-loss token accounting
             inference_params=inference_params,  # currently always None
             packed_seq_params=packed_seq_params,  # currently always None
             visual_pos_masks=visual_pos_masks,

From da9d92634b59b27ac1348ba365ceba51a63569e6 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Thu, 9 Apr 2026 16:20:35 +0800
Subject: [PATCH 02/51] fix: wire padding_mask through qwen3-vl gpt forward
 path

Allow Qwen3VLGPTModel and Qwen3VLTransformerBlock to accept and propagate padding_mask so THD MoE routing can consume valid-token masking without crashing.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../models/qwen_vl/modelling_qwen3_vl/text_model.py      | 6 +++++-
 .../qwen_vl/modelling_qwen3_vl/transformer_block.py      | 9 +++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py
index d067655a1b..5e05ca9ec7 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py
@@ -120,6 +120,7 @@ def forward(
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         loss_mask: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
         # args for deepstack
         visual_pos_masks: Optional[torch.Tensor] = None,
         deepstack_visual_embeds: Optional[list[torch.Tensor]] = None,
@@ -148,6 +149,7 @@ def forward(
             decoder_input=decoder_input,
             inference_context=inference_context,
             packed_seq_params=packed_seq_params,
+            padding_mask=padding_mask,
         )
 
         (
@@ -156,7 +158,8 @@ def forward(
             rotary_pos_cos,
             rotary_pos_sin,
             sequence_len_offset,
-        ) = preproc_output[:5]
+            padding_mask,
+        ) = preproc_output[:6]
 
         # Run decoder.
         hidden_states = self.decoder(
@@ -170,6 +173,7 @@ def forward(
             # the standard components only.
             packed_seq_params=packed_seq_params,
             sequence_len_offset=sequence_len_offset,
+            padding_mask=padding_mask,
             visual_pos_masks=visual_pos_masks,
             deepstack_visual_embeds=deepstack_visual_embeds,
             **(extra_block_kwargs or {}),
diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_block.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_block.py
index 6284617bed..440e45c873 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_block.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_block.py
@@ -503,6 +503,7 @@ def _checkpointed_forward(
         rotary_pos_emb: Tensor,
         attention_bias: Tensor,
         packed_seq_params: PackedSeqParams,
+        padding_mask: Optional[Tensor],
         use_inner_fp8_context: bool,
         # args for deepstack
         visual_pos_masks: Optional[torch.Tensor] = None,
@@ -517,6 +518,7 @@ def custom_forward(
                 context,
                 context_mask,
                 rotary_pos_emb,
+                padding_mask,
                 visual_pos_masks,
                 *deepstack_visual_embeds_args,
             ):
@@ -538,6 +540,7 @@ def custom_forward(
                             attention_bias=attention_bias,
                             inference_context=None,
                             packed_seq_params=packed_seq_params,
+                            padding_mask=padding_mask,
                         )
 
                         if self.pre_process and deepstack_visual_embeds is not None:
@@ -567,6 +570,7 @@ def checkpoint_handler(forward_func):
                     context,
                     context_mask,
                     rotary_pos_emb,
+                    padding_mask,
                     visual_pos_masks,
                     *deepstack_visual_embeds_tuple,
                 )
@@ -579,6 +583,7 @@ def checkpoint_handler(forward_func):
                     context,
                     context_mask,
                     rotary_pos_emb,
+                    padding_mask,
                     visual_pos_masks,
                     *deepstack_visual_embeds_tuple,
                 )
@@ -618,6 +623,7 @@ def checkpoint_handler(forward_func):
                         context,
                         context_mask,
                         rotary_pos_emb,
+                        padding_mask,
                         visual_pos_masks,
                         *deepstack_visual_embeds_tuple,
                     )
@@ -639,6 +645,7 @@ def forward(
         inference_context: Optional[BaseInferenceContext] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
         sequence_len_offset: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
         *,
         inference_params: Optional[BaseInferenceContext] = None,
         # args for deepstack
@@ -730,6 +737,7 @@ def forward(
                     rotary_pos_emb=rotary_pos_emb,
                     attention_bias=attention_bias,
                     packed_seq_params=packed_seq_params,
+                    padding_mask=padding_mask,
                     use_inner_fp8_context=use_inner_fp8_context,
                     visual_pos_masks=visual_pos_masks,
                     deepstack_visual_embeds=deepstack_visual_embeds,
@@ -754,6 +762,7 @@ def forward(
                             inference_context=inference_context,
                             packed_seq_params=packed_seq_params,
                             sequence_len_offset=sequence_len_offset,
+                            padding_mask=padding_mask,
                         )
 
                         if self.pre_process and deepstack_visual_embeds is not None:

From 653dd62fea5349a4a48eeaf9b9a4cf4bf62dad6a Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Thu, 9 Apr 2026 17:08:33 +0800
Subject: [PATCH 03/51] fix: pass padding_mask into qwen gpt postprocess mtp
 path

Forward padding_mask from Qwen3VLGPTModel into _postprocess so MTP MoE global aux-loss uses masked token accounting consistently.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py       | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py
index 5e05ca9ec7..d6781ced3a 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py
@@ -208,6 +208,7 @@ def _sp_scatter_embedding(input_ids, position_ids):
             loss_mask=loss_mask,
             decoder_input=decoder_input,
             attention_mask=attention_mask,
+            padding_mask=padding_mask,
             inference_params=inference_params,
             packed_seq_params=packed_seq_params,
             sequence_len_offset=sequence_len_offset,

From 88a8ac209c91e5c6a80fb0ca4e34ea574eeea61b Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Thu, 9 Apr 2026 19:30:19 +0800
Subject: [PATCH 04/51] fix: update Megatron-LM pointer for main-based THD
 support

Point submodule to the main-based Megatron-LM branch commit that propagates padding_mask through MTP MoE routing.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 3rdparty/Megatron-LM | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM
index 23265d2db3..ec7ee7928f 160000
--- a/3rdparty/Megatron-LM
+++ b/3rdparty/Megatron-LM
@@ -1 +1 @@
-Subproject commit 23265d2db37efa883b3983bce321c631898822f3
+Subproject commit ec7ee7928f78db084d55c0fa1bf8b6ac7c2a19a5

From 924497c7743ce73acc82942ae523780464abd44f Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Fri, 10 Apr 2026 10:43:27 +0800
Subject: [PATCH 05/51] fix: make diffusion recipe imports optional

Guard diffusion recipe imports in recipes.__init__ so missing optional dependencies like diffusers do not block LLM/VLM training startup.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 src/megatron/bridge/recipes/__init__.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/megatron/bridge/recipes/__init__.py b/src/megatron/bridge/recipes/__init__.py
index 46d51fdb64..4428a47191 100644
--- a/src/megatron/bridge/recipes/__init__.py
+++ b/src/megatron/bridge/recipes/__init__.py
@@ -18,8 +18,14 @@
 This module exposes all recipe configurations from all model families.
 """
 
-from megatron.bridge.diffusion.recipes.flux.flux import *
-from megatron.bridge.diffusion.recipes.wan.wan import *
+# Diffusion recipes require optional dependencies (e.g. diffusers). Keep
+# non-diffusion training paths (LLM/VLM) usable when those deps are absent.
+try:
+    from megatron.bridge.diffusion.recipes.flux.flux import *
+    from megatron.bridge.diffusion.recipes.wan.wan import *
+except ModuleNotFoundError:
+    pass
+
 from megatron.bridge.recipes.deepseek import *
 from megatron.bridge.recipes.gemma import *
 from megatron.bridge.recipes.gemma3_vl import *

From 552282d7b123a0e7e0a1aa7e3975cf7debb51d09 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Fri, 10 Apr 2026 11:38:28 +0800
Subject: [PATCH 06/51] fix: guard optional deps for Qwen-VL recipe startup

Avoid import-time crashes in training environments missing optional diffusion/modelopt/transformers pieces by adding graceful fallbacks for quantization and processor loading.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../bridge/models/conversion/auto_bridge.py    |  9 ++++++++-
 .../bridge/recipes/qwen_vl/qwen3_vl.py         | 18 ++++++++++++++----
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/src/megatron/bridge/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py
index 0c61b75d47..8dbcb921d1 100644
--- a/src/megatron/bridge/models/conversion/auto_bridge.py
+++ b/src/megatron/bridge/models/conversion/auto_bridge.py
@@ -30,7 +30,6 @@
 
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.transformer_config import MLATransformerConfig, TransformerConfig
-from modelopt.torch.quantization.utils import is_quantized
 from safetensors.torch import save_file
 from transformers.configuration_utils import PretrainedConfig
 from typing_extensions import Unpack
@@ -51,6 +50,14 @@
 
 logger = logging.getLogger(__name__)
 
+try:
+    from modelopt.torch.quantization.utils import is_quantized
+except Exception:
+    # modelopt/scipy is optional for recipe import and training startup.
+    # Keep quantization-only export path available when dependency exists.
+    def is_quantized(_model: object) -> bool:
+        return False
+
 MegatronModelT = TypeVar("MegatronModelT", bound=MegatronModule)
 DataclassT = TypeVar("DataclassT")
 
diff --git a/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py b/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py
index 5f84009620..5c445784aa 100644
--- a/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py
+++ b/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py
@@ -23,9 +23,14 @@
 from typing import Optional, Union
 
 import torch
-from transformers import AutoTokenizer, Qwen3VLProcessor
+from transformers import AutoTokenizer
 from typing_extensions import TypedDict, Unpack
 
+try:
+    from transformers import Qwen3VLProcessor
+except Exception:
+    Qwen3VLProcessor = None
+
 from megatron.bridge import AutoBridge
 from megatron.bridge.data.energon.energon_provider import EnergonProvider
 from megatron.bridge.data.vlm_datasets import MockVLMConversationProvider
@@ -270,9 +275,14 @@ def _make_energon_dataset(
 ) -> EnergonProvider:
     """Create an EnergonProvider dataset config for Qwen3-VL recipes."""
     tokenizer = AutoTokenizer.from_pretrained(hf_path)
-    # Use Qwen3VLProcessor to match the HF flow (which uses AutoProcessor).
-    # This processor accepts both images and videos kwargs.
-    image_processor = Qwen3VLProcessor.from_pretrained(hf_path)
+    # Prefer Qwen3VLProcessor on newer transformers; fall back to AutoProcessor
+    # for older container images that do not yet expose that symbol.
+    if Qwen3VLProcessor is not None:
+        image_processor = Qwen3VLProcessor.from_pretrained(hf_path)
+    else:
+        from transformers import AutoProcessor
+
+        image_processor = AutoProcessor.from_pretrained(hf_path)
     task_encoder = QwenVLTaskEncoder(
         tokenizer=tokenizer,
         image_processor=image_processor,

From 24647716e1a176858a0cdbbdcd4a2cdaa2eb57da Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Fri, 10 Apr 2026 18:04:36 +0800
Subject: [PATCH 07/51] feat: add qwen35_vl_35b_a3b_sft_energon_config recipe

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../bridge/recipes/qwen_vl/__init__.py        |  2 ++
 .../bridge/recipes/qwen_vl/qwen35_vl.py       | 19 ++++++++++++++++++-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/src/megatron/bridge/recipes/qwen_vl/__init__.py b/src/megatron/bridge/recipes/qwen_vl/__init__.py
index 9bc536c45d..0c45ee8dc2 100644
--- a/src/megatron/bridge/recipes/qwen_vl/__init__.py
+++ b/src/megatron/bridge/recipes/qwen_vl/__init__.py
@@ -52,6 +52,7 @@
     qwen35_vl_35b_a3b_peft_config,
     qwen35_vl_35b_a3b_pretrain_mock_config,
     qwen35_vl_35b_a3b_sft_config,
+    qwen35_vl_35b_a3b_sft_energon_config,
     qwen35_vl_122b_a10b_peft_config,
     qwen35_vl_122b_a10b_pretrain_mock_config,
     qwen35_vl_122b_a10b_sft_config,
@@ -77,6 +78,7 @@
     "qwen35_vl_27b_sft_config",
     # Qwen3.5-VL SFT configs — MoE
     "qwen35_vl_35b_a3b_sft_config",
+    "qwen35_vl_35b_a3b_sft_energon_config",
     "qwen35_vl_35b_a3b_fsdp_sft_config",
     "qwen35_vl_122b_a10b_sft_config",
     "qwen35_vl_397b_a17b_sft_config",
diff --git a/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py b/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
index af089cc28b..c1bb8fd4ca 100644
--- a/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
+++ b/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
@@ -28,7 +28,11 @@
 from megatron.bridge import AutoBridge
 from megatron.bridge.peft.base import PEFT
 from megatron.bridge.recipes.common import _peft_common_vlm, _sft_common_vlm
-from megatron.bridge.recipes.qwen_vl.qwen3_vl import Qwen3VLCommonKwargs, _qwen3_vl_common
+from megatron.bridge.recipes.qwen_vl.qwen3_vl import (
+    Qwen3VLCommonKwargs,
+    _make_energon_dataset,
+    _qwen3_vl_common,
+)
 from megatron.bridge.recipes.utils.finetune_utils import default_peft_config
 from megatron.bridge.recipes.utils.optimizer_utils import distributed_fused_adam_with_cosine_annealing
 from megatron.bridge.training.config import ConfigContainer
@@ -385,6 +389,19 @@ def qwen35_vl_35b_a3b_sft_config(hf_path: str = "Qwen/Qwen3.5-35B-A3B") -> Confi
     return cfg
 
 
+def qwen35_vl_35b_a3b_sft_energon_config(hf_path: str = "Qwen/Qwen3.5-35B-A3B") -> ConfigContainer:
+    """Return a full SFT config for Qwen3.5-VL 35B-A3B (MoE) with Energon dataset.
+
+    Same as qwen35_vl_35b_a3b_sft_config but uses EnergonProvider instead of HF dataset.
+    Set the dataset path via CLI override: dataset.path=/path/to/energon/dataset
+    """
+    cfg = qwen35_vl_35b_a3b_sft_config(hf_path=hf_path)
+    cfg.dataset = _make_energon_dataset(
+        hf_path, cfg.dataset.seq_length, cfg.train.micro_batch_size, cfg.train.global_batch_size
+    )
+    return cfg
+
+
 def qwen35_vl_35b_a3b_fsdp_sft_config(hf_path: str = "Qwen/Qwen3.5-35B-A3B") -> ConfigContainer:
     """Return a full SFT config for Qwen3.5-VL 35B-A3B (MoE) with Megatron FSDP.
 

From 7d4502dcbe6202855a26c5122254be41a69c628b Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Mon, 13 Apr 2026 08:44:19 +0800
Subject: [PATCH 08/51] fix: add packing_buffer_size to EnergonProvider and
 harden Qwen3VLProcessor fallback

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 src/megatron/bridge/data/energon/energon_provider.py |  2 ++
 src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py      | 11 ++++++++---
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/megatron/bridge/data/energon/energon_provider.py b/src/megatron/bridge/data/energon/energon_provider.py
index f33ea48dc1..2afcd0d1b9 100644
--- a/src/megatron/bridge/data/energon/energon_provider.py
+++ b/src/megatron/bridge/data/energon/energon_provider.py
@@ -35,6 +35,7 @@ class EnergonProvider(DatasetProvider):
     task_encoder: Optional[Any] = None
     # Enable batch-level online sequence packing
     pack_sequences_in_batch: bool = False
+    packing_buffer_size: Optional[int] = None
 
     def build_datasets(self, context: DatasetBuildContext):
         assert self.path, "EnergonProvider.path must be set. Use CLI override: dataset.path=<path>"
@@ -47,6 +48,7 @@ def build_datasets(self, context: DatasetBuildContext):
             micro_batch_size=self.micro_batch_size,
             global_batch_size=self.global_batch_size,
             num_workers=self.num_workers,
+            packing_buffer_size=self.packing_buffer_size,
             pg_collection=context.pg_collection,
         )
         return (
diff --git a/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py b/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py
index 5c445784aa..64b126974a 100644
--- a/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py
+++ b/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py
@@ -276,10 +276,15 @@ def _make_energon_dataset(
     """Create an EnergonProvider dataset config for Qwen3-VL recipes."""
     tokenizer = AutoTokenizer.from_pretrained(hf_path)
     # Prefer Qwen3VLProcessor on newer transformers; fall back to AutoProcessor
-    # for older container images that do not yet expose that symbol.
+    # for older container images that do not yet expose that symbol or lack
+    # the video_processing sub-module required by from_pretrained.
+    image_processor = None
     if Qwen3VLProcessor is not None:
-        image_processor = Qwen3VLProcessor.from_pretrained(hf_path)
-    else:
+        try:
+            image_processor = Qwen3VLProcessor.from_pretrained(hf_path)
+        except Exception:
+            pass
+    if image_processor is None:
         from transformers import AutoProcessor
 
         image_processor = AutoProcessor.from_pretrained(hf_path)

From 942240792bb1ff0a02a17243179c5fb99b55397e Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 31 Mar 2026 10:51:01 +0800
Subject: [PATCH 09/51] feat: Add Energon THD packing support for QwenVL
 TaskEncoder and vlm_step

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../qwen_vl/data/energon/task_encoder.py      | 360 ++++++++++++++++--
 src/megatron/bridge/training/vlm_step.py      |  62 ++-
 2 files changed, 381 insertions(+), 41 deletions(-)

diff --git a/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py b/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
index 50178a0749..13e90877f0 100644
--- a/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
+++ b/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
@@ -12,16 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import bisect
 import dataclasses
 import logging
 import re
 from collections import defaultdict
-from dataclasses import dataclass
-from typing import Dict, List
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Union
 
 import numpy as np
 import torch
-from megatron.energon import Batch, DefaultTaskEncoder
+from megatron.energon import Batch, DefaultTaskEncoder, stateless
 from transformers import BatchEncoding
 
 from megatron.bridge.data.energon.task_encoder_utils import (
@@ -39,6 +40,46 @@
 from megatron.bridge.training.utils.visual_inputs import Qwen2_5_VLVisualInputs
 
 
+logger = logging.getLogger(__name__)
+
+
+def _search_for_fit(numbers: List[int], capacity: int) -> int:
+    """Binary search for the largest number that fits within capacity."""
+    index = bisect.bisect(numbers, capacity)
+    return -1 if index == 0 else (index - 1)
+
+
+def greedy_knapsack(item_sizes: List[int], samples: List, max_capacity: int) -> List:
+    """Greedy bin-packing with binary search.
+
+    Sorts samples by length ascending, then greedily fills each bin by picking
+    the largest item that still fits (via binary search). Returns a list of bins,
+    each bin being a list of samples.
+    """
+    assert len(item_sizes) == len(samples)
+    if not item_sizes:
+        return []
+
+    sorted_sizes, sorted_samples = zip(*sorted(zip(item_sizes, samples), key=lambda x: x[0]))
+    sorted_sizes = list(sorted_sizes)
+    sorted_samples = list(sorted_samples)
+
+    if sorted_sizes[-1] > max_capacity:
+        raise ValueError(f"Sample size {sorted_sizes[-1]} exceeds max_capacity {max_capacity}")
+
+    knapsacks = []
+    while sorted_sizes:
+        current_knapsack = []
+        remaining = max_capacity
+        while True:
+            idx = _search_for_fit(sorted_sizes, remaining)
+            if idx == -1:
+                break
+            remaining -= sorted_sizes[idx]
+            sorted_sizes.pop(idx)
+            current_knapsack.append(sorted_samples.pop(idx))
+        knapsacks.append(current_knapsack)
+    return knapsacks
 def process_vision(
     processor, images, videos, fps=None, model_version: str = "qwen-vl", min_pixels=None, max_pixels=None
 ):
@@ -103,6 +144,31 @@ class QwenVLTaskSample:
     video_input_mask: torch.Tensor
     text: torch.Tensor
     target: torch.Tensor
+    total_len: int = 0
+
+
+@dataclass
+class QwenVLTaskSamplePacked:
+    """Packed sample: multiple QwenVLTaskSample concatenated into one sequence."""
+
+    __key__: str
+    __subflavors__: Dict
+    tokens: torch.Tensor
+    target: torch.Tensor
+    imgs: list
+    videos: list
+    image_thw_grids: list
+    video_thw_grids: list
+    image_input_mask: torch.Tensor
+    video_input_mask: torch.Tensor
+    cu_lengths: torch.Tensor
+    max_length: int
+    num_sub_samples: int
+    sub_sample_lengths: list = field(default_factory=list)
+    num_image_tokens: int = 0
+    num_video_tokens: int = 0
+    num_text_tokens: int = 0
+    num_vit_patches: int = 0
 
 
 @dataclass
@@ -125,6 +191,10 @@ class QwenVLTaskBatch(Batch):
     position_ids: torch.Tensor
     labels: torch.Tensor
     loss_mask: torch.Tensor
+    # Packing metadata (set when using Energon packing)
+    cu_lengths: Optional[torch.Tensor] = None
+    max_lengths: Optional[torch.Tensor] = None
+    cu_seqlens_argmin: Optional[torch.Tensor] = None
 
 
 def convert_to_qwenvl_content(user_input: str, image_pattern: str = "<image>", video_pattern: str = "<video>"):
@@ -339,55 +409,276 @@ def encode_sample(self, sample: ChatMLSample):
             video_input_mask=video_input_mask,
             text=torch.from_numpy(final_input_ids),
             target=torch.from_numpy(target),
+            total_len=int(target_length),
         )
 
-    def batch(self, samples: List[QwenVLTaskSample]) -> QwenVLTaskBatch:
+    def select_samples_to_pack(self, samples: List[QwenVLTaskSample]) -> List[List[QwenVLTaskSample]]:
+        """Select groups of samples to pack together using greedy knapsack.
+
+        Called by Energon's PackingDataset. Returns list of bins, each bin
+        is a list of samples whose total length <= self.seq_len.
         """
-        Put encoded sample into Batch, do padding, add labels and visual input masks
+        lengths = [s.total_len for s in samples]
+        return greedy_knapsack(lengths, samples, self.seq_len)
 
-        Args:
-            samples: List of encoded samples
+    @stateless
+    def pack_selected_samples(self, samples: List[QwenVLTaskSample]) -> QwenVLTaskSamplePacked:
+        """Concatenate multiple QwenVLTaskSample into one packed sample.
 
-        Returns:
-            Batch with necessary fields
+        Called by Energon's PackingDataset after select_samples_to_pack.
+        Produces a QwenVLTaskSamplePacked with cu_lengths for THD attention.
         """
-        imgs, image_thw_grids = [], []
+        all_tokens, all_targets = [], []
+        all_imgs, all_videos = [], []
+        all_image_thw_grids, all_video_thw_grids = [], []
+        all_image_masks, all_video_masks = [], []
+        sub_lengths = []
+        total_image_tokens = 0
+        total_video_tokens = 0
+        total_text_tokens = 0
+        total_vit_patches = 0
+        merge_length = self.merge_size ** 2
+
         for s in samples:
+            seq_len = len(s.text)
+            sub_lengths.append(seq_len)
+            all_tokens.append(s.text)
+            all_targets.append(s.target)
+            all_image_masks.append(s.image_input_mask)
+            all_video_masks.append(s.video_input_mask)
+
+            n_img = int(s.image_input_mask.sum().item())
+            n_vid = int(s.video_input_mask.sum().item())
+            n_txt = seq_len - n_img - n_vid
+            total_image_tokens += n_img
+            total_video_tokens += n_vid
+            total_text_tokens += n_txt
+
             if len(s.imgs) > 0:
-                s_imgs = [img for img in s.imgs.unsqueeze(0)]
-                cat_imgs = torch.cat([img for img in s_imgs])
-                imgs.append(cat_imgs)
+                if isinstance(s.imgs, torch.Tensor):
+                    all_imgs.append(s.imgs.unsqueeze(0) if s.imgs.dim() == 3 else s.imgs)
+                else:
+                    all_imgs.append(s.imgs)
             if len(s.image_thw_grids) > 0:
-                s_image_thw_grids = [thw_grids for thw_grids in s.image_thw_grids]
-                image_thw_grids.extend(s_image_thw_grids)
-        videos, video_thw_grids = [], []
-        for s in samples:
+                thw = s.image_thw_grids
+                if isinstance(thw, np.ndarray):
+                    for row in thw:
+                        all_image_thw_grids.append(row)
+                        total_vit_patches += int(np.prod(row))
+                elif isinstance(thw, torch.Tensor):
+                    for row in thw:
+                        all_image_thw_grids.append(row.numpy())
+                        total_vit_patches += int(row.prod().item())
+                else:
+                    all_image_thw_grids.extend(thw)
+
             if len(s.videos) > 0:
-                s_videos = [video for video in s.videos.unsqueeze(0)]
-                cat_videos = torch.cat([video for video in s_videos])
-                videos.append(cat_videos)
+                if isinstance(s.videos, torch.Tensor):
+                    all_videos.append(s.videos.unsqueeze(0) if s.videos.dim() == 3 else s.videos)
+                else:
+                    all_videos.append(s.videos)
             if len(s.video_thw_grids) > 0:
-                s_video_thw_grids = [thw_grids for thw_grids in s.video_thw_grids]
-                video_thw_grids.extend(s_video_thw_grids)
-                # assert s_video_thw_grids.prod(dim=-1).sum() == s_videos.shape[0]
+                thw = s.video_thw_grids
+                if isinstance(thw, np.ndarray):
+                    for row in thw:
+                        all_video_thw_grids.append(row)
+                        total_vit_patches += int(np.prod(row))
+                elif isinstance(thw, torch.Tensor):
+                    for row in thw:
+                        all_video_thw_grids.append(row.numpy())
+                        total_vit_patches += int(row.prod().item())
+                else:
+                    all_video_thw_grids.extend(thw)
+
+        packed_tokens = torch.cat(all_tokens, dim=0)
+        packed_targets = torch.cat(all_targets, dim=0)
+        packed_image_mask = torch.cat(all_image_masks, dim=0)
+        packed_video_mask = torch.cat(all_video_masks, dim=0)
+
+        cu_lengths_list = [0]
+        for l in sub_lengths:
+            cu_lengths_list.append(cu_lengths_list[-1] + l)
+        cu_lengths = torch.tensor(cu_lengths_list, dtype=torch.int32)
+
+        total_packed = int(cu_lengths[-1].item())
+        max_sub_len = max(sub_lengths)
+
+        logger.info(
+            f"[PackingStats] packed {len(samples)} samples → {total_packed} tokens "
+            f"(seq_len={self.seq_len}, utilization={total_packed / self.seq_len * 100:.1f}%), "
+            f"sub_lengths={sub_lengths}, "
+            f"image_tokens={total_image_tokens}, video_tokens={total_video_tokens}, "
+            f"text_tokens={total_text_tokens}, vit_patches={total_vit_patches}"
+        )
+
+        return QwenVLTaskSamplePacked(
+            __key__="+".join(s.__key__ for s in samples),
+            __subflavors__=samples[0].__subflavors__,
+            tokens=packed_tokens,
+            target=packed_targets,
+            imgs=all_imgs,
+            videos=all_videos,
+            image_thw_grids=all_image_thw_grids,
+            video_thw_grids=all_video_thw_grids,
+            image_input_mask=packed_image_mask,
+            video_input_mask=packed_video_mask,
+            cu_lengths=cu_lengths,
+            max_length=max_sub_len,
+            num_sub_samples=len(samples),
+            sub_sample_lengths=sub_lengths,
+            num_image_tokens=total_image_tokens,
+            num_video_tokens=total_video_tokens,
+            num_text_tokens=total_text_tokens,
+            num_vit_patches=total_vit_patches,
+        )
+
+    def _collect_visual_data(self, samples):
+        """Gather pixel_values and grid_thw from a list of samples (packed or unpacked)."""
+        imgs, image_thw_grids = [], []
+        videos, video_thw_grids = [], []
+
+        for s in samples:
+            if isinstance(s, QwenVLTaskSamplePacked):
+                for img_block in s.imgs:
+                    if isinstance(img_block, torch.Tensor):
+                        imgs.append(img_block if img_block.dim() == 4 else img_block.unsqueeze(0))
+                for thw in s.image_thw_grids:
+                    image_thw_grids.append(np.asarray(thw))
+                for vid_block in s.videos:
+                    if isinstance(vid_block, torch.Tensor):
+                        videos.append(vid_block if vid_block.dim() == 4 else vid_block.unsqueeze(0))
+                for thw in s.video_thw_grids:
+                    video_thw_grids.append(np.asarray(thw))
+            else:
+                if len(s.imgs) > 0:
+                    s_imgs = s.imgs.unsqueeze(0) if isinstance(s.imgs, torch.Tensor) and s.imgs.dim() == 3 else s.imgs
+                    if isinstance(s_imgs, torch.Tensor):
+                        imgs.append(s_imgs)
+                if len(s.image_thw_grids) > 0:
+                    if isinstance(s.image_thw_grids, np.ndarray):
+                        for row in s.image_thw_grids:
+                            image_thw_grids.append(row)
+                    else:
+                        image_thw_grids.extend(s.image_thw_grids)
+                if len(s.videos) > 0:
+                    s_vids = s.videos.unsqueeze(0) if isinstance(s.videos, torch.Tensor) and s.videos.dim() == 3 else s.videos
+                    if isinstance(s_vids, torch.Tensor):
+                        videos.append(s_vids)
+                if len(s.video_thw_grids) > 0:
+                    if isinstance(s.video_thw_grids, np.ndarray):
+                        for row in s.video_thw_grids:
+                            video_thw_grids.append(row)
+                    else:
+                        video_thw_grids.extend(s.video_thw_grids)
+
+        return imgs, image_thw_grids, videos, video_thw_grids
+
+    def batch(self, samples: List[Union[QwenVLTaskSample, QwenVLTaskSamplePacked]]) -> QwenVLTaskBatch:
+        """
+        Put encoded sample into Batch, do padding, add labels and visual input masks.
+
+        Handles both unpacked (QwenVLTaskSample) and packed (QwenVLTaskSamplePacked) inputs.
+        For packed inputs, pads to self.seq_len and passes cu_lengths through.
+        """
+        is_packed = isinstance(samples[0], QwenVLTaskSamplePacked)
+
+        imgs, image_thw_grids, videos, video_thw_grids = self._collect_visual_data(samples)
 
-        # use the max sample lengths in the batch.
+        pad_token_id = self.hf_tokenizer.pad_token_id
+
+        if is_packed:
+            batch_size = len(samples)
+            fixed_len = self.seq_len
+
+            text_mat = np.full((batch_size, fixed_len), pad_token_id, dtype=np.int64)
+            target_mat = np.full((batch_size, fixed_len), pad_token_id, dtype=np.int64)
+            image_input_masks = np.zeros((batch_size, fixed_len), dtype=bool)
+            video_input_masks = np.zeros((batch_size, fixed_len), dtype=bool)
+
+            cu_lengths_list = []
+            max_lengths_list = []
+
+            for i, s in enumerate(samples):
+                content_len = len(s.tokens)
+                actual_len = min(content_len, fixed_len)
+                text_mat[i, :actual_len] = np.array(s.tokens)[:actual_len]
+                target_mat[i, :actual_len] = np.array(s.target)[:actual_len]
+                image_input_masks[i, :actual_len] = np.array(s.image_input_mask)[:actual_len]
+                video_input_masks[i, :actual_len] = np.array(s.video_input_mask)[:actual_len]
+
+                cu = s.cu_lengths.clone()
+                cu[cu > fixed_len] = fixed_len
+                cu_lengths_list.append(cu)
+                max_lengths_list.append(s.max_length)
+
+                pad_len = fixed_len - actual_len
+                logger.info(
+                    f"[BatchStats] sample {i}: {s.num_sub_samples} sub-seqs packed, "
+                    f"content={actual_len}, pad={pad_len}, "
+                    f"image_toks={s.num_image_tokens}, video_toks={s.num_video_tokens}, "
+                    f"text_toks={s.num_text_tokens}, vit_patches={s.num_vit_patches}, "
+                    f"utilization={actual_len / fixed_len * 100:.1f}%"
+                )
+
+            tokens = torch.from_numpy(text_mat)
+            tokens[tokens == pad_token_id] = 0
+            labels = torch.from_numpy(target_mat)
+            labels[labels == pad_token_id] = IGNORE_INDEX
+
+            attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
+                data=tokens,
+                eod_token=self.hf_tokenizer.eos_token_id,
+                eod_mask_loss=False,
+                reset_attention_mask=False,
+                reset_position_ids=False,
+                compute_attention_mask=False,
+            )
+            loss_mask[labels < 0] = 0.0
+
+            max_cu_len = max(len(cu) for cu in cu_lengths_list)
+            padded_cu = torch.zeros((batch_size, max_cu_len), dtype=torch.int32)
+            for i, cu in enumerate(cu_lengths_list):
+                padded_cu[i, : len(cu)] = cu
+
+            cu_seqlens_argmin_list = [len(cu) for cu in cu_lengths_list]
+            cu_seqlens_argmin = torch.tensor(cu_seqlens_argmin_list, dtype=torch.int32)
+            max_lengths = torch.tensor(max_lengths_list, dtype=torch.int32)
+
+            batch_obj = QwenVLTaskBatch(
+                __keys__=[s.__key__ for s in samples],
+                __subflavors__=[s.__subflavors__ for s in samples],
+                pixel_values=torch.vstack(imgs) if len(imgs) > 0 else None,
+                pixel_values_videos=torch.vstack(videos) if len(videos) > 0 else None,
+                image_grid_thw=torch.from_numpy(np.array(image_thw_grids)) if len(image_thw_grids) > 0 else None,
+                video_grid_thw=torch.from_numpy(np.array(video_thw_grids)) if len(video_thw_grids) > 0 else None,
+                image_input_mask=torch.from_numpy(image_input_masks),
+                video_input_mask=torch.from_numpy(video_input_masks),
+                input_ids=tokens,
+                attention_mask=None,
+                position_ids=position_ids,
+                labels=labels,
+                loss_mask=loss_mask,
+                cu_lengths=padded_cu,
+                max_lengths=max_lengths,
+                cu_seqlens_argmin=cu_seqlens_argmin,
+            )
+            return batch_obj
+
+        # --- Original unpacked path ---
         max_seq_len = max(len(s.text) for s in samples)
         if max_seq_len > self.seq_len:
             logging.warning("max sequence length larger than passed parameter")
 
-        text_mat = np.full((len(samples), max_seq_len), self.hf_tokenizer.pad_token_id, dtype=np.int64)
-        target_mat = np.full((len(samples), max_seq_len), self.hf_tokenizer.pad_token_id, dtype=np.int64)
+        text_mat = np.full((len(samples), max_seq_len), pad_token_id, dtype=np.int64)
+        target_mat = np.full((len(samples), max_seq_len), pad_token_id, dtype=np.int64)
 
         image_input_masks = np.zeros_like(text_mat, dtype=bool)
         video_input_masks = np.zeros_like(text_mat, dtype=bool)
         for i, s in enumerate(samples):
-            # If the sample/target length exceeds the target sequence length, then truncate.
             text_len = min(max_seq_len, len(s.text))
             target_len = min(max_seq_len, len(s.target))
 
             text_mat[i, :text_len] = np.array(s.text)[:text_len]
-            # NOTE: we should assert user input sequence will not be truncated
             if s.image_input_mask is not None:
                 image_input_masks[i, :text_len] = np.array(s.image_input_mask)[:text_len]
             if s.video_input_mask is not None:
@@ -395,10 +686,10 @@ def batch(self, samples: List[QwenVLTaskSample]) -> QwenVLTaskBatch:
             target_mat[i, :target_len] = np.array(s.target)[:target_len]
 
         tokens = torch.from_numpy(text_mat)
-        tokens[tokens == self.hf_tokenizer.pad_token_id] = 0
+        tokens[tokens == pad_token_id] = 0
 
         labels = torch.from_numpy(target_mat)
-        labels[labels == self.hf_tokenizer.pad_token_id] = IGNORE_INDEX
+        labels[labels == pad_token_id] = IGNORE_INDEX
 
         attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
             data=tokens,
@@ -410,7 +701,7 @@ def batch(self, samples: List[QwenVLTaskSample]) -> QwenVLTaskBatch:
 
         loss_mask[labels < 0] = 0.0
 
-        batch = QwenVLTaskBatch(
+        batch_obj = QwenVLTaskBatch(
             __keys__=[s.__key__ for s in samples],
             __subflavors__=[s.__subflavors__ for s in samples],
             pixel_values=torch.vstack(imgs) if len(imgs) > 0 else None,
@@ -425,7 +716,7 @@ def batch(self, samples: List[QwenVLTaskSample]) -> QwenVLTaskBatch:
             labels=labels,
             loss_mask=loss_mask,
         )
-        return batch
+        return batch_obj
 
     def encode_batch(self, batch: QwenVLTaskBatch) -> dict:
         """Encode batch in dict"""
@@ -438,4 +729,9 @@ def encode_batch(self, batch: QwenVLTaskBatch) -> dict:
             image_grid_thw=batch.image_grid_thw,
         )
 
+        if batch.cu_lengths is not None:
+            raw["cu_seqlens"] = batch.cu_lengths
+            raw["max_seqlen"] = batch.max_lengths
+            raw["cu_seqlens_argmin"] = batch.cu_seqlens_argmin
+
         return raw
diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index c3622c9815..2481d7d80f 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -335,8 +335,51 @@ def _ceil_to_mult(n: int, mult: int) -> int:
     tp_size = pg_collection.tp.size() if pg_collection is not None and pg_collection.tp is not None else 1
     has_sp = getattr(cfg.model, "sequence_parallel", False)
 
+    # Energon pre-packed path: cu_seqlens already present from TaskEncoder packing
+    energon_cu_seqlens = batch.get("cu_seqlens")
+    if energon_cu_seqlens is not None:
+        tokens_or_input = batch.get("tokens") if batch.get("tokens") is not None else batch.get("input_ids")
+        energon_max_seqlen = batch.get("max_seqlen")
+        energon_cu_argmin = batch.get("cu_seqlens_argmin")
+
+        # Log detailed iteration stats for FLOPs analysis
+        if tokens_or_input is not None:
+            _seq_dim = tokens_or_input.shape[-1]
+            if energon_cu_seqlens.dim() == 1:
+                _content_len = energon_cu_seqlens[-1].item()
+            elif energon_cu_argmin is not None:
+                _argmin_val = energon_cu_argmin.item() if energon_cu_argmin.dim() == 0 else energon_cu_argmin[0].item()
+                _content_len = energon_cu_seqlens[0, _argmin_val - 1].item()
+            else:
+                _content_len = energon_cu_seqlens[0, -1].item()
+            _pad_len = _seq_dim - _content_len
+            _n_img_toks = int((tokens_or_input == 151655).sum().item())
+            _vit_shape = (
+                visual_inputs.pixel_values.shape
+                if visual_inputs is not None and visual_inputs.pixel_values is not None
+                else "None"
+            )
+            logger.info(
+                f"[IterStats] decoder_input={list(tokens_or_input.shape)}, "
+                f"content={_content_len}, pad={_pad_len}, "
+                f"image_tokens_in_decoder={_n_img_toks}, "
+                f"vit_input={_vit_shape}"
+            )
+
+        return (
+            tokens_or_input,
+            batch.get("labels"),
+            batch.get("loss_mask"),
+            batch.get("attention_mask"),
+            batch.get("position_ids"),
+            energon_cu_seqlens,
+            energon_max_seqlen,
+            visual_inputs,
+            energon_cu_argmin,
+        )
+
     if enable_packing:
-        # Pack sequences
+        # In-batch packing (concatenate all micro-batch sequences)
         tokens_or_input = batch.get("tokens") if batch.get("tokens") is not None else batch.get("input_ids")
 
         # Compute pad_to_multiple_of as lcm of CP and SP constraints.
@@ -366,7 +409,6 @@ def _ceil_to_mult(n: int, mult: int) -> int:
             padding_mask=batch.get("_padding_mask"),
         )
 
-        # Update batch dict with packed tensors
         if batch.get("tokens") is not None:
             batch["tokens"] = packed_tokens
         else:
@@ -376,10 +418,8 @@ def _ceil_to_mult(n: int, mult: int) -> int:
         batch["attention_mask"] = packed_attention_mask
         batch["position_ids"] = packed_position_ids
 
-        # # Add packing metadata
-        logger.debug(f"Packed batch: cu_seqlens={cu_seqlens.tolist()}, max_seqlen={max_seqlen}")
+        logger.debug(f"In-batch packed: cu_seqlens={cu_seqlens.tolist()}, max_seqlen={max_seqlen}")
     else:
-        # No packing, use dummy values
         cu_seqlens = None
         max_seqlen = None
 
@@ -392,6 +432,7 @@ def _ceil_to_mult(n: int, mult: int) -> int:
         cu_seqlens,
         max_seqlen,
         visual_inputs,
+        None,  # cu_seqlens_argmin (only set for Energon pre-packed path)
     )
 
 
@@ -427,6 +468,7 @@ def forward_step(
             cu_seqlens,
             max_seqlen,
             visual_inputs,
+            cu_seqlens_argmin,
         ) = get_batch(data_iterator, state.cfg, use_mtp, pg_collection=pg_collection)
     timers("batch-generator").stop()
 
@@ -443,13 +485,15 @@ def forward_step(
 
     # Add packed sequence support
     if cu_seqlens is not None:
-        cu_seqlens_argmin = torch.tensor(len(cu_seqlens))  # no padding in cu_seqlens since packing is done in-batch
-        packed_seq_params = {
+        packed_seq_dict = {
             "cu_seqlens": cu_seqlens,
             "max_seqlen": max_seqlen,
-            "cu_seqlens_argmin": cu_seqlens_argmin,
         }
-        forward_args["packed_seq_params"] = get_packed_seq_params(packed_seq_params)
+        if cu_seqlens_argmin is not None:
+            packed_seq_dict["cu_seqlens_argmin"] = cu_seqlens_argmin
+        elif cu_seqlens.dim() == 1:
+            packed_seq_dict["cu_seqlens_argmin"] = torch.tensor(len(cu_seqlens))
+        forward_args["packed_seq_params"] = get_packed_seq_params(packed_seq_dict)
 
     if loss_mask is not None:
         loss_mask = loss_mask.contiguous()

From 6af4f0f2b2d613c5e3d115e47472f868b22881fc Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 31 Mar 2026 15:22:30 +0800
Subject: [PATCH 10/51] feat: Add Energon SFT recipe for Qwen3.5-VL 35B-A3B
 with THD packing

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py b/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
index c1bb8fd4ca..44c9fbd57c 100644
--- a/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
+++ b/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
@@ -26,6 +26,7 @@
 from typing_extensions import Unpack
 
 from megatron.bridge import AutoBridge
+from megatron.bridge.data.energon.energon_provider import EnergonProvider
 from megatron.bridge.peft.base import PEFT
 from megatron.bridge.recipes.common import _peft_common_vlm, _sft_common_vlm
 from megatron.bridge.recipes.qwen_vl.qwen3_vl import (

From fd4215f71eee90adf3c9ed5864e8303c2367f0a8 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Wed, 1 Apr 2026 13:54:15 +0800
Subject: [PATCH 11/51] fix: QwenVLTaskBatch missing Energon 7.x
 __key__/__restore_key__ and sync TaskEncoder seq_len from CLI override

- QwenVLTaskBatch: pass __key__ and __restore_key__ to satisfy Energon Batch base class
- energon_provider: sync TaskEncoder.seq_len from EnergonProvider.seq_length at build time
  so CLI overrides (e.g. dataset.seq_length=8192) propagate to packing capacity and padding

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../bridge/data/energon/energon_provider.py    |  3 +++
 .../qwen_vl/data/energon/task_encoder.py       | 18 +++++++++++++++---
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/src/megatron/bridge/data/energon/energon_provider.py b/src/megatron/bridge/data/energon/energon_provider.py
index 2afcd0d1b9..07de6a9328 100644
--- a/src/megatron/bridge/data/energon/energon_provider.py
+++ b/src/megatron/bridge/data/energon/energon_provider.py
@@ -39,6 +39,9 @@ class EnergonProvider(DatasetProvider):
 
     def build_datasets(self, context: DatasetBuildContext):
         assert self.path, "EnergonProvider.path must be set. Use CLI override: dataset.path=<path>"
+        if self.task_encoder is not None and hasattr(self.task_encoder, "seq_len"):
+            self.task_encoder.seq_len = self.seq_length
+            self.task_encoder.seq_length = self.seq_length
         dataset = EnergonMultiModalDataModule(
             path=self.path,
             tokenizer=context.tokenizer if context.tokenizer is not None else self.tokenizer,
diff --git a/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py b/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
index 13e90877f0..9795da52db 100644
--- a/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
+++ b/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
@@ -173,7 +173,13 @@ class QwenVLTaskSamplePacked:
 
 @dataclass
 class QwenVLTaskBatch(Batch):
-    """Encoded Batch Format For QwenVL"""
+    """Encoded Batch Format For QwenVL.
+
+    Energon 7.x ``Batch`` requires ``__key__`` and ``__restore_key__`` as
+    keyword-only init fields.  We keep ``__keys__`` (plural) for downstream
+    consumption but satisfy the parent via ``__key__`` (joined) and
+    ``__restore_key__`` (empty tuple placeholder).
+    """
 
     __keys__: List[str]
     __subflavors__: List[Dict]
@@ -644,8 +650,11 @@ def batch(self, samples: List[Union[QwenVLTaskSample, QwenVLTaskSamplePacked]])
             cu_seqlens_argmin = torch.tensor(cu_seqlens_argmin_list, dtype=torch.int32)
             max_lengths = torch.tensor(max_lengths_list, dtype=torch.int32)
 
+            sample_keys = [s.__key__ for s in samples]
             batch_obj = QwenVLTaskBatch(
-                __keys__=[s.__key__ for s in samples],
+                __key__="+".join(sample_keys),
+                __restore_key__=(),
+                __keys__=sample_keys,
                 __subflavors__=[s.__subflavors__ for s in samples],
                 pixel_values=torch.vstack(imgs) if len(imgs) > 0 else None,
                 pixel_values_videos=torch.vstack(videos) if len(videos) > 0 else None,
@@ -701,8 +710,11 @@ def batch(self, samples: List[Union[QwenVLTaskSample, QwenVLTaskSamplePacked]])
 
         loss_mask[labels < 0] = 0.0
 
+        sample_keys = [s.__key__ for s in samples]
         batch_obj = QwenVLTaskBatch(
-            __keys__=[s.__key__ for s in samples],
+            __key__="+".join(sample_keys),
+            __restore_key__=(),
+            __keys__=sample_keys,
             __subflavors__=[s.__subflavors__ for s in samples],
             pixel_values=torch.vstack(imgs) if len(imgs) > 0 else None,
             pixel_values_videos=torch.vstack(videos) if len(videos) > 0 else None,

From cfffa6aa652d7cf7b3e76b266c81e410a4f29f93 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Thu, 2 Apr 2026 14:03:05 +0800
Subject: [PATCH 12/51] fix: patch MTP attention with Qwen3VLSelfAttention for
 correct M-RoPE

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../models/qwen_vl/qwen35_vl_provider.py      | 37 ++++++++++++++++++-
 1 file changed, 35 insertions(+), 2 deletions(-)

diff --git a/src/megatron/bridge/models/qwen_vl/qwen35_vl_provider.py b/src/megatron/bridge/models/qwen_vl/qwen35_vl_provider.py
index d2dcec2704..18c0472b24 100644
--- a/src/megatron/bridge/models/qwen_vl/qwen35_vl_provider.py
+++ b/src/megatron/bridge/models/qwen_vl/qwen35_vl_provider.py
@@ -195,6 +195,10 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> Qwen3VL
         )
         _patch_standard_attention_specs(block_spec, Qwen3VLSelfAttention)
 
+        mtp_spec = mtp_block_spec(self, vp_stage=vp_stage)
+        if mtp_spec is not None:
+            _patch_mtp_attention_specs(mtp_spec, Qwen3VLSelfAttention)
+
         model = Qwen3VLModel(
             language_transformer_config=language_transformer_config,
             language_transformer_layer_spec=block_spec,
@@ -202,7 +206,7 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> Qwen3VL
             pre_process=pre_process,
             post_process=post_process,
             pg_collection=self._pg_collection,
-            mtp_block_spec=mtp_block_spec(self, vp_stage=vp_stage),
+            mtp_block_spec=mtp_spec,
             vp_stage=vp_stage,
         )
 
@@ -383,6 +387,10 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> Qwen3VL
         # with Qwen3VLSelfAttention for mRoPE support. GDN layers are left as-is.
         _patch_standard_attention_specs(block_spec, Qwen3VLSelfAttention)
 
+        mtp_spec = mtp_block_spec(self, vp_stage=vp_stage)
+        if mtp_spec is not None:
+            _patch_mtp_attention_specs(mtp_spec, Qwen3VLSelfAttention)
+
         model = Qwen3VLModel(
             language_transformer_config=language_transformer_config,
             language_transformer_layer_spec=block_spec,
@@ -390,7 +398,7 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> Qwen3VL
             pre_process=pre_process,
             post_process=post_process,
             pg_collection=self._pg_collection,
-            mtp_block_spec=mtp_block_spec(self, vp_stage=vp_stage),
+            mtp_block_spec=mtp_spec,
             vp_stage=vp_stage,
         )
 
@@ -436,3 +444,28 @@ def _patch_standard_attention_specs(
             isinstance(attn_spec.module, type) and issubclass(attn_spec.module, SelfAttention)
         ):
             attn_spec.module = attention_cls
+
+
+def _patch_mtp_attention_specs(
+    mtp_block_spec,
+    attention_cls,
+) -> None:
+    """Patch MTP layer specs so their inner attention uses *attention_cls*.
+
+    MTP layers are independently constructed from the main decoder spec and
+    therefore miss the ``Qwen3VLSelfAttention`` override applied to the decoder.
+    Without this patch, MTP attention falls back to the base ``SelfAttention``
+    which uses ``_apply_rotary_pos_emb_thd`` (position-resetting per sub-sequence)
+    instead of the Qwen3VL absolute M-RoPE implementation.
+    """
+    from megatron.core.transformer.attention import SelfAttention
+
+    for mtp_layer_spec in mtp_block_spec.layer_specs:
+        inner_layer = mtp_layer_spec.submodules.mtp_model_layer
+        if inner_layer is None or not hasattr(inner_layer, 'submodules'):
+            continue
+        attn_spec = inner_layer.submodules.self_attention
+        if attn_spec.module is SelfAttention or (
+            isinstance(attn_spec.module, type) and issubclass(attn_spec.module, SelfAttention)
+        ):
+            attn_spec.module = attention_cls

From a6cd85f8cc40c79836f9f38d0350fe3f0854e01b Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 7 Apr 2026 10:16:47 +0800
Subject: [PATCH 13/51] fix: correctly pass padded/unpadded cu_seqlens for
 packed sequence THD format

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 src/megatron/bridge/training/vlm_step.py | 54 +++++++++++++++++++++---
 1 file changed, 47 insertions(+), 7 deletions(-)

diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index 2481d7d80f..e515101567 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -70,8 +70,12 @@ def get_batch_from_iterator(
 
     if "cu_seqlens" in batch:
         required_device_keys.add("cu_seqlens")
+        if "cu_seqlens_unpadded" in batch:
+            required_device_keys.add("cu_seqlens_unpadded")
         required_host_keys.add("cu_seqlens_argmin")
         required_host_keys.add("max_seqlen")
+        if "cu_seqlens_unpadded_argmin" in batch:
+            required_host_keys.add("cu_seqlens_unpadded_argmin")
 
     required_device_keys.update(("tokens", "input_ids", "position_ids"))
     if is_last_pp_stage:
@@ -485,14 +489,50 @@ def forward_step(
 
     # Add packed sequence support
     if cu_seqlens is not None:
-        packed_seq_dict = {
-            "cu_seqlens": cu_seqlens,
-            "max_seqlen": max_seqlen,
-        }
+        physical_seq_len = tokens.shape[-1]
+
+        # Determine the number of valid entries in cu_seqlens.
+        cu = cu_seqlens.squeeze()
         if cu_seqlens_argmin is not None:
-            packed_seq_dict["cu_seqlens_argmin"] = cu_seqlens_argmin
-        elif cu_seqlens.dim() == 1:
-            packed_seq_dict["cu_seqlens_argmin"] = torch.tensor(len(cu_seqlens))
+            n_valid = int(cu_seqlens_argmin.item())
+        elif cu.dim() == 1:
+            n_valid = len(cu)
+        else:
+            n_valid = int(torch.argmin(cu).item())
+
+        cu_clean = cu[:n_valid]
+        last_boundary = int(cu_clean[-1].item())
+
+        if last_boundary < physical_seq_len:
+            # Trailing padding exists: build padded/unpadded pair following sft.py convention.
+            # Padded: append physical_seq_len (padding segment covers the tail).
+            # Unpadded: repeat last boundary (padding segment has 0 real tokens).
+            device, dtype = cu_clean.device, cu_clean.dtype
+            cu_padded = torch.cat([cu_clean, torch.tensor([physical_seq_len], device=device, dtype=dtype)])
+            cu_unpadded = torch.cat([cu_clean, torch.tensor([last_boundary], device=device, dtype=dtype)])
+            new_n_valid = torch.tensor(n_valid + 1)
+
+            padded_diffs = cu_padded[1:] - cu_padded[:-1]
+            max_seqlen_padded = padded_diffs.max()
+
+            packed_seq_dict = {
+                "cu_seqlens": cu_padded,
+                "cu_seqlens_argmin": new_n_valid,
+                "max_seqlen": max_seqlen_padded,
+                "cu_seqlens_unpadded": cu_unpadded,
+                "cu_seqlens_unpadded_argmin": new_n_valid,
+            }
+        else:
+            # No trailing padding (content fills the full sequence length).
+            packed_seq_dict = {
+                "cu_seqlens": cu_seqlens,
+                "max_seqlen": max_seqlen,
+            }
+            if cu_seqlens_argmin is not None:
+                packed_seq_dict["cu_seqlens_argmin"] = cu_seqlens_argmin
+            elif cu_seqlens.dim() == 1:
+                packed_seq_dict["cu_seqlens_argmin"] = torch.tensor(len(cu_seqlens))
+
         forward_args["packed_seq_params"] = get_packed_seq_params(packed_seq_dict)
 
     if loss_mask is not None:

From 1bafcda215b758b598152310801dae26838276bf Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 7 Apr 2026 10:54:50 +0800
Subject: [PATCH 14/51] fix: extend last segment boundary instead of adding
 zero-length padding segment to avoid GDN NaN

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 src/megatron/bridge/training/vlm_step.py | 19 +++++++++----------
 1 file changed, 9 insertions(+), 10 deletions(-)

diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index e515101567..76f30ae9bd 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -504,23 +504,22 @@ def forward_step(
         last_boundary = int(cu_clean[-1].item())
 
         if last_boundary < physical_seq_len:
-            # Trailing padding exists: build padded/unpadded pair following sft.py convention.
-            # Padded: append physical_seq_len (padding segment covers the tail).
-            # Unpadded: repeat last boundary (padding segment has 0 real tokens).
-            device, dtype = cu_clean.device, cu_clean.dtype
-            cu_padded = torch.cat([cu_clean, torch.tensor([physical_seq_len], device=device, dtype=dtype)])
-            cu_unpadded = torch.cat([cu_clean, torch.tensor([last_boundary], device=device, dtype=dtype)])
-            new_n_valid = torch.tensor(n_valid + 1)
+            # Trailing padding exists: extend the last real segment's padded
+            # boundary to cover the tail padding instead of creating an extra
+            # zero-length segment (which would cause GDN FLA kernels to process
+            # an isolated padding-only segment and diverge to NaN).
+            cu_padded = cu_clean.clone()
+            cu_padded[-1] = physical_seq_len
 
             padded_diffs = cu_padded[1:] - cu_padded[:-1]
             max_seqlen_padded = padded_diffs.max()
 
             packed_seq_dict = {
                 "cu_seqlens": cu_padded,
-                "cu_seqlens_argmin": new_n_valid,
+                "cu_seqlens_argmin": torch.tensor(n_valid),
                 "max_seqlen": max_seqlen_padded,
-                "cu_seqlens_unpadded": cu_unpadded,
-                "cu_seqlens_unpadded_argmin": new_n_valid,
+                "cu_seqlens_unpadded": cu_clean,
+                "cu_seqlens_unpadded_argmin": torch.tensor(n_valid),
             }
         else:
             # No trailing padding (content fills the full sequence length).

From be4105d2f94e0dc885492b681d7818c4375b2839 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 7 Apr 2026 11:32:57 +0800
Subject: [PATCH 15/51] fix: use padded cu_seqlens for all fields following
 Megatron-LM data_schedule.py convention

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../bridge/training/utils/packed_seq_utils.py        |  6 ++++++
 src/megatron/bridge/training/vlm_step.py             | 12 +++++++-----
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/src/megatron/bridge/training/utils/packed_seq_utils.py b/src/megatron/bridge/training/utils/packed_seq_utils.py
index 1194fc7ed2..78d5dd17cd 100644
--- a/src/megatron/bridge/training/utils/packed_seq_utils.py
+++ b/src/megatron/bridge/training/utils/packed_seq_utils.py
@@ -71,9 +71,15 @@ def get_packed_seq_params(batch: dict[str, torch.Tensor]) -> PackedSeqParams:
             qkv_format="thd",
         )
     else:
+        # Follow Megatron-LM data_schedule.py convention: set all four
+        # cu_seqlens fields to padded values so every downstream op (RoPE,
+        # TE attention, GDN) sees consistent boundaries without relying on
+        # fallback logic.
         return PackedSeqParams(
             cu_seqlens_q=cu_seqlens_padded,
             cu_seqlens_kv=cu_seqlens_padded,
+            cu_seqlens_q_padded=cu_seqlens_padded,
+            cu_seqlens_kv_padded=cu_seqlens_padded,
             max_seqlen_q=max_seqlen,
             max_seqlen_kv=max_seqlen,
             qkv_format="thd",
diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index 76f30ae9bd..7835bd5292 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -505,9 +505,13 @@ def forward_step(
 
         if last_boundary < physical_seq_len:
             # Trailing padding exists: extend the last real segment's padded
-            # boundary to cover the tail padding instead of creating an extra
-            # zero-length segment (which would cause GDN FLA kernels to process
-            # an isolated padding-only segment and diverge to NaN).
+            # boundary to cover the tail so that all layers (TE attention, GDN,
+            # etc.) see the same sequence boundaries.  We intentionally do NOT
+            # set cu_seqlens_unpadded here so that cu_seqlens_q_padded stays
+            # None — this forces every layer to treat pad tokens identically to
+            # real tokens, avoiding undefined attention output at pad positions
+            # that would otherwise propagate NaN through subsequent layers and
+            # into the backward pass.  loss_mask already zeros pad positions.
             cu_padded = cu_clean.clone()
             cu_padded[-1] = physical_seq_len
 
@@ -518,8 +522,6 @@ def forward_step(
                 "cu_seqlens": cu_padded,
                 "cu_seqlens_argmin": torch.tensor(n_valid),
                 "max_seqlen": max_seqlen_padded,
-                "cu_seqlens_unpadded": cu_clean,
-                "cu_seqlens_unpadded_argmin": torch.tensor(n_valid),
             }
         else:
             # No trailing padding (content fills the full sequence length).

From 32ae832b28e5e49f9f9f2adc81bcfc21bc98b7ec Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Wed, 8 Apr 2026 09:26:15 +0800
Subject: [PATCH 16/51] fix: compute MRoPE per sub-sequence for packed
 sequences to fix 22x loss regression

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../qwen_vl/modelling_qwen3_vl/model.py       | 73 +++++++++++++++----
 src/megatron/bridge/training/vlm_step.py      | 43 +++++++++++
 2 files changed, 101 insertions(+), 15 deletions(-)

diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
index 5b0e8424e7..46f01fea76 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
@@ -590,22 +590,50 @@ def forward(
                 )
 
         if position_ids is None:
-            # BSHD
-            # Megatron uses 4D bool masks ([B|1,1,S,S], True=masked); HF uses 2D keep masks ([B,S], 1=keep)
-            # For simplicity, we set hf_attention_mask to None.
-            hf_attention_mask = None
-            position_ids, _ = get_rope_index(
-                self.config.spatial_merge_size,
-                self.image_token_id,
-                self.video_token_id,
-                self.vision_start_token_id,
-                input_ids,
-                image_grid_thw=image_grid_thw,
-                video_grid_thw=video_grid_thw,
-                attention_mask=hf_attention_mask,
-            )  #  [3*b*s]
             if packed_seq_params is not None:
-                # convert position_ids to THD format
+                # Packed sequences: compute MRoPE per sub-sequence independently
+                # so that each sub-sequence's positions restart from 0.
+                cu = packed_seq_params.cu_seqlens_q
+                if cu.dim() > 1:
+                    cu = cu.squeeze()
+                seq_lens = (cu[1:] - cu[:-1]).tolist()
+                num_seqs = len(seq_lens)
+                max_subseq_len = int(max(seq_lens))
+                total_len = input_ids.shape[1]
+                flat_ids = input_ids.squeeze(0)
+
+                sub_ids = torch.zeros(
+                    num_seqs, max_subseq_len, dtype=input_ids.dtype, device=input_ids.device
+                )
+                sub_mask = torch.zeros(
+                    num_seqs, max_subseq_len, dtype=torch.int32, device=input_ids.device
+                )
+                for i, sl in enumerate(seq_lens):
+                    start = int(cu[i].item())
+                    sl_int = int(sl)
+                    sub_ids[i, :sl_int] = flat_ids[start : start + sl_int]
+                    sub_mask[i, :sl_int] = 1
+
+                position_ids, _ = get_rope_index(
+                    self.config.spatial_merge_size,
+                    self.image_token_id,
+                    self.video_token_id,
+                    self.vision_start_token_id,
+                    sub_ids,
+                    image_grid_thw=image_grid_thw,
+                    video_grid_thw=video_grid_thw,
+                    attention_mask=sub_mask,
+                )  # [3, num_seqs, max_subseq_len]
+
+                packed_pos = torch.zeros(
+                    3, 1, total_len, dtype=position_ids.dtype, device=position_ids.device
+                )
+                for i, sl in enumerate(seq_lens):
+                    start = int(cu[i].item())
+                    sl_int = int(sl)
+                    packed_pos[:, 0, start : start + sl_int] = position_ids[:, i, :sl_int]
+                position_ids = packed_pos
+
                 position_ids = (
                     preprocess_packed_seqs(
                         position_ids.permute(1, 2, 0),
@@ -618,6 +646,21 @@ def forward(
                 )
                 attention_mask = None
                 self.language_model.rotary_pos_emb.is_thd_format = True
+            else:
+                # BSHD
+                # Megatron uses 4D bool masks ([B|1,1,S,S], True=masked); HF uses 2D keep masks ([B,S], 1=keep)
+                # For get_rope_index we pass None to avoid semantic mismatch.
+                hf_attention_mask = None
+                position_ids, _ = get_rope_index(
+                    self.config.spatial_merge_size,
+                    self.image_token_id,
+                    self.video_token_id,
+                    self.vision_start_token_id,
+                    input_ids,
+                    image_grid_thw=image_grid_thw,
+                    video_grid_thw=video_grid_thw,
+                    attention_mask=hf_attention_mask,
+                )  #  [3*b*s]
 
         torch.cuda.nvtx.range_pop()
         torch.cuda.nvtx.range_push("Qwen3VLModel.forward.language_model")
diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index 7835bd5292..074639a42c 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -476,6 +476,49 @@ def forward_step(
         ) = get_batch(data_iterator, state.cfg, use_mtp, pg_collection=pg_collection)
     timers("batch-generator").stop()
 
+    # ---- Unified debug dump for first 2 iterations (both BSHD and Energon) ----
+    _iter = getattr(state, "_debug_iter_count", 0)
+    if _iter < 2 and tokens is not None:
+        _is_energon = cu_seqlens is not None
+        _mode = "Energon-THD" if _is_energon else "HF-BSHD"
+        _t = tokens.squeeze()
+        _lm = loss_mask.squeeze() if loss_mask is not None else None
+        _lab = labels.squeeze() if labels is not None else None
+        _sup_count = int((_lab != -100).sum().item()) if _lab is not None else 0
+        _lm_sum = int(_lm.sum().item()) if _lm is not None else 0
+        logger.info(
+            f"\n{'='*70}\n"
+            f"[DEBUG {_mode}] iteration {_iter}\n"
+            f"  tokens.shape        = {list(tokens.shape)}\n"
+            f"  tokens[:20]         = {_t[:20].tolist()}\n"
+            f"  tokens[-10:]        = {_t[-10:].tolist()}\n"
+            f"  labels[:20]         = {_lab[:20].tolist() if _lab is not None else None}\n"
+            f"  loss_mask sum/total = {_lm_sum} / {_t.shape[0]}\n"
+            f"  supervised labels   = {_sup_count}\n"
+            f"  cu_seqlens (raw)    = {cu_seqlens}\n"
+            f"  position_ids[:10]   = {position_ids.squeeze()[:10].tolist() if position_ids is not None else None}\n"
+            f"  visual_inputs       = {'present' if visual_inputs is not None else 'None'}\n"
+            f"{'='*70}"
+        )
+        # Label alignment check: labels[i] should equal tokens[i+1] at supervised positions
+        if _lab is not None:
+            _sup_idx = (_lab != -100).nonzero(as_tuple=True)[0]
+            if len(_sup_idx) > 0:
+                _valid = _sup_idx[_sup_idx < len(_t) - 1]
+                _match = int((_lab[_valid] == _t[_valid + 1]).sum().item())
+                logger.info(
+                    f"  label alignment: labels[i]==tokens[i+1] at {_match}/{len(_valid)} supervised positions "
+                    f"(first_sup_pos={_sup_idx[0].item()}, last_sup_pos={_sup_idx[-1].item()})"
+                )
+                if _match < len(_valid):
+                    _mismatch_idx = _valid[_lab[_valid] != _t[_valid + 1]][:5]
+                    for _mi in _mismatch_idx:
+                        _mi = _mi.item()
+                        logger.info(
+                            f"    MISMATCH pos {_mi}: label={_lab[_mi].item()} vs tokens[{_mi+1}]={_t[_mi+1].item()}"
+                        )
+        state._debug_iter_count = _iter + 1
+
     forward_args = {
         "input_ids": tokens,
         "position_ids": position_ids,

From 198b12e6c595fe429b680f4930f0249b4e7fb78e Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Wed, 8 Apr 2026 09:55:54 +0800
Subject: [PATCH 17/51] fix: force Qwen-VL to recompute MRoPE in-model for THD

Pass position_ids=None for VLM batches so Qwen3VLModel runs its packed MRoPE path, and add short-lived debug logs to verify packed boundary position resets are actually applied.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../qwen_vl/modelling_qwen3_vl/model.py       | 23 +++++++++++++++++++
 src/megatron/bridge/training/vlm_step.py      | 11 ++++++++-
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
index 46f01fea76..7839123615 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 from typing import Optional
 
 import torch
@@ -52,6 +53,9 @@
 from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.vision_model import Qwen3VLVisionModel
 
 
+logger = logging.getLogger(__name__)
+
+
 class Qwen3VLModel(MegatronModule):
     """Qwen3VL multi-modal model.
 
@@ -203,6 +207,7 @@ def __init__(
             assert self.config.calculate_per_token_loss, (
                 "Qwen3-VL model only supports context parallelism with calculate_per_token_loss enabled"
             )
+        self._mrope_debug_count = 0
 
         self._expose_language_model_for_cuda_graph_helper()
 
@@ -632,6 +637,24 @@ def forward(
                     start = int(cu[i].item())
                     sl_int = int(sl)
                     packed_pos[:, 0, start : start + sl_int] = position_ids[:, i, :sl_int]
+
+                # Debug only the first two forwards on rank-0 to verify packed MRoPE reset behavior.
+                if self._mrope_debug_count < 2:
+                    rank = (
+                        torch.distributed.get_rank()
+                        if torch.distributed.is_available() and torch.distributed.is_initialized()
+                        else 0
+                    )
+                    if rank == 0:
+                        boundary_pos0 = [int(cu[i].item()) for i in range(len(seq_lens))]
+                        boundary_posvals = [int(packed_pos[0, 0, p].item()) for p in boundary_pos0]
+                        logger.info(
+                            "[MRoPE-DEBUG] packed branch hit: cu=%s seq_lens=%s boundary_pos0=%s",
+                            cu.tolist(),
+                            [int(x) for x in seq_lens],
+                            boundary_posvals,
+                        )
+                    self._mrope_debug_count += 1
                 position_ids = packed_pos
 
                 position_ids = (
diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index 074639a42c..906969da45 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -519,9 +519,18 @@ def forward_step(
                         )
         state._debug_iter_count = _iter + 1
 
+    # For Qwen-VL, let model compute MRoPE position_ids from multimodal tokens.
+    # Passing 1D sequential position_ids from dataloader would bypass that path.
+    model_position_ids = None if visual_inputs is not None else position_ids
+    if _iter < 2:
+        logger.info(
+            "  forward position_ids mode = %s",
+            "None (compute MRoPE in model)" if model_position_ids is None else "provided-from-dataloader",
+        )
+
     forward_args = {
         "input_ids": tokens,
-        "position_ids": position_ids,
+        "position_ids": model_position_ids,
         "attention_mask": attention_mask,
         "labels": labels,
         "loss_mask": loss_mask,  # Pass full loss_mask so model can slice it consistently with labels

From cf7c7c51f0a599a1d891fc23f6efcea3d03acdb1 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Mon, 13 Apr 2026 13:55:27 +0800
Subject: [PATCH 18/51] fix: restore THD diagnostic helper symbols in Qwen3VL
 model

Add local helper functions used by THD diagnostic logging after cherry-picking THD MRoPE fixes, preventing runtime NameError for _thd_diag_enabled/_rank0.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../bridge/models/qwen_vl/modelling_qwen3_vl/model.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
index 7839123615..40a61cbe71 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import logging
+import os
 from typing import Optional
 
 import torch
@@ -56,6 +57,16 @@
 logger = logging.getLogger(__name__)
 
 
+def _thd_diag_enabled() -> bool:
+    return os.environ.get("THD_DIAG", "0") not in ("0", "", "false", "False")
+
+
+def _rank0() -> bool:
+    if not torch.distributed.is_available() or not torch.distributed.is_initialized():
+        return True
+    return torch.distributed.get_rank() == 0
+
+
 class Qwen3VLModel(MegatronModule):
     """Qwen3VL multi-modal model.
 

From cbeb228b38bd7958c692e513655be910d0e574a6 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Mon, 13 Apr 2026 15:08:31 +0800
Subject: [PATCH 19/51] fix: harden Energon QwenVL batch assembly for packed
 visual inputs

Add missing Energon batch key fields and normalize/concat visual blocks along token dimension so packed multimodal batches no longer fail on mixed visual lengths.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../qwen_vl/data/energon/task_encoder.py      | 57 ++++++++++++++++---
 1 file changed, 49 insertions(+), 8 deletions(-)

diff --git a/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py b/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
index 9795da52db..882de547b6 100644
--- a/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
+++ b/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
@@ -181,6 +181,8 @@ class QwenVLTaskBatch(Batch):
     ``__restore_key__`` (empty tuple placeholder).
     """
 
+    __key__: str
+    __restore_key__: tuple
     __keys__: List[str]
     __subflavors__: List[Dict]
     # (num_tiles, c, h, w)
@@ -538,6 +540,45 @@ def pack_selected_samples(self, samples: List[QwenVLTaskSample]) -> QwenVLTaskSa
             num_vit_patches=total_vit_patches,
         )
 
+    @staticmethod
+    def _normalize_visual_block(block: torch.Tensor) -> torch.Tensor:
+        """Normalize visual block shape before batch concat.
+
+        Expected downstream format is either:
+        - [N, C, H, W] (patch tiles), or
+        - [N, D] (flattened patch features).
+        """
+        if not isinstance(block, torch.Tensor):
+            raise TypeError(f"Visual block must be torch.Tensor, got {type(block)}")
+
+        # Common case for processor output: [1, N, D] -> [N, D]
+        if block.dim() == 3 and block.shape[0] == 1:
+            return block.squeeze(0)
+
+        # Single tile CHW -> add batch dimension.
+        if block.dim() == 3 and block.shape[0] in (1, 3):
+            return block.unsqueeze(0)
+
+        # Keep already batched formats as-is.
+        if block.dim() in (2, 4):
+            return block
+
+        return block
+
+    @staticmethod
+    def _concat_visual_blocks(blocks: List[torch.Tensor], name: str) -> Optional[torch.Tensor]:
+        """Concatenate visual blocks along first dimension with shape diagnostics."""
+        if not blocks:
+            return None
+
+        normalized = [QwenVLTaskEncoder._normalize_visual_block(b) for b in blocks]
+        try:
+            return torch.cat(normalized, dim=0)
+        except RuntimeError as exc:
+            shapes = [tuple(t.shape) for t in normalized]
+            logger.error(f"Failed to concat {name} blocks, shapes={shapes}, error={exc}")
+            raise
+
     def _collect_visual_data(self, samples):
         """Gather pixel_values and grid_thw from a list of samples (packed or unpacked)."""
         imgs, image_thw_grids = [], []
@@ -547,17 +588,17 @@ def _collect_visual_data(self, samples):
             if isinstance(s, QwenVLTaskSamplePacked):
                 for img_block in s.imgs:
                     if isinstance(img_block, torch.Tensor):
-                        imgs.append(img_block if img_block.dim() == 4 else img_block.unsqueeze(0))
+                        imgs.append(self._normalize_visual_block(img_block))
                 for thw in s.image_thw_grids:
                     image_thw_grids.append(np.asarray(thw))
                 for vid_block in s.videos:
                     if isinstance(vid_block, torch.Tensor):
-                        videos.append(vid_block if vid_block.dim() == 4 else vid_block.unsqueeze(0))
+                        videos.append(self._normalize_visual_block(vid_block))
                 for thw in s.video_thw_grids:
                     video_thw_grids.append(np.asarray(thw))
             else:
                 if len(s.imgs) > 0:
-                    s_imgs = s.imgs.unsqueeze(0) if isinstance(s.imgs, torch.Tensor) and s.imgs.dim() == 3 else s.imgs
+                    s_imgs = self._normalize_visual_block(s.imgs) if isinstance(s.imgs, torch.Tensor) else s.imgs
                     if isinstance(s_imgs, torch.Tensor):
                         imgs.append(s_imgs)
                 if len(s.image_thw_grids) > 0:
@@ -567,7 +608,7 @@ def _collect_visual_data(self, samples):
                     else:
                         image_thw_grids.extend(s.image_thw_grids)
                 if len(s.videos) > 0:
-                    s_vids = s.videos.unsqueeze(0) if isinstance(s.videos, torch.Tensor) and s.videos.dim() == 3 else s.videos
+                    s_vids = self._normalize_visual_block(s.videos) if isinstance(s.videos, torch.Tensor) else s.videos
                     if isinstance(s_vids, torch.Tensor):
                         videos.append(s_vids)
                 if len(s.video_thw_grids) > 0:
@@ -656,8 +697,8 @@ def batch(self, samples: List[Union[QwenVLTaskSample, QwenVLTaskSamplePacked]])
                 __restore_key__=(),
                 __keys__=sample_keys,
                 __subflavors__=[s.__subflavors__ for s in samples],
-                pixel_values=torch.vstack(imgs) if len(imgs) > 0 else None,
-                pixel_values_videos=torch.vstack(videos) if len(videos) > 0 else None,
+                pixel_values=self._concat_visual_blocks(imgs, "image"),
+                pixel_values_videos=self._concat_visual_blocks(videos, "video"),
                 image_grid_thw=torch.from_numpy(np.array(image_thw_grids)) if len(image_thw_grids) > 0 else None,
                 video_grid_thw=torch.from_numpy(np.array(video_thw_grids)) if len(video_thw_grids) > 0 else None,
                 image_input_mask=torch.from_numpy(image_input_masks),
@@ -716,8 +757,8 @@ def batch(self, samples: List[Union[QwenVLTaskSample, QwenVLTaskSamplePacked]])
             __restore_key__=(),
             __keys__=sample_keys,
             __subflavors__=[s.__subflavors__ for s in samples],
-            pixel_values=torch.vstack(imgs) if len(imgs) > 0 else None,
-            pixel_values_videos=torch.vstack(videos) if len(videos) > 0 else None,
+            pixel_values=self._concat_visual_blocks(imgs, "image"),
+            pixel_values_videos=self._concat_visual_blocks(videos, "video"),
             image_grid_thw=torch.from_numpy(np.array(image_thw_grids)) if len(image_thw_grids) > 0 else None,
             video_grid_thw=torch.from_numpy(np.array(video_thw_grids)) if len(video_thw_grids) > 0 else None,
             image_input_mask=torch.from_numpy(image_input_masks),

From 0f51a2cbdc71e5d22aa0886cd5efd93097ff0a9c Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Mon, 13 Apr 2026 15:36:05 +0800
Subject: [PATCH 20/51] fix: remove temporary THD debug logging noise

Drop hardcoded first-iteration debug dumps from vlm_step and Qwen3VL MRoPE path so training logs stay clean while preserving core packed-sequence behavior.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../qwen_vl/modelling_qwen3_vl/model.py       | 18 -------
 src/megatron/bridge/training/vlm_step.py      | 48 -------------------
 2 files changed, 66 deletions(-)

diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
index 40a61cbe71..ecae8ef875 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
@@ -218,7 +218,6 @@ def __init__(
             assert self.config.calculate_per_token_loss, (
                 "Qwen3-VL model only supports context parallelism with calculate_per_token_loss enabled"
             )
-        self._mrope_debug_count = 0
 
         self._expose_language_model_for_cuda_graph_helper()
 
@@ -649,23 +648,6 @@ def forward(
                     sl_int = int(sl)
                     packed_pos[:, 0, start : start + sl_int] = position_ids[:, i, :sl_int]
 
-                # Debug only the first two forwards on rank-0 to verify packed MRoPE reset behavior.
-                if self._mrope_debug_count < 2:
-                    rank = (
-                        torch.distributed.get_rank()
-                        if torch.distributed.is_available() and torch.distributed.is_initialized()
-                        else 0
-                    )
-                    if rank == 0:
-                        boundary_pos0 = [int(cu[i].item()) for i in range(len(seq_lens))]
-                        boundary_posvals = [int(packed_pos[0, 0, p].item()) for p in boundary_pos0]
-                        logger.info(
-                            "[MRoPE-DEBUG] packed branch hit: cu=%s seq_lens=%s boundary_pos0=%s",
-                            cu.tolist(),
-                            [int(x) for x in seq_lens],
-                            boundary_posvals,
-                        )
-                    self._mrope_debug_count += 1
                 position_ids = packed_pos
 
                 position_ids = (
diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index 906969da45..5e89c17031 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -476,57 +476,9 @@ def forward_step(
         ) = get_batch(data_iterator, state.cfg, use_mtp, pg_collection=pg_collection)
     timers("batch-generator").stop()
 
-    # ---- Unified debug dump for first 2 iterations (both BSHD and Energon) ----
-    _iter = getattr(state, "_debug_iter_count", 0)
-    if _iter < 2 and tokens is not None:
-        _is_energon = cu_seqlens is not None
-        _mode = "Energon-THD" if _is_energon else "HF-BSHD"
-        _t = tokens.squeeze()
-        _lm = loss_mask.squeeze() if loss_mask is not None else None
-        _lab = labels.squeeze() if labels is not None else None
-        _sup_count = int((_lab != -100).sum().item()) if _lab is not None else 0
-        _lm_sum = int(_lm.sum().item()) if _lm is not None else 0
-        logger.info(
-            f"\n{'='*70}\n"
-            f"[DEBUG {_mode}] iteration {_iter}\n"
-            f"  tokens.shape        = {list(tokens.shape)}\n"
-            f"  tokens[:20]         = {_t[:20].tolist()}\n"
-            f"  tokens[-10:]        = {_t[-10:].tolist()}\n"
-            f"  labels[:20]         = {_lab[:20].tolist() if _lab is not None else None}\n"
-            f"  loss_mask sum/total = {_lm_sum} / {_t.shape[0]}\n"
-            f"  supervised labels   = {_sup_count}\n"
-            f"  cu_seqlens (raw)    = {cu_seqlens}\n"
-            f"  position_ids[:10]   = {position_ids.squeeze()[:10].tolist() if position_ids is not None else None}\n"
-            f"  visual_inputs       = {'present' if visual_inputs is not None else 'None'}\n"
-            f"{'='*70}"
-        )
-        # Label alignment check: labels[i] should equal tokens[i+1] at supervised positions
-        if _lab is not None:
-            _sup_idx = (_lab != -100).nonzero(as_tuple=True)[0]
-            if len(_sup_idx) > 0:
-                _valid = _sup_idx[_sup_idx < len(_t) - 1]
-                _match = int((_lab[_valid] == _t[_valid + 1]).sum().item())
-                logger.info(
-                    f"  label alignment: labels[i]==tokens[i+1] at {_match}/{len(_valid)} supervised positions "
-                    f"(first_sup_pos={_sup_idx[0].item()}, last_sup_pos={_sup_idx[-1].item()})"
-                )
-                if _match < len(_valid):
-                    _mismatch_idx = _valid[_lab[_valid] != _t[_valid + 1]][:5]
-                    for _mi in _mismatch_idx:
-                        _mi = _mi.item()
-                        logger.info(
-                            f"    MISMATCH pos {_mi}: label={_lab[_mi].item()} vs tokens[{_mi+1}]={_t[_mi+1].item()}"
-                        )
-        state._debug_iter_count = _iter + 1
-
     # For Qwen-VL, let model compute MRoPE position_ids from multimodal tokens.
     # Passing 1D sequential position_ids from dataloader would bypass that path.
     model_position_ids = None if visual_inputs is not None else position_ids
-    if _iter < 2:
-        logger.info(
-            "  forward position_ids mode = %s",
-            "None (compute MRoPE in model)" if model_position_ids is None else "provided-from-dataloader",
-        )
 
     forward_args = {
         "input_ids": tokens,

From c57be633ab40ea43205c5c09eecf3ddbea8b8b00 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Mon, 13 Apr 2026 16:52:10 +0800
Subject: [PATCH 21/51] fix: gate THD diagnostics and remove hardcoded
 IterStats token id

Route THD diagnostic logs behind THD_DIAG and compute image token counts from the active model config so IterStats reflects real tokenizer ids.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../qwen_vl/data/energon/task_encoder.py      | 35 +++++++++++--------
 src/megatron/bridge/training/vlm_step.py      | 27 ++++++++++++--
 2 files changed, 45 insertions(+), 17 deletions(-)

diff --git a/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py b/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
index 882de547b6..8538648d34 100644
--- a/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
+++ b/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
@@ -15,6 +15,7 @@
 import bisect
 import dataclasses
 import logging
+import os
 import re
 from collections import defaultdict
 from dataclasses import dataclass, field
@@ -43,6 +44,10 @@
 logger = logging.getLogger(__name__)
 
 
+def _thd_diag_enabled() -> bool:
+    return os.environ.get("THD_DIAG", "0") not in ("0", "", "false", "False")
+
+
 def _search_for_fit(numbers: List[int], capacity: int) -> int:
     """Binary search for the largest number that fits within capacity."""
     index = bisect.bisect(numbers, capacity)
@@ -511,13 +516,14 @@ def pack_selected_samples(self, samples: List[QwenVLTaskSample]) -> QwenVLTaskSa
         total_packed = int(cu_lengths[-1].item())
         max_sub_len = max(sub_lengths)
 
-        logger.info(
-            f"[PackingStats] packed {len(samples)} samples → {total_packed} tokens "
-            f"(seq_len={self.seq_len}, utilization={total_packed / self.seq_len * 100:.1f}%), "
-            f"sub_lengths={sub_lengths}, "
-            f"image_tokens={total_image_tokens}, video_tokens={total_video_tokens}, "
-            f"text_tokens={total_text_tokens}, vit_patches={total_vit_patches}"
-        )
+        if _thd_diag_enabled():
+            logger.info(
+                f"[PackingStats] packed {len(samples)} samples → {total_packed} tokens "
+                f"(seq_len={self.seq_len}, utilization={total_packed / self.seq_len * 100:.1f}%), "
+                f"sub_lengths={sub_lengths}, "
+                f"image_tokens={total_image_tokens}, video_tokens={total_video_tokens}, "
+                f"text_tokens={total_text_tokens}, vit_patches={total_vit_patches}"
+            )
 
         return QwenVLTaskSamplePacked(
             __key__="+".join(s.__key__ for s in samples),
@@ -659,13 +665,14 @@ def batch(self, samples: List[Union[QwenVLTaskSample, QwenVLTaskSamplePacked]])
                 max_lengths_list.append(s.max_length)
 
                 pad_len = fixed_len - actual_len
-                logger.info(
-                    f"[BatchStats] sample {i}: {s.num_sub_samples} sub-seqs packed, "
-                    f"content={actual_len}, pad={pad_len}, "
-                    f"image_toks={s.num_image_tokens}, video_toks={s.num_video_tokens}, "
-                    f"text_toks={s.num_text_tokens}, vit_patches={s.num_vit_patches}, "
-                    f"utilization={actual_len / fixed_len * 100:.1f}%"
-                )
+                if _thd_diag_enabled():
+                    logger.info(
+                        f"[BatchStats] sample {i}: {s.num_sub_samples} sub-seqs packed, "
+                        f"content={actual_len}, pad={pad_len}, "
+                        f"image_toks={s.num_image_tokens}, video_toks={s.num_video_tokens}, "
+                        f"text_toks={s.num_text_tokens}, vit_patches={s.num_vit_patches}, "
+                        f"utilization={actual_len / fixed_len * 100:.1f}%"
+                    )
 
             tokens = torch.from_numpy(text_mat)
             tokens[tokens == pad_token_id] = 0
diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index 5e89c17031..b56069478c 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import logging
+import os
 import math
 from functools import partial
 from typing import Any, Iterable
@@ -39,6 +40,21 @@
 logger = logging.getLogger(__name__)
 
 
+def _thd_diag_enabled() -> bool:
+    return os.environ.get("THD_DIAG", "0") not in ("0", "", "false", "False")
+
+
+def _resolve_cfg_token_id(cfg: ConfigContainer, key: str) -> int | None:
+    model_cfg = getattr(cfg, "model", None)
+    token_id = getattr(model_cfg, key, None)
+    if token_id is None:
+        return None
+    try:
+        return int(token_id)
+    except (TypeError, ValueError):
+        return None
+
+
 def get_batch_from_iterator(
     data_iterator: Iterable,
     use_mtp: bool = False,
@@ -346,8 +362,8 @@ def _ceil_to_mult(n: int, mult: int) -> int:
         energon_max_seqlen = batch.get("max_seqlen")
         energon_cu_argmin = batch.get("cu_seqlens_argmin")
 
-        # Log detailed iteration stats for FLOPs analysis
-        if tokens_or_input is not None:
+        # Log detailed packed-iteration diagnostics when explicitly enabled.
+        if tokens_or_input is not None and _thd_diag_enabled():
             _seq_dim = tokens_or_input.shape[-1]
             if energon_cu_seqlens.dim() == 1:
                 _content_len = energon_cu_seqlens[-1].item()
@@ -357,7 +373,11 @@ def _ceil_to_mult(n: int, mult: int) -> int:
             else:
                 _content_len = energon_cu_seqlens[0, -1].item()
             _pad_len = _seq_dim - _content_len
-            _n_img_toks = int((tokens_or_input == 151655).sum().item())
+            image_token_id = _resolve_cfg_token_id(cfg, "image_token_id")
+            if image_token_id is not None:
+                _n_img_toks: int | str = int((tokens_or_input == image_token_id).sum().item())
+            else:
+                _n_img_toks = "n/a"
             _vit_shape = (
                 visual_inputs.pixel_values.shape
                 if visual_inputs is not None and visual_inputs.pixel_values is not None
@@ -367,6 +387,7 @@ def _ceil_to_mult(n: int, mult: int) -> int:
                 f"[IterStats] decoder_input={list(tokens_or_input.shape)}, "
                 f"content={_content_len}, pad={_pad_len}, "
                 f"image_tokens_in_decoder={_n_img_toks}, "
+                f"image_token_id={image_token_id}, "
                 f"vit_input={_vit_shape}"
             )
 

From 8e07070a5a786257a9bb5ef7f267a4a24875d8fd Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Mon, 13 Apr 2026 17:20:37 +0800
Subject: [PATCH 22/51] chore: add THD low-content dataloader diagnostics

Add switch-gated low-content packed-sample diagnostics and gate remaining THD debug logs so all related debug output is controlled by THD_DIAG.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../qwen_vl/data/energon/task_encoder.py      | 42 ++++++++++++++++++-
 src/megatron/bridge/training/vlm_step.py      | 10 +++--
 2 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py b/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
index 8538648d34..9c4cb3b211 100644
--- a/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
+++ b/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
@@ -42,12 +42,29 @@
 
 
 logger = logging.getLogger(__name__)
+_LOW_CONTENT_LOG_COUNT = 0
 
 
 def _thd_diag_enabled() -> bool:
     return os.environ.get("THD_DIAG", "0") not in ("0", "", "false", "False")
 
 
+def _thd_diag_low_content_threshold() -> int:
+    raw = os.environ.get("THD_DIAG_LOW_CONTENT_TOKENS", "2048")
+    try:
+        return max(0, int(raw))
+    except ValueError:
+        return 2048
+
+
+def _thd_diag_max_low_content_logs() -> int:
+    raw = os.environ.get("THD_DIAG_MAX_LOW_CONTENT_LOGS", "100")
+    try:
+        return max(0, int(raw))
+    except ValueError:
+        return 100
+
+
 def _search_for_fit(numbers: List[int], capacity: int) -> int:
     """Binary search for the largest number that fits within capacity."""
     index = bisect.bisect(numbers, capacity)
@@ -642,6 +659,9 @@ def batch(self, samples: List[Union[QwenVLTaskSample, QwenVLTaskSamplePacked]])
         if is_packed:
             batch_size = len(samples)
             fixed_len = self.seq_len
+            low_content_threshold = _thd_diag_low_content_threshold()
+            max_low_content_logs = _thd_diag_max_low_content_logs()
+            diag_enabled = _thd_diag_enabled()
 
             text_mat = np.full((batch_size, fixed_len), pad_token_id, dtype=np.int64)
             target_mat = np.full((batch_size, fixed_len), pad_token_id, dtype=np.int64)
@@ -665,7 +685,7 @@ def batch(self, samples: List[Union[QwenVLTaskSample, QwenVLTaskSamplePacked]])
                 max_lengths_list.append(s.max_length)
 
                 pad_len = fixed_len - actual_len
-                if _thd_diag_enabled():
+                if diag_enabled:
                     logger.info(
                         f"[BatchStats] sample {i}: {s.num_sub_samples} sub-seqs packed, "
                         f"content={actual_len}, pad={pad_len}, "
@@ -673,6 +693,26 @@ def batch(self, samples: List[Union[QwenVLTaskSample, QwenVLTaskSamplePacked]])
                         f"text_toks={s.num_text_tokens}, vit_patches={s.num_vit_patches}, "
                         f"utilization={actual_len / fixed_len * 100:.1f}%"
                     )
+                    global _LOW_CONTENT_LOG_COUNT
+                    if (
+                        low_content_threshold > 0
+                        and actual_len < low_content_threshold
+                        and _LOW_CONTENT_LOG_COUNT < max_low_content_logs
+                    ):
+                        logger.warning(
+                            "[LowContentPackedSample] key=%s content=%d pad=%d seq_len=%d utilization=%.1f%% "
+                            "sub_samples=%d sub_lengths=%s image_toks=%d text_toks=%d",
+                            s.__key__,
+                            actual_len,
+                            pad_len,
+                            fixed_len,
+                            actual_len / fixed_len * 100.0,
+                            s.num_sub_samples,
+                            s.sub_sample_lengths,
+                            s.num_image_tokens,
+                            s.num_text_tokens,
+                        )
+                        _LOW_CONTENT_LOG_COUNT += 1
 
             tokens = torch.from_numpy(text_mat)
             tokens[tokens == pad_token_id] = 0
diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index b56069478c..8e9aa1a679 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -253,9 +253,10 @@ def pack_batch_sequences(
             )
         offset += padded_len
 
-    logger.debug(
-        f"Packed {len(valid_sequences)} sequences: lengths={seq_lengths}, total_len={total_len}, max_len={max_seqlen}"
-    )
+    if _thd_diag_enabled():
+        logger.debug(
+            f"Packed {len(valid_sequences)} sequences: lengths={seq_lengths}, total_len={total_len}, max_len={max_seqlen}"
+        )
 
     # Attention mask is not used with packed sequences (handled by cu_seqlens)
     packed_attention_mask = None
@@ -443,7 +444,8 @@ def _ceil_to_mult(n: int, mult: int) -> int:
         batch["attention_mask"] = packed_attention_mask
         batch["position_ids"] = packed_position_ids
 
-        logger.debug(f"In-batch packed: cu_seqlens={cu_seqlens.tolist()}, max_seqlen={max_seqlen}")
+        if _thd_diag_enabled():
+            logger.debug(f"In-batch packed: cu_seqlens={cu_seqlens.tolist()}, max_seqlen={max_seqlen}")
     else:
         cu_seqlens = None
         max_seqlen = None

From be70e4f9fcce5151c96e6c7846477236114a4f15 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 14 Apr 2026 09:10:35 +0800
Subject: [PATCH 23/51] fix: wire THD runtime switches into vlm step paths

Connect THD_FORCE_BSHD, THD_FORCE_SINGLE_SEGMENT_CU, and THD_SKIP_PREPROCESS_PACKED_POS to actual vlm_step control flow, and log switch activation only when enabled.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 src/megatron/bridge/training/vlm_step.py | 53 ++++++++++++++++++++++--
 1 file changed, 50 insertions(+), 3 deletions(-)

diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index 8e9aa1a679..84f0e4ef01 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -38,10 +38,31 @@
 
 
 logger = logging.getLogger(__name__)
+_SWITCH_LOGGED: set[str] = set()
 
 
 def _thd_diag_enabled() -> bool:
-    return os.environ.get("THD_DIAG", "0") not in ("0", "", "false", "False")
+    return _switch_enabled("THD_DIAG")
+
+
+def _switch_enabled(name: str) -> bool:
+    enabled = os.environ.get(name, "0") not in ("0", "", "false", "False")
+    if enabled and name not in _SWITCH_LOGGED:
+        logger.info("[THD_SWITCH] %s=1 enabled", name)
+        _SWITCH_LOGGED.add(name)
+    return enabled
+
+
+def _thd_force_bshd_enabled() -> bool:
+    return _switch_enabled("THD_FORCE_BSHD")
+
+
+def _thd_force_single_segment_cu_enabled() -> bool:
+    return _switch_enabled("THD_FORCE_SINGLE_SEGMENT_CU")
+
+
+def _thd_skip_preprocess_packed_pos_enabled() -> bool:
+    return _switch_enabled("THD_SKIP_PREPROCESS_PACKED_POS")
 
 
 def _resolve_cfg_token_id(cfg: ConfigContainer, key: str) -> int | None:
@@ -297,6 +318,13 @@ def get_batch(data_iterator: Iterable, cfg: ConfigContainer, use_mtp: bool = Fal
         is_last_pp_stage=is_last,
     )
     enable_packing = getattr(cfg.dataset, "pack_sequences_in_batch", False)
+    force_bshd = _thd_force_bshd_enabled()
+    force_single_segment_cu = _thd_force_single_segment_cu_enabled()
+
+    if force_bshd and force_single_segment_cu:
+        logger.warning(
+            "[THD_SWITCH] THD_FORCE_BSHD=1 overrides THD_FORCE_SINGLE_SEGMENT_CU=1 (single-segment CU ignored)."
+        )
 
     if not enable_packing:
         # When using pipeline parallelism, ensure fixed shapes equal to cfg.model.seq_length
@@ -356,13 +384,26 @@ def _ceil_to_mult(n: int, mult: int) -> int:
     tp_size = pg_collection.tp.size() if pg_collection is not None and pg_collection.tp is not None else 1
     has_sp = getattr(cfg.model, "sequence_parallel", False)
 
+    if force_bshd:
+        # Explicitly disable packed paths and ignore pre-packed cu_seqlens metadata.
+        enable_packing = False
+
     # Energon pre-packed path: cu_seqlens already present from TaskEncoder packing
     energon_cu_seqlens = batch.get("cu_seqlens")
-    if energon_cu_seqlens is not None:
+    if energon_cu_seqlens is not None and not force_bshd:
         tokens_or_input = batch.get("tokens") if batch.get("tokens") is not None else batch.get("input_ids")
         energon_max_seqlen = batch.get("max_seqlen")
         energon_cu_argmin = batch.get("cu_seqlens_argmin")
 
+        if force_single_segment_cu and tokens_or_input is not None:
+            seq_len = int(tokens_or_input.shape[-1])
+            device = energon_cu_seqlens.device if isinstance(energon_cu_seqlens, torch.Tensor) else tokens_or_input.device
+            energon_cu_seqlens = torch.tensor([0, seq_len], dtype=torch.int32, device=device)
+            energon_max_seqlen = torch.tensor(seq_len, dtype=torch.int32, device=device)
+            energon_cu_argmin = torch.tensor(2, dtype=torch.int64)
+            if _thd_diag_enabled():
+                logger.info("[THD_SWITCH] force single-segment cu_seqlens=[0, %d]", seq_len)
+
         # Log detailed packed-iteration diagnostics when explicitly enabled.
         if tokens_or_input is not None and _thd_diag_enabled():
             _seq_dim = tokens_or_input.shape[-1]
@@ -530,7 +571,7 @@ def forward_step(
         cu_clean = cu[:n_valid]
         last_boundary = int(cu_clean[-1].item())
 
-        if last_boundary < physical_seq_len:
+        if last_boundary < physical_seq_len and not _thd_skip_preprocess_packed_pos_enabled():
             # Trailing padding exists: extend the last real segment's padded
             # boundary to cover the tail so that all layers (TE attention, GDN,
             # etc.) see the same sequence boundaries.  We intentionally do NOT
@@ -552,6 +593,12 @@ def forward_step(
             }
         else:
             # No trailing padding (content fills the full sequence length).
+            if last_boundary < physical_seq_len and _thd_diag_enabled():
+                logger.info(
+                    "[THD_SWITCH] skip packed-pos preprocess: preserve original cu_seqlens last_boundary=%d physical_seq_len=%d",
+                    last_boundary,
+                    physical_seq_len,
+                )
             packed_seq_dict = {
                 "cu_seqlens": cu_seqlens,
                 "max_seqlen": max_seqlen,

From 24b05038275facc61b23cb8f90e4e655cdfa6763 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 14 Apr 2026 09:34:19 +0800
Subject: [PATCH 24/51] fix: split packed boundaries for rope and kernels

Keep unpadded cu_seqlens for MRoPE sub-sequence semantics while passing padded boundaries for TE/GDN kernels, and limit THD switch logs to rank0.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 src/megatron/bridge/training/vlm_step.py | 60 ++++++++++++------------
 1 file changed, 30 insertions(+), 30 deletions(-)

diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index 84f0e4ef01..7a60e956fe 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -41,13 +41,19 @@
 _SWITCH_LOGGED: set[str] = set()
 
 
+def _rank0() -> bool:
+    if not torch.distributed.is_available() or not torch.distributed.is_initialized():
+        return True
+    return torch.distributed.get_rank() == 0
+
+
 def _thd_diag_enabled() -> bool:
     return _switch_enabled("THD_DIAG")
 
 
 def _switch_enabled(name: str) -> bool:
     enabled = os.environ.get(name, "0") not in ("0", "", "false", "False")
-    if enabled and name not in _SWITCH_LOGGED:
+    if enabled and _rank0() and name not in _SWITCH_LOGGED:
         logger.info("[THD_SWITCH] %s=1 enabled", name)
         _SWITCH_LOGGED.add(name)
     return enabled
@@ -321,7 +327,7 @@ def get_batch(data_iterator: Iterable, cfg: ConfigContainer, use_mtp: bool = Fal
     force_bshd = _thd_force_bshd_enabled()
     force_single_segment_cu = _thd_force_single_segment_cu_enabled()
 
-    if force_bshd and force_single_segment_cu:
+    if force_bshd and force_single_segment_cu and _rank0():
         logger.warning(
             "[THD_SWITCH] THD_FORCE_BSHD=1 overrides THD_FORCE_SINGLE_SEGMENT_CU=1 (single-segment CU ignored)."
         )
@@ -401,7 +407,7 @@ def _ceil_to_mult(n: int, mult: int) -> int:
             energon_cu_seqlens = torch.tensor([0, seq_len], dtype=torch.int32, device=device)
             energon_max_seqlen = torch.tensor(seq_len, dtype=torch.int32, device=device)
             energon_cu_argmin = torch.tensor(2, dtype=torch.int64)
-            if _thd_diag_enabled():
+            if _thd_diag_enabled() and _rank0():
                 logger.info("[THD_SWITCH] force single-segment cu_seqlens=[0, %d]", seq_len)
 
         # Log detailed packed-iteration diagnostics when explicitly enabled.
@@ -571,42 +577,36 @@ def forward_step(
         cu_clean = cu[:n_valid]
         last_boundary = int(cu_clean[-1].item())
 
+        # Keep unpadded boundaries for MRoPE/sub-sequence semantics, while
+        # providing padded boundaries for TE/GDN kernels when needed.
+        cu_unpadded = cu_clean
+
         if last_boundary < physical_seq_len and not _thd_skip_preprocess_packed_pos_enabled():
-            # Trailing padding exists: extend the last real segment's padded
-            # boundary to cover the tail so that all layers (TE attention, GDN,
-            # etc.) see the same sequence boundaries.  We intentionally do NOT
-            # set cu_seqlens_unpadded here so that cu_seqlens_q_padded stays
-            # None — this forces every layer to treat pad tokens identically to
-            # real tokens, avoiding undefined attention output at pad positions
-            # that would otherwise propagate NaN through subsequent layers and
-            # into the backward pass.  loss_mask already zeros pad positions.
+            # Trailing padding exists: extend padded boundary to full physical
+            # sequence length so packed kernels cover all tokens.
             cu_padded = cu_clean.clone()
             cu_padded[-1] = physical_seq_len
-
-            padded_diffs = cu_padded[1:] - cu_padded[:-1]
-            max_seqlen_padded = padded_diffs.max()
-
-            packed_seq_dict = {
-                "cu_seqlens": cu_padded,
-                "cu_seqlens_argmin": torch.tensor(n_valid),
-                "max_seqlen": max_seqlen_padded,
-            }
+            max_seqlen_out = (cu_padded[1:] - cu_padded[:-1]).max()
         else:
-            # No trailing padding (content fills the full sequence length).
-            if last_boundary < physical_seq_len and _thd_diag_enabled():
+            if last_boundary < physical_seq_len and _thd_diag_enabled() and _rank0():
                 logger.info(
                     "[THD_SWITCH] skip packed-pos preprocess: preserve original cu_seqlens last_boundary=%d physical_seq_len=%d",
                     last_boundary,
                     physical_seq_len,
                 )
-            packed_seq_dict = {
-                "cu_seqlens": cu_seqlens,
-                "max_seqlen": max_seqlen,
-            }
-            if cu_seqlens_argmin is not None:
-                packed_seq_dict["cu_seqlens_argmin"] = cu_seqlens_argmin
-            elif cu_seqlens.dim() == 1:
-                packed_seq_dict["cu_seqlens_argmin"] = torch.tensor(len(cu_seqlens))
+            cu_padded = cu_clean
+            if max_seqlen is not None:
+                max_seqlen_out = max_seqlen.squeeze()
+            else:
+                max_seqlen_out = (cu_padded[1:] - cu_padded[:-1]).max()
+
+        packed_seq_dict = {
+            "cu_seqlens": cu_padded,
+            "cu_seqlens_argmin": torch.tensor(len(cu_padded)),
+            "cu_seqlens_unpadded": cu_unpadded,
+            "cu_seqlens_unpadded_argmin": torch.tensor(len(cu_unpadded)),
+            "max_seqlen": max_seqlen_out,
+        }
 
         forward_args["packed_seq_params"] = get_packed_seq_params(packed_seq_dict)
 

From 583f1850e005f2409e34b0f49135138a3c22f649 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 14 Apr 2026 09:55:02 +0800
Subject: [PATCH 25/51] fix: keep THD kernels padded while using unpadded rope
 boundaries

Avoid switching to the unstable unpadded TE packed path by keeping PackedSeqParams padded-only, and pass unpadded cu_seqlens separately for Qwen MRoPE sub-sequence indexing.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../bridge/models/qwen_vl/modelling_qwen3_vl/model.py    | 3 ++-
 src/megatron/bridge/training/vlm_step.py                 | 9 +++++----
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
index ecae8ef875..049408d35d 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
@@ -364,6 +364,7 @@ def forward(
         video_input_mask: torch.Tensor = None,
         cp_img_num: list[int] = None,
         images_padded: list[bool] = None,
+        rope_cu_seqlens: torch.Tensor = None,
         inference_context: object | None = None,
         runtime_gather_output: bool | None = None,
         mm_token_type_ids: torch.Tensor = None,
@@ -608,7 +609,7 @@ def forward(
             if packed_seq_params is not None:
                 # Packed sequences: compute MRoPE per sub-sequence independently
                 # so that each sub-sequence's positions restart from 0.
-                cu = packed_seq_params.cu_seqlens_q
+                cu = rope_cu_seqlens if rope_cu_seqlens is not None else packed_seq_params.cu_seqlens_q
                 if cu.dim() > 1:
                     cu = cu.squeeze()
                 seq_lens = (cu[1:] - cu[:-1]).tolist()
diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index 7a60e956fe..3ad8cd77b0 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -577,8 +577,9 @@ def forward_step(
         cu_clean = cu[:n_valid]
         last_boundary = int(cu_clean[-1].item())
 
-        # Keep unpadded boundaries for MRoPE/sub-sequence semantics, while
-        # providing padded boundaries for TE/GDN kernels when needed.
+        # Keep an unpadded copy for Qwen MRoPE sub-sequence semantics.
+        # Do not feed unpadded cu_seqlens into PackedSeqParams, which can
+        # switch TE kernels and destabilize training.
         cu_unpadded = cu_clean
 
         if last_boundary < physical_seq_len and not _thd_skip_preprocess_packed_pos_enabled():
@@ -603,12 +604,12 @@ def forward_step(
         packed_seq_dict = {
             "cu_seqlens": cu_padded,
             "cu_seqlens_argmin": torch.tensor(len(cu_padded)),
-            "cu_seqlens_unpadded": cu_unpadded,
-            "cu_seqlens_unpadded_argmin": torch.tensor(len(cu_unpadded)),
             "max_seqlen": max_seqlen_out,
         }
 
         forward_args["packed_seq_params"] = get_packed_seq_params(packed_seq_dict)
+        # Pass unpadded boundaries only to Qwen model's MRoPE construction.
+        forward_args["rope_cu_seqlens"] = cu_unpadded
 
     if loss_mask is not None:
         loss_mask = loss_mask.contiguous()

From 5813386921b2bc4669581d347444bd85cc7f58f0 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 14 Apr 2026 10:32:24 +0800
Subject: [PATCH 26/51] chore: add THD mask alignment diagnostics

Add rank0 THD_DIAG logs to correlate packed boundaries, MoE padding mask counts,
and loss-mask valid tokens without changing training behavior.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../qwen_vl/modelling_qwen3_vl/model.py       | 14 ++++++--
 src/megatron/bridge/training/vlm_step.py      | 33 ++++++++++++++++++-
 2 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
index 049408d35d..bf386d5d07 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
@@ -689,13 +689,21 @@ def forward(
         if packed_seq_params is not None and lm_input_ids is not None:
             padding_mask_for_moe = lm_input_ids.eq(0)
             if _thd_diag_enabled() and _rank0():
+                input_zero_cnt = int(input_ids.eq(0).sum().item()) if input_ids is not None else -1
+                lm_zero_cnt = int(lm_input_ids.eq(0).sum().item())
                 pad_cnt = int(padding_mask_for_moe.sum().item())
                 tok_cnt = int(padding_mask_for_moe.numel())
+                loss_valid = int((loss_mask > 0).sum().item()) if loss_mask is not None else -1
+                loss_total = int(loss_mask.numel()) if loss_mask is not None else -1
                 logger.info(
-                    "[THD_DIAG][model] padding_mask_for_moe: total_tokens=%d padding_tokens=%d valid_tokens=%d",
+                    "[THD_DIAG][model] padding_mask_for_moe: total_tokens=%d padding_tokens=%d valid_tokens=%d input_zero_tokens=%d lm_zero_tokens=%d loss_valid_tokens=%d loss_total_tokens=%d",
                     tok_cnt,
                     pad_cnt,
                     tok_cnt - pad_cnt,
+                    input_zero_cnt,
+                    lm_zero_cnt,
+                    loss_valid,
+                    loss_total,
                 )
 
         output = self.language_model(
@@ -706,8 +714,8 @@ def forward(
             labels=labels,  # only not None in the last decoder PP stage
             loss_mask=loss_mask,  # Added for THD training compatibility
             padding_mask=padding_mask_for_moe,  # for MoE routing/aux-loss token accounting
-            inference_params=inference_params,  # currently always None
-            packed_seq_params=packed_seq_params,  # currently always None
+            inference_params=inference_params,  # training path keeps this as None
+            packed_seq_params=packed_seq_params,
             visual_pos_masks=visual_pos_masks,
             deepstack_visual_embeds=deepstack_visual_embeds,
             **(extra_block_kwargs or {}),
diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index 3ad8cd77b0..a145d766af 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -239,7 +239,7 @@ def pack_batch_sequences(
         else:
             padded_len = length
         padded_seq_lengths.append(padded_len)
-        # Use padded lengths for cu_seqlens so THD RoPE splits sum correctly under CP.
+        # Use padded lengths for packed kernels under CP/SP alignment constraints.
         cu_seqlens.append(cu_seqlens[-1] + padded_len)
 
     cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device)
@@ -610,6 +610,15 @@ def forward_step(
         forward_args["packed_seq_params"] = get_packed_seq_params(packed_seq_dict)
         # Pass unpadded boundaries only to Qwen model's MRoPE construction.
         forward_args["rope_cu_seqlens"] = cu_unpadded
+        if _thd_diag_enabled() and _rank0():
+            logger.info(
+                "[THD_DIAG][packed] physical_seq_len=%d cu_unpadded_last=%d cu_padded_last=%d implicit_pad=%d max_seqlen=%d",
+                int(physical_seq_len),
+                int(cu_unpadded[-1].item()),
+                int(cu_padded[-1].item()),
+                int(physical_seq_len - int(cu_unpadded[-1].item())),
+                int(max_seqlen_out.item()) if torch.is_tensor(max_seqlen_out) else int(max_seqlen_out),
+            )
 
     if loss_mask is not None:
         loss_mask = loss_mask.contiguous()
@@ -634,6 +643,28 @@ def forward_step(
             else:
                 output_tensor = model_output
 
+    if _thd_diag_enabled() and _rank0() and loss_mask is not None:
+        loss_mask_flat = loss_mask.view(-1)
+        valid_loss_tokens = int((loss_mask_flat > 0).sum().item())
+        total_loss_tokens = int(loss_mask_flat.numel())
+        if labels is not None:
+            labels_flat = labels.view(-1)
+            ignore_tokens = int((labels_flat == -100).sum().item())
+            logger.info(
+                "[THD_DIAG][loss] total=%d valid=%d masked=%d labels_ignore=%d",
+                total_loss_tokens,
+                valid_loss_tokens,
+                total_loss_tokens - valid_loss_tokens,
+                ignore_tokens,
+            )
+        else:
+            logger.info(
+                "[THD_DIAG][loss] total=%d valid=%d masked=%d labels_ignore=n/a",
+                total_loss_tokens,
+                valid_loss_tokens,
+                total_loss_tokens - valid_loss_tokens,
+            )
+
     loss_function = _create_loss_function(loss_mask, check_for_nan_in_loss, check_for_spiky_loss)
 
     return output_tensor, loss_function

From bacf81317c1f6dbb92f0268c71238c700eb6bc90 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 14 Apr 2026 11:19:39 +0800
Subject: [PATCH 27/51] fix: use explicit THD MoE padding mask

Build MoE padding mask from packed boundaries and pass it explicitly through Qwen3VL
instead of inferring padding from lm_input_ids.eq(0), with fallback kept for compatibility.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../qwen_vl/modelling_qwen3_vl/model.py       | 27 +++++++++++++++++--
 src/megatron/bridge/training/vlm_step.py      | 13 ++++++++-
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
index bf386d5d07..4dc1b39f2f 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
@@ -365,6 +365,7 @@ def forward(
         cp_img_num: list[int] = None,
         images_padded: list[bool] = None,
         rope_cu_seqlens: torch.Tensor = None,
+        moe_padding_mask: torch.Tensor = None,
         inference_context: object | None = None,
         runtime_gather_output: bool | None = None,
         mm_token_type_ids: torch.Tensor = None,
@@ -414,6 +415,7 @@ def forward(
         # so it must be a real tensor. For packed sequences we use the THD-format
         # input_ids_thd (updated below); for regular sequences we use input_ids as-is.
         lm_input_ids = input_ids
+        moe_padding_mask_for_lm = moe_padding_mask
 
         if self.pre_process:
             # can reorganize_inputs at dataset
@@ -525,6 +527,13 @@ def forward(
                     input_ids, attention_mask, pre_process=True, pg_collection=self.pg_collection
                 )
                 lm_input_ids = input_ids_thd
+                if moe_padding_mask_for_lm is not None:
+                    moe_padding_mask_for_lm = preprocess_packed_seqs(
+                        moe_padding_mask_for_lm.to(dtype=torch.int32),
+                        attention_mask,
+                        pre_process=True,
+                        pg_collection=self.pg_collection,
+                    )[0].bool()
                 _, _, vision_mask_thd = reorganize_inputs(
                     input_ids=input_ids_thd,
                     pixel_values=pixel_values,
@@ -580,6 +589,13 @@ def forward(
                 lm_input_ids, _ = preprocess_packed_seqs(
                     input_ids, attention_mask, pre_process=True, pg_collection=self.pg_collection
                 )
+                if moe_padding_mask_for_lm is not None:
+                    moe_padding_mask_for_lm = preprocess_packed_seqs(
+                        moe_padding_mask_for_lm.to(dtype=torch.int32),
+                        attention_mask,
+                        pre_process=True,
+                        pg_collection=self.pg_collection,
+                    )[0].bool()
 
         visual_pos_masks = vision_mask
         deepstack_visual_embeds = deepstack_feature_lists
@@ -687,7 +703,13 @@ def forward(
         # a dedicated padding_mask (True=padding) to GPTModel/Router.
         padding_mask_for_moe = None
         if packed_seq_params is not None and lm_input_ids is not None:
-            padding_mask_for_moe = lm_input_ids.eq(0)
+            if moe_padding_mask_for_lm is not None:
+                padding_mask_for_moe = moe_padding_mask_for_lm.bool()
+                mask_source = "explicit"
+            else:
+                # Fallback for old call sites that do not provide explicit packed padding mask.
+                padding_mask_for_moe = lm_input_ids.eq(0)
+                mask_source = "token_eq_0_fallback"
             if _thd_diag_enabled() and _rank0():
                 input_zero_cnt = int(input_ids.eq(0).sum().item()) if input_ids is not None else -1
                 lm_zero_cnt = int(lm_input_ids.eq(0).sum().item())
@@ -696,7 +718,8 @@ def forward(
                 loss_valid = int((loss_mask > 0).sum().item()) if loss_mask is not None else -1
                 loss_total = int(loss_mask.numel()) if loss_mask is not None else -1
                 logger.info(
-                    "[THD_DIAG][model] padding_mask_for_moe: total_tokens=%d padding_tokens=%d valid_tokens=%d input_zero_tokens=%d lm_zero_tokens=%d loss_valid_tokens=%d loss_total_tokens=%d",
+                    "[THD_DIAG][model] padding_mask_for_moe: source=%s total_tokens=%d padding_tokens=%d valid_tokens=%d input_zero_tokens=%d lm_zero_tokens=%d loss_valid_tokens=%d loss_total_tokens=%d",
+                    mask_source,
                     tok_cnt,
                     pad_cnt,
                     tok_cnt - pad_cnt,
diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index a145d766af..18407ea132 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -606,18 +606,29 @@ def forward_step(
             "cu_seqlens_argmin": torch.tensor(len(cu_padded)),
             "max_seqlen": max_seqlen_out,
         }
+        # Build an explicit padding mask for MoE routing/aux-loss accounting.
+        # True means padding token that should be excluded from MoE statistics.
+        moe_padding_mask = torch.zeros_like(tokens, dtype=torch.bool)
+        if cu_unpadded.numel() == cu_padded.numel() and cu_padded.numel() > 1:
+            for i in range(int(cu_padded.numel()) - 1):
+                seg_unpadded_end = int(cu_unpadded[i + 1].item())
+                seg_padded_end = int(cu_padded[i + 1].item())
+                if seg_padded_end > seg_unpadded_end:
+                    moe_padding_mask[..., seg_unpadded_end:seg_padded_end] = True
 
         forward_args["packed_seq_params"] = get_packed_seq_params(packed_seq_dict)
         # Pass unpadded boundaries only to Qwen model's MRoPE construction.
         forward_args["rope_cu_seqlens"] = cu_unpadded
+        forward_args["moe_padding_mask"] = moe_padding_mask
         if _thd_diag_enabled() and _rank0():
             logger.info(
-                "[THD_DIAG][packed] physical_seq_len=%d cu_unpadded_last=%d cu_padded_last=%d implicit_pad=%d max_seqlen=%d",
+                "[THD_DIAG][packed] physical_seq_len=%d cu_unpadded_last=%d cu_padded_last=%d implicit_pad=%d max_seqlen=%d moe_padding_tokens=%d",
                 int(physical_seq_len),
                 int(cu_unpadded[-1].item()),
                 int(cu_padded[-1].item()),
                 int(physical_seq_len - int(cu_unpadded[-1].item())),
                 int(max_seqlen_out.item()) if torch.is_tensor(max_seqlen_out) else int(max_seqlen_out),
+                int(moe_padding_mask.sum().item()),
             )
 
     if loss_mask is not None:

From 552d4a7870514793c15e2db268cad2c1e1bca377 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 14 Apr 2026 11:36:57 +0800
Subject: [PATCH 28/51] fix: cast attention mask to bool for THD mask remap

Ensure packed mask remapping uses boolean indexing semantics when preprocessing
moe_padding_mask, preventing explicit THD padding mask from collapsing to zero.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../bridge/models/qwen_vl/modelling_qwen3_vl/model.py       | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
index 4dc1b39f2f..a688c2bd87 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
@@ -528,9 +528,10 @@ def forward(
                 )
                 lm_input_ids = input_ids_thd
                 if moe_padding_mask_for_lm is not None:
+                    attn_mask_bool = attention_mask.bool()
                     moe_padding_mask_for_lm = preprocess_packed_seqs(
                         moe_padding_mask_for_lm.to(dtype=torch.int32),
-                        attention_mask,
+                        attn_mask_bool,
                         pre_process=True,
                         pg_collection=self.pg_collection,
                     )[0].bool()
@@ -590,9 +591,10 @@ def forward(
                     input_ids, attention_mask, pre_process=True, pg_collection=self.pg_collection
                 )
                 if moe_padding_mask_for_lm is not None:
+                    attn_mask_bool = attention_mask.bool()
                     moe_padding_mask_for_lm = preprocess_packed_seqs(
                         moe_padding_mask_for_lm.to(dtype=torch.int32),
-                        attention_mask,
+                        attn_mask_bool,
                         pre_process=True,
                         pg_collection=self.pg_collection,
                     )[0].bool()

From d801e7b42314a48cf7e538e4c6c81649ebdd4fb5 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 14 Apr 2026 13:19:05 +0800
Subject: [PATCH 29/51] chore: add THD label/loss alignment diagnostics

Add a default-off THD_DIAG_ALIGN switch to log pre/post packed remap summaries for labels and loss_mask, so we can directly verify whether THD preprocessing preserves supervision alignment.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../qwen_vl/modelling_qwen3_vl/model.py       | 67 ++++++++++++++++++-
 1 file changed, 65 insertions(+), 2 deletions(-)

diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
index a688c2bd87..17dc82637d 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
@@ -61,6 +61,10 @@ def _thd_diag_enabled() -> bool:
     return os.environ.get("THD_DIAG", "0") not in ("0", "", "false", "False")
 
 
+def _thd_diag_align_enabled() -> bool:
+    return os.environ.get("THD_DIAG_ALIGN", "0") not in ("0", "", "false", "False")
+
+
 def _rank0() -> bool:
     if not torch.distributed.is_available() or not torch.distributed.is_initialized():
         return True
@@ -417,6 +421,18 @@ def forward(
         lm_input_ids = input_ids
         moe_padding_mask_for_lm = moe_padding_mask
 
+        def _mask_summary(mask: torch.Tensor, max_items: int = 8) -> tuple[int, str]:
+            flat_idx = torch.nonzero(mask.reshape(-1), as_tuple=False).view(-1)
+            count = int(flat_idx.numel())
+            if count == 0:
+                return count, "[]"
+            head = flat_idx[:max_items].tolist()
+            tail = flat_idx[-max_items:].tolist() if count > max_items else []
+            summary = f"head={head}"
+            if tail:
+                summary += f", tail={tail}"
+            return count, summary
+
         if self.pre_process:
             # can reorganize_inputs at dataset
             vision_data, vision_grid_thw, vision_mask = reorganize_inputs(
@@ -523,18 +539,65 @@ def forward(
             if packed_seq_params is not None:
                 if attention_mask is None:
                     attention_mask = torch.ones_like(input_ids, dtype=torch.int32, device=input_ids.device)
+                attn_mask_bool = attention_mask.bool()
                 input_ids_thd, _ = preprocess_packed_seqs(
                     input_ids, attention_mask, pre_process=True, pg_collection=self.pg_collection
                 )
                 lm_input_ids = input_ids_thd
                 if moe_padding_mask_for_lm is not None:
-                    attn_mask_bool = attention_mask.bool()
                     moe_padding_mask_for_lm = preprocess_packed_seqs(
                         moe_padding_mask_for_lm.to(dtype=torch.int32),
                         attn_mask_bool,
                         pre_process=True,
                         pg_collection=self.pg_collection,
                     )[0].bool()
+                if (
+                    _thd_diag_align_enabled()
+                    and _rank0()
+                    and labels is not None
+                    and loss_mask is not None
+                ):
+                    labels_thd = preprocess_packed_seqs(
+                        labels,
+                        attn_mask_bool,
+                        pre_process=True,
+                        pg_collection=self.pg_collection,
+                    )[0]
+                    loss_mask_thd = preprocess_packed_seqs(
+                        loss_mask.to(dtype=torch.float32),
+                        attn_mask_bool,
+                        pre_process=True,
+                        pg_collection=self.pg_collection,
+                    )[0]
+                    labels_valid_pre = labels.ne(-100)
+                    labels_valid_post = labels_thd.ne(-100)
+                    loss_valid_pre = loss_mask > 0
+                    loss_valid_post = loss_mask_thd > 0
+                    label_pre_cnt, label_pre_idx = _mask_summary(labels_valid_pre)
+                    label_post_cnt, label_post_idx = _mask_summary(labels_valid_post)
+                    loss_pre_cnt, loss_pre_idx = _mask_summary(loss_valid_pre)
+                    loss_post_cnt, loss_post_idx = _mask_summary(loss_valid_post)
+                    logger.info(
+                        "[THD_DIAG][align] labels_pre_count=%d labels_post_count=%d labels_same=%s labels_pre_%s labels_post_%s",
+                        label_pre_cnt,
+                        label_post_cnt,
+                        str(bool(torch.equal(labels_valid_pre, labels_valid_post))),
+                        label_pre_idx,
+                        label_post_idx,
+                    )
+                    logger.info(
+                        "[THD_DIAG][align] loss_pre_count=%d loss_post_count=%d loss_same=%s loss_pre_%s loss_post_%s",
+                        loss_pre_cnt,
+                        loss_post_cnt,
+                        str(bool(torch.equal(loss_valid_pre, loss_valid_post))),
+                        loss_pre_idx,
+                        loss_post_idx,
+                    )
+                    logger.info(
+                        "[THD_DIAG][align] pre_label_loss_same=%s post_label_loss_same=%s",
+                        str(bool(torch.equal(labels_valid_pre, loss_valid_pre))),
+                        str(bool(torch.equal(labels_valid_post, loss_valid_post))),
+                    )
                 _, _, vision_mask_thd = reorganize_inputs(
                     input_ids=input_ids_thd,
                     pixel_values=pixel_values,
@@ -587,11 +650,11 @@ def forward(
             if packed_seq_params is not None:
                 if attention_mask is None:
                     attention_mask = torch.ones_like(input_ids, dtype=torch.int32, device=input_ids.device)
+                attn_mask_bool = attention_mask.bool()
                 lm_input_ids, _ = preprocess_packed_seqs(
                     input_ids, attention_mask, pre_process=True, pg_collection=self.pg_collection
                 )
                 if moe_padding_mask_for_lm is not None:
-                    attn_mask_bool = attention_mask.bool()
                     moe_padding_mask_for_lm = preprocess_packed_seqs(
                         moe_padding_mask_for_lm.to(dtype=torch.int32),
                         attn_mask_bool,

From bab9f910153925391e498680603813c1b6cc65a2 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 14 Apr 2026 13:54:53 +0800
Subject: [PATCH 30/51] chore: add THD MRoPE position diagnostics

Add opt-in THD MRoPE pre/post remap diagnostics behind THD_DIAG_MROPE to inspect position semantics without affecting default behavior.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../qwen_vl/modelling_qwen3_vl/model.py       | 46 +++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
index 17dc82637d..42e741a90d 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
@@ -65,6 +65,10 @@ def _thd_diag_align_enabled() -> bool:
     return os.environ.get("THD_DIAG_ALIGN", "0") not in ("0", "", "false", "False")
 
 
+def _thd_diag_mrope_enabled() -> bool:
+    return os.environ.get("THD_DIAG_MROPE", "0") not in ("0", "", "false", "False")
+
+
 def _rank0() -> bool:
     if not torch.distributed.is_available() or not torch.distributed.is_initialized():
         return True
@@ -730,6 +734,32 @@ def _mask_summary(mask: torch.Tensor, max_items: int = 8) -> tuple[int, str]:
                     sl_int = int(sl)
                     packed_pos[:, 0, start : start + sl_int] = position_ids[:, i, :sl_int]
 
+                if _thd_diag_mrope_enabled() and _rank0():
+                    # Summarize per-subsequence MRoPE position behavior before THD remap.
+                    # This is diagnostics-only and intentionally does not affect semantics.
+                    show_n = 4
+                    head_stats = []
+                    for i, sl in enumerate(seq_lens[:show_n]):
+                        start = int(cu[i].item())
+                        sl_int = int(sl)
+                        if sl_int <= 0:
+                            head_stats.append(f"{i}:len=0")
+                            continue
+                        pos0 = packed_pos[0, 0, start : start + sl_int]
+                        nonmono = int((pos0[1:] < pos0[:-1]).sum().item()) if sl_int > 1 else 0
+                        head_stats.append(
+                            f"{i}:len={sl_int},start={int(pos0[0].item())},end={int(pos0[-1].item())},"
+                            f"min={int(pos0.min().item())},max={int(pos0.max().item())},dec={nonmono}"
+                        )
+                    logger.info(
+                        "[THD_DIAG][mrope] pre_thd cp_size=%d num_seqs=%d total_len=%d seq_lens_head=%s seg_stats_head=%s",
+                        int(cp_size),
+                        int(num_seqs),
+                        int(total_len),
+                        seq_lens[:show_n],
+                        "; ".join(head_stats),
+                    )
+
                 position_ids = packed_pos
 
                 position_ids = (
@@ -742,6 +772,22 @@ def _mask_summary(mask: torch.Tensor, max_items: int = 8) -> tuple[int, str]:
                     .permute(2, 0, 1)
                     .contiguous()
                 )
+                if _thd_diag_mrope_enabled() and _rank0():
+                    pos0_post = position_ids[0, 0]
+                    post_len = int(pos0_post.numel())
+                    head_vals = pos0_post[:12].tolist()
+                    tail_vals = pos0_post[-12:].tolist() if post_len > 12 else []
+                    nonmono_post = int((pos0_post[1:] < pos0_post[:-1]).sum().item()) if post_len > 1 else 0
+                    logger.info(
+                        "[THD_DIAG][mrope] post_thd cp_size=%d len=%d min=%d max=%d dec=%d head=%s tail=%s",
+                        int(cp_size),
+                        post_len,
+                        int(pos0_post.min().item()) if post_len > 0 else -1,
+                        int(pos0_post.max().item()) if post_len > 0 else -1,
+                        nonmono_post,
+                        head_vals,
+                        tail_vals,
+                    )
                 attention_mask = None
                 self.language_model.rotary_pos_emb.is_thd_format = True
             else:

From 10f6d780a010602558217eefdcd3815145086473 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 14 Apr 2026 14:09:24 +0800
Subject: [PATCH 31/51] fix: enforce bool masks for THD packed remap

Use boolean attention masks at all THD preprocess_packed_seqs call sites and add a function-level bool cast guard to prevent silent positional remap corruption.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../bridge/models/qwen_vl/modelling_qwen3_vl/model.py | 11 ++++++-----
 .../bridge/models/qwen_vl/modelling_qwen3_vl/utils.py |  5 +++++
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
index 42e741a90d..dea92df20c 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
@@ -545,7 +545,7 @@ def _mask_summary(mask: torch.Tensor, max_items: int = 8) -> tuple[int, str]:
                     attention_mask = torch.ones_like(input_ids, dtype=torch.int32, device=input_ids.device)
                 attn_mask_bool = attention_mask.bool()
                 input_ids_thd, _ = preprocess_packed_seqs(
-                    input_ids, attention_mask, pre_process=True, pg_collection=self.pg_collection
+                    input_ids, attn_mask_bool, pre_process=True, pg_collection=self.pg_collection
                 )
                 lm_input_ids = input_ids_thd
                 if moe_padding_mask_for_lm is not None:
@@ -622,7 +622,7 @@ def _mask_summary(mask: torch.Tensor, max_items: int = 8) -> tuple[int, str]:
                         tmp_embeddings[vision_mask] = deepstack_visual_embed
                         tmp_embeddings_thd = preprocess_packed_seqs(
                             tmp_embeddings.contiguous(),
-                            attention_mask,
+                            attn_mask_bool,
                             pre_process=True,
                             pg_collection=self.pg_collection,
                         )[0]
@@ -634,7 +634,7 @@ def _mask_summary(mask: torch.Tensor, max_items: int = 8) -> tuple[int, str]:
                 combined_embeddings_thd = (
                     preprocess_packed_seqs(
                         combined_embeddings.transpose(0, 1).contiguous(),
-                        attention_mask,
+                        attn_mask_bool,
                         pre_process=True,
                         pg_collection=self.pg_collection,
                     )[0]
@@ -656,7 +656,7 @@ def _mask_summary(mask: torch.Tensor, max_items: int = 8) -> tuple[int, str]:
                     attention_mask = torch.ones_like(input_ids, dtype=torch.int32, device=input_ids.device)
                 attn_mask_bool = attention_mask.bool()
                 lm_input_ids, _ = preprocess_packed_seqs(
-                    input_ids, attention_mask, pre_process=True, pg_collection=self.pg_collection
+                    input_ids, attn_mask_bool, pre_process=True, pg_collection=self.pg_collection
                 )
                 if moe_padding_mask_for_lm is not None:
                     moe_padding_mask_for_lm = preprocess_packed_seqs(
@@ -762,10 +762,11 @@ def _mask_summary(mask: torch.Tensor, max_items: int = 8) -> tuple[int, str]:
 
                 position_ids = packed_pos
 
+                attn_mask_bool = attention_mask.bool()
                 position_ids = (
                     preprocess_packed_seqs(
                         position_ids.permute(1, 2, 0),
-                        attention_mask,
+                        attn_mask_bool,
                         pre_process=True,
                         pg_collection=self.pg_collection,
                     )[0]
diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py
index 5cbc2b1251..0ae6abe92c 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py
@@ -694,6 +694,11 @@ def preprocess_packed_seqs(
     See https://github.com/NVIDIA/TransformerEngine/issues/1368
     """
     batch_size = input_ids.shape[0]
+    # Boolean mask is required for advanced indexing semantics below:
+    # input_ids[i, attention_mask[i]]
+    # If mask is int (0/1), PyTorch treats it as positional index selection
+    # rather than keep-mask, which can silently corrupt THD remapping.
+    attention_mask = attention_mask.bool()
 
     seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
     if pg_collection is not None:

From 03711a9d814d7612b364fef683732c81bc4a832d Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 14 Apr 2026 14:25:26 +0800
Subject: [PATCH 32/51] chore: add strict THD MRoPE remap check

Add opt-in strict diagnostics to compare THD-remapped position ids against explicit bool-mask gather semantics for cp_size==1.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../qwen_vl/modelling_qwen3_vl/model.py       | 36 +++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
index dea92df20c..51dd7ca74a 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
@@ -789,6 +789,42 @@ def _mask_summary(mask: torch.Tensor, max_items: int = 8) -> tuple[int, str]:
                         head_vals,
                         tail_vals,
                     )
+                    # Strict check: verify THD-remapped position_ids equals explicit bool-mask gather
+                    # semantics for cp_size==1. This is diagnostics-only.
+                    if int(cp_size) == 1:
+                        tp_size = int(self.pg_collection.tp.size())
+                        align_size = tp_size
+                        packed_pos_bsd = packed_pos.permute(1, 2, 0).contiguous()  # [1, S, 3]
+                        valid_len = int(attn_mask_bool[0].sum().item())
+                        padded_len = valid_len + ((align_size - (valid_len % align_size)) % align_size)
+                        expected_bsd = torch.zeros(
+                            (1, padded_len, packed_pos_bsd.size(-1)),
+                            dtype=packed_pos_bsd.dtype,
+                            device=packed_pos_bsd.device,
+                        )
+                        expected_bsd[0, :valid_len] = packed_pos_bsd[0, attn_mask_bool[0]]
+                        expected_pid = expected_bsd.permute(2, 0, 1).contiguous()
+
+                        same = bool(torch.equal(position_ids, expected_pid))
+                        mismatch = (position_ids != expected_pid).any(dim=0).squeeze(0)
+                        mismatch_cnt = int(mismatch.sum().item())
+                        mismatch_idx = torch.nonzero(mismatch, as_tuple=False).flatten()
+                        head_idx = mismatch_idx[:8].tolist()
+                        tail_idx = mismatch_idx[-8:].tolist() if mismatch_idx.numel() > 8 else []
+                        logger.info(
+                            "[THD_DIAG][mrope] strict_match=%s mismatch_count=%d expected_len=%d actual_len=%d mismatch_head=%s mismatch_tail=%s",
+                            str(same),
+                            mismatch_cnt,
+                            int(expected_pid.size(-1)),
+                            int(position_ids.size(-1)),
+                            head_idx,
+                            tail_idx,
+                        )
+                    else:
+                        logger.info(
+                            "[THD_DIAG][mrope] strict_match_skipped cp_size=%d (currently only checks cp_size==1)",
+                            int(cp_size),
+                        )
                 attention_mask = None
                 self.language_model.rotary_pos_emb.is_thd_format = True
             else:

From c68367cd4c93f10797affd04c8639b6a800a7262 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 14 Apr 2026 15:24:13 +0800
Subject: [PATCH 33/51] chore: add THD boundary diagnostics switch

Add opt-in THD boundary diagnostics to compare kernel vs rope cu summaries and report boundary-window label/loss/moe mask consistency, while keeping all new logic disabled by default.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 src/megatron/bridge/training/vlm_step.py | 52 ++++++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index 18407ea132..471660dda3 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -51,6 +51,10 @@ def _thd_diag_enabled() -> bool:
     return _switch_enabled("THD_DIAG")
 
 
+def _thd_diag_boundary_enabled() -> bool:
+    return _switch_enabled("THD_DIAG_BOUNDARY")
+
+
 def _switch_enabled(name: str) -> bool:
     enabled = os.environ.get(name, "0") not in ("0", "", "false", "False")
     if enabled and _rank0() and name not in _SWITCH_LOGGED:
@@ -630,6 +634,54 @@ def forward_step(
                 int(max_seqlen_out.item()) if torch.is_tensor(max_seqlen_out) else int(max_seqlen_out),
                 int(moe_padding_mask.sum().item()),
             )
+        if _thd_diag_boundary_enabled() and _rank0():
+            cu_padded_cpu = cu_padded.detach().cpu()
+            cu_unpadded_cpu = cu_unpadded.detach().cpu()
+            seg_lens_padded = (cu_padded_cpu[1:] - cu_padded_cpu[:-1]).tolist()
+            seg_lens_unpadded = (cu_unpadded_cpu[1:] - cu_unpadded_cpu[:-1]).tolist()
+            logger.info(
+                "[THD_DIAG][cu] kernel_num_segs=%d rope_num_segs=%d kernel_cu_head=%s rope_cu_head=%s kernel_seg_lens_head=%s rope_seg_lens_head=%s",
+                max(0, int(cu_padded_cpu.numel()) - 1),
+                max(0, int(cu_unpadded_cpu.numel()) - 1),
+                cu_padded_cpu[:8].tolist(),
+                cu_unpadded_cpu[:8].tolist(),
+                seg_lens_padded[:8],
+                seg_lens_unpadded[:8],
+            )
+            # Boundary windows are centered on each segment join index.
+            # This helps detect edge-only mismatch that can be hidden by global stats.
+            if labels is not None and loss_mask is not None and cu_unpadded_cpu.numel() > 2:
+                win_k = 8
+                seq_len = int(tokens.shape[-1])
+                boundary_window = torch.zeros(seq_len, dtype=torch.bool, device=tokens.device)
+                boundary_positions = []
+                for i in range(1, int(cu_unpadded_cpu.numel()) - 1):
+                    join_idx = int(cu_unpadded_cpu[i].item())
+                    boundary_positions.append(join_idx)
+                    left = max(0, join_idx - win_k)
+                    right = min(seq_len, join_idx + win_k)
+                    if right > left:
+                        boundary_window[left:right] = True
+                labels_flat = labels.view(-1)
+                loss_mask_flat = loss_mask.view(-1)
+                moe_mask_flat = moe_padding_mask.view(-1)
+                bw = boundary_window.view(-1)
+                bw_tokens = int(bw.sum().item())
+                bw_valid_loss = int(((loss_mask_flat > 0) & bw).sum().item())
+                bw_valid_labels = int(((labels_flat != -100) & bw).sum().item())
+                bw_moe_padding = int((moe_mask_flat & bw).sum().item())
+                bw_label_loss_mismatch = int((((labels_flat != -100) != (loss_mask_flat > 0)) & bw).sum().item())
+                logger.info(
+                    "[THD_DIAG][boundary] win_k=%d joins=%d join_head=%s window_tokens=%d valid_loss=%d valid_labels=%d moe_padding=%d label_loss_mismatch=%d",
+                    win_k,
+                    len(boundary_positions),
+                    boundary_positions[:8],
+                    bw_tokens,
+                    bw_valid_loss,
+                    bw_valid_labels,
+                    bw_moe_padding,
+                    bw_label_loss_mismatch,
+                )
 
     if loss_mask is not None:
         loss_mask = loss_mask.contiguous()

From 85a54c260a2724ebb1718a5fc3787a37ccbb7028 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 14 Apr 2026 17:58:04 +0800
Subject: [PATCH 34/51] chore: clarify packed-seq boundary documentation

Document how argmin-based trimming and padded boundary fields are used so THD packed-kernel boundary expectations are explicit.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../bridge/training/utils/packed_seq_utils.py     | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/src/megatron/bridge/training/utils/packed_seq_utils.py b/src/megatron/bridge/training/utils/packed_seq_utils.py
index 78d5dd17cd..37956bb135 100644
--- a/src/megatron/bridge/training/utils/packed_seq_utils.py
+++ b/src/megatron/bridge/training/utils/packed_seq_utils.py
@@ -21,9 +21,8 @@
 def get_packed_seq_params(batch: dict[str, torch.Tensor]) -> PackedSeqParams:
     """Build packed sequence parameters from a batch dictionary.
 
-    The function squeezes possible batch dimensions and removes any padding
-    marked by -1 values. It returns a `PackedSeqParams` instance suitable for
-    packed sequence attention kernels.
+    The function trims optional sentinel tails via argmin metadata and returns
+    a `PackedSeqParams` instance suitable for packed sequence kernels.
 
     Args:
         batch: A dictionary containing packed-sequence metadata. Expected keys:
@@ -31,8 +30,9 @@ def get_packed_seq_params(batch: dict[str, torch.Tensor]) -> PackedSeqParams:
             optional `max_seqlen`.
 
     Returns:
-        PackedSeqParams with identical q/kv parameters and `qkv_format` set to
-        "thd".
+        PackedSeqParams for THD (`qkv_format="thd"`). If unpadded boundaries
+        are provided, they are used for q/kv while padded boundaries remain in
+        `*_padded` fields for kernel compatibility.
     """
 
     cu_seqlens_padded = batch["cu_seqlens"].squeeze()
@@ -72,9 +72,8 @@ def get_packed_seq_params(batch: dict[str, torch.Tensor]) -> PackedSeqParams:
         )
     else:
         # Follow Megatron-LM data_schedule.py convention: set all four
-        # cu_seqlens fields to padded values so every downstream op (RoPE,
-        # TE attention, GDN) sees consistent boundaries without relying on
-        # fallback logic.
+        # cu_seqlens fields to padded values so packed kernels consume a
+        # single consistent boundary set.
         return PackedSeqParams(
             cu_seqlens_q=cu_seqlens_padded,
             cu_seqlens_kv=cu_seqlens_padded,

From 03ca9c7b2e33ae7f940198138cf1945aeb6030f8 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 14 Apr 2026 18:06:51 +0800
Subject: [PATCH 35/51] feat: expose energon shuffle buffer override

Allow EnergonProvider to accept and forward shuffle_buffer_size so training scripts can explicitly control sample shuffling behavior in THD/BSHD comparison runs.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 src/megatron/bridge/data/energon/energon_provider.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/megatron/bridge/data/energon/energon_provider.py b/src/megatron/bridge/data/energon/energon_provider.py
index 07de6a9328..58f9762f5d 100644
--- a/src/megatron/bridge/data/energon/energon_provider.py
+++ b/src/megatron/bridge/data/energon/energon_provider.py
@@ -36,6 +36,7 @@ class EnergonProvider(DatasetProvider):
     # Enable batch-level online sequence packing
     pack_sequences_in_batch: bool = False
     packing_buffer_size: Optional[int] = None
+    shuffle_buffer_size: int = 100
 
     def build_datasets(self, context: DatasetBuildContext):
         assert self.path, "EnergonProvider.path must be set. Use CLI override: dataset.path=<path>"
@@ -52,6 +53,7 @@ def build_datasets(self, context: DatasetBuildContext):
             global_batch_size=self.global_batch_size,
             num_workers=self.num_workers,
             packing_buffer_size=self.packing_buffer_size,
+            shuffle_buffer_size=self.shuffle_buffer_size,
             pg_collection=context.pg_collection,
         )
         return (

From a0b5fd52fec4944196d390fa5b74231e4887fd86 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 14 Apr 2026 19:23:32 +0800
Subject: [PATCH 36/51] chore: sync nodebug training updates without submodule
 bump

Apply nodebug recipe/provider updates while intentionally keeping the Megatron-LM submodule pointer unchanged.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../bridge/data/energon/energon_provider.py   | 17 ++++++++++--
 .../bridge/recipes/qwen_vl/__init__.py        |  2 ++
 .../bridge/recipes/qwen_vl/qwen35_vl.py       | 27 +++++++++++++++++++
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/src/megatron/bridge/data/energon/energon_provider.py b/src/megatron/bridge/data/energon/energon_provider.py
index 58f9762f5d..96f6723130 100644
--- a/src/megatron/bridge/data/energon/energon_provider.py
+++ b/src/megatron/bridge/data/energon/energon_provider.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 from dataclasses import dataclass
 from typing import Any, Optional
 
@@ -37,14 +38,26 @@ class EnergonProvider(DatasetProvider):
     pack_sequences_in_batch: bool = False
     packing_buffer_size: Optional[int] = None
     shuffle_buffer_size: int = 100
+    # Optional bin selector for datasets split into bin directories.
+    # Used to pin data selection and align comparisons with Energon BSHD.
+    cord_bins_root: Optional[str] = None
+    cord_bin_prefix: str = "cord_bin_"
+    cord_bin_id: Optional[str] = None
 
     def build_datasets(self, context: DatasetBuildContext):
-        assert self.path, "EnergonProvider.path must be set. Use CLI override: dataset.path=<path>"
+        resolved_path = self.path
+        if self.cord_bin_id is not None and self.cord_bin_id != "":
+            assert (
+                self.cord_bins_root
+            ), "EnergonProvider.cord_bins_root must be set when dataset.cord_bin_id is provided."
+            resolved_path = os.path.join(self.cord_bins_root, f"{self.cord_bin_prefix}{self.cord_bin_id}")
+
+        assert resolved_path, "EnergonProvider.path must be set. Use CLI override: dataset.path=<path>"
         if self.task_encoder is not None and hasattr(self.task_encoder, "seq_len"):
             self.task_encoder.seq_len = self.seq_length
             self.task_encoder.seq_length = self.seq_length
         dataset = EnergonMultiModalDataModule(
-            path=self.path,
+            path=resolved_path,
             tokenizer=context.tokenizer if context.tokenizer is not None else self.tokenizer,
             image_processor=self.image_processor,
             seq_length=self.seq_length,
diff --git a/src/megatron/bridge/recipes/qwen_vl/__init__.py b/src/megatron/bridge/recipes/qwen_vl/__init__.py
index 0c45ee8dc2..93dc981f59 100644
--- a/src/megatron/bridge/recipes/qwen_vl/__init__.py
+++ b/src/megatron/bridge/recipes/qwen_vl/__init__.py
@@ -53,6 +53,7 @@
     qwen35_vl_35b_a3b_pretrain_mock_config,
     qwen35_vl_35b_a3b_sft_config,
     qwen35_vl_35b_a3b_sft_energon_config,
+    qwen35_vl_35b_a3b_sft_energon_nopack_config,
     qwen35_vl_122b_a10b_peft_config,
     qwen35_vl_122b_a10b_pretrain_mock_config,
     qwen35_vl_122b_a10b_sft_config,
@@ -79,6 +80,7 @@
     # Qwen3.5-VL SFT configs — MoE
     "qwen35_vl_35b_a3b_sft_config",
     "qwen35_vl_35b_a3b_sft_energon_config",
+    "qwen35_vl_35b_a3b_sft_energon_nopack_config",
     "qwen35_vl_35b_a3b_fsdp_sft_config",
     "qwen35_vl_122b_a10b_sft_config",
     "qwen35_vl_397b_a17b_sft_config",
diff --git a/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py b/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
index 44c9fbd57c..4e94f779e3 100644
--- a/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
+++ b/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
@@ -394,12 +394,39 @@ def qwen35_vl_35b_a3b_sft_energon_config(hf_path: str = "Qwen/Qwen3.5-35B-A3B")
     """Return a full SFT config for Qwen3.5-VL 35B-A3B (MoE) with Energon dataset.
 
     Same as qwen35_vl_35b_a3b_sft_config but uses EnergonProvider instead of HF dataset.
+    This recipe enables batch-level packing by default (THD path).
     Set the dataset path via CLI override: dataset.path=/path/to/energon/dataset
+    Optional fixed-bin overrides:
+      - dataset.cord_bins_root=/path/to/cord_bins
+      - dataset.cord_bin_prefix=cord_bin_
+      - dataset.cord_bin_id=0
     """
     cfg = qwen35_vl_35b_a3b_sft_config(hf_path=hf_path)
     cfg.dataset = _make_energon_dataset(
         hf_path, cfg.dataset.seq_length, cfg.train.micro_batch_size, cfg.train.global_batch_size
     )
+    # THD path: enable batch-level sequence packing in Energon.
+    cfg.dataset.pack_sequences_in_batch = True
+    # Default packing window; can still be overridden via CLI:
+    # dataset.packing_buffer_size=<N>
+    cfg.dataset.packing_buffer_size = 128
+    return cfg
+
+
+def qwen35_vl_35b_a3b_sft_energon_nopack_config(hf_path: str = "Qwen/Qwen3.5-35B-A3B") -> ConfigContainer:
+    """Return 35B-A3B Energon SFT config with packing disabled (BSHD baseline path).
+
+    This is useful as a non-THD reference when comparing with the THD recipe.
+    Set the dataset path via CLI override: dataset.path=/path/to/energon/dataset
+    Supports the same optional fixed-bin overrides as the THD Energon recipe.
+    """
+    cfg = qwen35_vl_35b_a3b_sft_config(hf_path=hf_path)
+    cfg.dataset = _make_energon_dataset(
+        hf_path, cfg.dataset.seq_length, cfg.train.micro_batch_size, cfg.train.global_batch_size
+    )
+    # Baseline path: keep Energon data flow but disable packing.
+    cfg.dataset.pack_sequences_in_batch = False
+    cfg.dataset.packing_buffer_size = None
     return cfg
 
 

From 0525c55b02e97b591ddbddab06d9a874d91aae21 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 14 Apr 2026 19:48:22 +0800
Subject: [PATCH 37/51] refactor: split batch-level and in-batch packing
 switches

Introduce independent batch-level packing control for THD dataloader, keep in-batch packing semantics isolated, and update validation/step logic so the two modes cannot be mixed.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../bridge/data/energon/energon_provider.py   |  8 +++++--
 .../bridge/models/qwen_vl/qwen3_vl_step.py    |  6 +++--
 .../bridge/recipes/qwen_vl/qwen35_vl.py       |  7 ++++--
 src/megatron/bridge/training/config.py        | 22 ++++++++++++++-----
 src/megatron/bridge/training/vlm_step.py      |  5 +++--
 5 files changed, 35 insertions(+), 13 deletions(-)

diff --git a/src/megatron/bridge/data/energon/energon_provider.py b/src/megatron/bridge/data/energon/energon_provider.py
index 96f6723130..ca550acbed 100644
--- a/src/megatron/bridge/data/energon/energon_provider.py
+++ b/src/megatron/bridge/data/energon/energon_provider.py
@@ -34,8 +34,11 @@ class EnergonProvider(DatasetProvider):
     num_workers: int_repr
     dataloader_type: str = "external"
     task_encoder: Optional[Any] = None
-    # Enable batch-level online sequence packing
+    # Legacy in-batch packing switch. Keep for compatibility with existing recipes.
     pack_sequences_in_batch: bool = False
+    # THD dataloader switch for batch-level online packing.
+    # This is independent from pack_sequences_in_batch.
+    batch_level_packing: bool = False
     packing_buffer_size: Optional[int] = None
     shuffle_buffer_size: int = 100
     # Optional bin selector for datasets split into bin directories.
@@ -56,6 +59,7 @@ def build_datasets(self, context: DatasetBuildContext):
         if self.task_encoder is not None and hasattr(self.task_encoder, "seq_len"):
             self.task_encoder.seq_len = self.seq_length
             self.task_encoder.seq_length = self.seq_length
+        effective_packing_buffer_size = self.packing_buffer_size if self.batch_level_packing else None
         dataset = EnergonMultiModalDataModule(
             path=resolved_path,
             tokenizer=context.tokenizer if context.tokenizer is not None else self.tokenizer,
@@ -65,7 +69,7 @@ def build_datasets(self, context: DatasetBuildContext):
             micro_batch_size=self.micro_batch_size,
             global_batch_size=self.global_batch_size,
             num_workers=self.num_workers,
-            packing_buffer_size=self.packing_buffer_size,
+            packing_buffer_size=effective_packing_buffer_size,
             shuffle_buffer_size=self.shuffle_buffer_size,
             pg_collection=context.pg_collection,
         )
diff --git a/src/megatron/bridge/models/qwen_vl/qwen3_vl_step.py b/src/megatron/bridge/models/qwen_vl/qwen3_vl_step.py
index c87202262b..b0107104f7 100644
--- a/src/megatron/bridge/models/qwen_vl/qwen3_vl_step.py
+++ b/src/megatron/bridge/models/qwen_vl/qwen3_vl_step.py
@@ -248,7 +248,9 @@ def forward_step(
 
     # To be compatible with qwen3vl, we move the sequence padding and packing to forward_step function.
     # Qwen3VL model need the original input and do cp and sp split in model.forward.
-    pack_sequences_in_batch = getattr(state.cfg.dataset, "pack_sequences_in_batch", False)
+    in_batch_pack_enabled = getattr(state.cfg.dataset, "pack_sequences_in_batch", False)
+    batch_level_pack_enabled = getattr(state.cfg.dataset, "batch_level_packing", False)
+    enable_sequence_packing = in_batch_pack_enabled or batch_level_pack_enabled
 
     tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params = pack_or_pad_batch_sequences(
         tokens,
@@ -275,7 +277,7 @@ def forward_step(
     forward_args["input_ids"] = original_tokens
     # calculate position_ids in model forward
     forward_args["position_ids"] = None
-    if pack_sequences_in_batch:
+    if enable_sequence_packing:
         if forward_args["labels"] is not None:
             # When using pp, labels could be None
             forward_args["labels"] = forward_args["labels"].reshape(1, -1)
diff --git a/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py b/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
index 4e94f779e3..f5893f0c7c 100644
--- a/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
+++ b/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
@@ -405,8 +405,10 @@ def qwen35_vl_35b_a3b_sft_energon_config(hf_path: str = "Qwen/Qwen3.5-35B-A3B")
     cfg.dataset = _make_energon_dataset(
         hf_path, cfg.dataset.seq_length, cfg.train.micro_batch_size, cfg.train.global_batch_size
     )
-    # THD path: enable batch-level sequence packing in Energon.
-    cfg.dataset.pack_sequences_in_batch = True
+    # THD path: use dedicated batch-level packing switch.
+    # Keep legacy pack_sequences_in_batch for in-batch semantics only.
+    cfg.dataset.pack_sequences_in_batch = False
+    cfg.dataset.batch_level_packing = True
     # Default packing window; can still be overridden via CLI:
     # dataset.packing_buffer_size=<N>
     cfg.dataset.packing_buffer_size = 128
@@ -426,6 +428,7 @@ def qwen35_vl_35b_a3b_sft_energon_nopack_config(hf_path: str = "Qwen/Qwen3.5-35B
     )
     # Baseline path: keep Energon data flow but disable packing.
     cfg.dataset.pack_sequences_in_batch = False
+    cfg.dataset.batch_level_packing = False
     cfg.dataset.packing_buffer_size = None
     return cfg
 
diff --git a/src/megatron/bridge/training/config.py b/src/megatron/bridge/training/config.py
index 32de2a619f..83ba6bace7 100644
--- a/src/megatron/bridge/training/config.py
+++ b/src/megatron/bridge/training/config.py
@@ -1051,9 +1051,20 @@ def validate(self) -> None:
         Ensures compatibility between different configuration settings.
         """
 
-        # Propagate in-batch packing flag to model config so TransformerConfig.finalize()
+        in_batch_pack_enabled = getattr(self.dataset, "pack_sequences_in_batch", False)
+        batch_level_pack_enabled = getattr(self.dataset, "batch_level_packing", False)
+
+        # Keep in-batch packing and batch-level packing as two independent modes.
+        # They serve different semantics and should not be enabled together.
+        if in_batch_pack_enabled and batch_level_pack_enabled:
+            raise ValueError(
+                "dataset.pack_sequences_in_batch and dataset.batch_level_packing "
+                "cannot both be true. Keep in-batch packing and batch-level packing separate."
+            )
+
+        # Propagate packed-batch flags to model config so TransformerConfig.finalize()
         # can enable variable_seq_lengths for pipeline parallelism.
-        if getattr(self.dataset, "pack_sequences_in_batch", False):
+        if in_batch_pack_enabled or batch_level_pack_enabled:
             self.model._pack_sequences_in_batch = True
 
         if hasattr(self.dataset, "finalize"):
@@ -1250,10 +1261,11 @@ def validate(self) -> None:
                 f"https://docs.nvidia.com/nemo-framework/user-guide/latest/sft_peft/packed_sequence.html"
             )
 
-        if getattr(self.dataset, "pack_sequences_in_batch", False) and self.train.micro_batch_size == 1:
+        if in_batch_pack_enabled and self.train.micro_batch_size == 1:
             raise ValueError(
-                "micro_batch_size should be greater than 1 when using pack_sequences_in_batch=True. "
-                "In-batch packing concatenates multiple sequences within a microbatch, so at least 2 sequences "
+                "micro_batch_size should be greater than 1 when using in-batch packing "
+                "(dataset.pack_sequences_in_batch=True). In-batch packing concatenates multiple sequences "
+                "within a microbatch, so at least 2 sequences "
                 "are required per micro-batch."
             )
 
diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index 471660dda3..109fd308b6 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -327,7 +327,9 @@ def get_batch(data_iterator: Iterable, cfg: ConfigContainer, use_mtp: bool = Fal
         is_first_pp_stage=is_first,
         is_last_pp_stage=is_last,
     )
-    enable_packing = getattr(cfg.dataset, "pack_sequences_in_batch", False)
+    in_batch_pack_enabled = getattr(cfg.dataset, "pack_sequences_in_batch", False)
+    batch_level_pack_enabled = getattr(cfg.dataset, "batch_level_packing", False)
+    enable_packing = in_batch_pack_enabled or batch_level_pack_enabled
     force_bshd = _thd_force_bshd_enabled()
     force_single_segment_cu = _thd_force_single_segment_cu_enabled()
 
@@ -335,7 +337,6 @@ def get_batch(data_iterator: Iterable, cfg: ConfigContainer, use_mtp: bool = Fal
         logger.warning(
             "[THD_SWITCH] THD_FORCE_BSHD=1 overrides THD_FORCE_SINGLE_SEGMENT_CU=1 (single-segment CU ignored)."
         )
-
     if not enable_packing:
         # When using pipeline parallelism, ensure fixed shapes equal to cfg.model.seq_length
         if getattr(cfg.model, "pipeline_model_parallel_size", 1) > 1:

From 0f5ee33a31ea840b91682ee9e4b784f818df7ace Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 14 Apr 2026 19:59:03 +0800
Subject: [PATCH 38/51] fix: disable packing defaults in energon recipe

Keep both in-batch and THD batch-level packing disabled by default, requiring explicit CLI overrides to enable either mode and avoid unintended coupling.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py b/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
index f5893f0c7c..638a268f18 100644
--- a/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
+++ b/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
@@ -405,10 +405,12 @@ def qwen35_vl_35b_a3b_sft_energon_config(hf_path: str = "Qwen/Qwen3.5-35B-A3B")
     cfg.dataset = _make_energon_dataset(
         hf_path, cfg.dataset.seq_length, cfg.train.micro_batch_size, cfg.train.global_batch_size
     )
-    # THD path: use dedicated batch-level packing switch.
-    # Keep legacy pack_sequences_in_batch for in-batch semantics only.
+    # Keep both packing modes disabled by default.
+    # - in-batch packing: dataset.pack_sequences_in_batch
+    # - THD batch-level packing: dataset.batch_level_packing
+    # Enable either mode explicitly via CLI override when needed.
     cfg.dataset.pack_sequences_in_batch = False
-    cfg.dataset.batch_level_packing = True
+    cfg.dataset.batch_level_packing = False
     # Default packing window; can still be overridden via CLI:
     # dataset.packing_buffer_size=<N>
     cfg.dataset.packing_buffer_size = 128

From 2a018ec508af36d95f741f54528057834b16eda4 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 14 Apr 2026 21:37:46 +0800
Subject: [PATCH 39/51] fix: address optional import handling and recipe docs
 feedback

Narrow optional-import exception scopes, add fallback warnings, split diffusion optional imports, and align the qwen35 energon docstring with actual default packing behavior.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../bridge/models/conversion/auto_bridge.py      |  7 ++++++-
 src/megatron/bridge/recipes/__init__.py          |  4 ++++
 src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py |  3 ++-
 src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py  | 16 +++++++++++++---
 4 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/src/megatron/bridge/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py
index 8dbcb921d1..45e14876c2 100644
--- a/src/megatron/bridge/models/conversion/auto_bridge.py
+++ b/src/megatron/bridge/models/conversion/auto_bridge.py
@@ -52,9 +52,14 @@
 
 try:
     from modelopt.torch.quantization.utils import is_quantized
-except Exception:
+except ImportError as exc:
     # modelopt/scipy is optional for recipe import and training startup.
     # Keep quantization-only export path available when dependency exists.
+    logger.warning(
+        "modelopt quantization utils unavailable; quantized-export detection disabled: %s",
+        exc,
+    )
+
     def is_quantized(_model: object) -> bool:
         return False
 
diff --git a/src/megatron/bridge/recipes/__init__.py b/src/megatron/bridge/recipes/__init__.py
index 4428a47191..3685045ec7 100644
--- a/src/megatron/bridge/recipes/__init__.py
+++ b/src/megatron/bridge/recipes/__init__.py
@@ -22,6 +22,10 @@
 # non-diffusion training paths (LLM/VLM) usable when those deps are absent.
 try:
     from megatron.bridge.diffusion.recipes.flux.flux import *
+except ModuleNotFoundError:
+    pass
+
+try:
     from megatron.bridge.diffusion.recipes.wan.wan import *
 except ModuleNotFoundError:
     pass
diff --git a/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py b/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
index 638a268f18..2642e38bfb 100644
--- a/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
+++ b/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
@@ -394,7 +394,8 @@ def qwen35_vl_35b_a3b_sft_energon_config(hf_path: str = "Qwen/Qwen3.5-35B-A3B")
     """Return a full SFT config for Qwen3.5-VL 35B-A3B (MoE) with Energon dataset.
 
     Same as qwen35_vl_35b_a3b_sft_config but uses EnergonProvider instead of HF dataset.
-    This recipe enables batch-level packing by default (THD path).
+    Packing is disabled by default. Enable THD batch-level packing via CLI:
+    dataset.batch_level_packing=true
     Set the dataset path via CLI override: dataset.path=/path/to/energon/dataset
     Optional fixed-bin overrides:
       - dataset.cord_bins_root=/path/to/cord_bins
diff --git a/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py b/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py
index 64b126974a..5f3e165d49 100644
--- a/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py
+++ b/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py
@@ -19,6 +19,7 @@
 
 from __future__ import annotations
 
+import logging
 import os
 from typing import Optional, Union
 
@@ -28,7 +29,11 @@
 
 try:
     from transformers import Qwen3VLProcessor
-except Exception:
+except (ImportError, ModuleNotFoundError) as exc:
+    logger.warning(
+        "Qwen3VLProcessor import unavailable; falling back to AutoProcessor: %s",
+        exc,
+    )
     Qwen3VLProcessor = None
 
 from megatron.bridge import AutoBridge
@@ -54,6 +59,8 @@
 from megatron.bridge.training.flex_dispatcher_backend import apply_flex_dispatcher_backend
 from megatron.bridge.training.mixed_precision import MixedPrecisionConfig
 
+logger = logging.getLogger(__name__)
+
 
 class Qwen3VLCommonKwargs(TypedDict, total=False):
     """Typed options accepted by Qwen3-VL pretrain recipe helper."""
@@ -282,8 +289,11 @@ def _make_energon_dataset(
     if Qwen3VLProcessor is not None:
         try:
             image_processor = Qwen3VLProcessor.from_pretrained(hf_path)
-        except Exception:
-            pass
+        except (AttributeError, OSError, RuntimeError, ValueError) as exc:
+            logger.warning(
+                "Qwen3VLProcessor.from_pretrained failed, falling back to AutoProcessor: %s",
+                exc,
+            )
     if image_processor is None:
         from transformers import AutoProcessor
 

From 23d6bc3e8d3c90ae9ea09eb564da3b8fa53a79d3 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 14 Apr 2026 21:48:23 +0800
Subject: [PATCH 40/51] fix: harden packed seqlens trimming and gate
 unsupported qwen3 batch-level packing

Use safe cu_seqlens tail trimming when argmin metadata is absent, and prevent qwen3_vl_step from enabling batch-level packing before full THD metadata wiring.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../bridge/models/qwen_vl/qwen3_vl_step.py    | 13 ++++++++-
 .../bridge/training/utils/packed_seq_utils.py | 28 ++++++++++++-------
 2 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/src/megatron/bridge/models/qwen_vl/qwen3_vl_step.py b/src/megatron/bridge/models/qwen_vl/qwen3_vl_step.py
index b0107104f7..1a044a06dd 100644
--- a/src/megatron/bridge/models/qwen_vl/qwen3_vl_step.py
+++ b/src/megatron/bridge/models/qwen_vl/qwen3_vl_step.py
@@ -36,6 +36,7 @@
 
 
 logger = logging.getLogger(__name__)
+_BATCH_LEVEL_PACKING_WARNED = False
 
 
 def get_batch_from_iterator(
@@ -250,7 +251,17 @@ def forward_step(
     # Qwen3VL model need the original input and do cp and sp split in model.forward.
     in_batch_pack_enabled = getattr(state.cfg.dataset, "pack_sequences_in_batch", False)
     batch_level_pack_enabled = getattr(state.cfg.dataset, "batch_level_packing", False)
-    enable_sequence_packing = in_batch_pack_enabled or batch_level_pack_enabled
+    # qwen3_vl_step does not yet wire full THD metadata (rope/moe padding-mask)
+    # for batch-level packing. Keep this path disabled to avoid incorrect token
+    # accounting until parity with training.vlm_step is implemented.
+    global _BATCH_LEVEL_PACKING_WARNED
+    if batch_level_pack_enabled and not _BATCH_LEVEL_PACKING_WARNED:
+        logger.warning(
+            "dataset.batch_level_packing is not fully supported in qwen3_vl_step yet; "
+            "falling back to in-batch packing behavior."
+        )
+        _BATCH_LEVEL_PACKING_WARNED = True
+    enable_sequence_packing = in_batch_pack_enabled
 
     tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params = pack_or_pad_batch_sequences(
         tokens,
diff --git a/src/megatron/bridge/training/utils/packed_seq_utils.py b/src/megatron/bridge/training/utils/packed_seq_utils.py
index 37956bb135..28b97bc714 100644
--- a/src/megatron/bridge/training/utils/packed_seq_utils.py
+++ b/src/megatron/bridge/training/utils/packed_seq_utils.py
@@ -18,6 +18,20 @@
 from megatron.core.packed_seq_params import PackedSeqParams
 
 
+def _trim_cu_seqlens(cu_seqlens: torch.Tensor, cu_seqlens_argmin: torch.Tensor | None) -> torch.Tensor:
+    """Trim padded cu_seqlens tail safely.
+
+    If argmin metadata is available from dataloader, trust it. Otherwise, detect
+    the first zero in the tail region after index 0 (the leading 0 is valid).
+    """
+    if cu_seqlens_argmin is not None:
+        return cu_seqlens[: int(cu_seqlens_argmin.item())]
+
+    padded_tail = torch.nonzero(cu_seqlens[1:] == 0, as_tuple=False)
+    first_tail = int(padded_tail[0].item() + 1) if padded_tail.numel() else cu_seqlens.numel()
+    return cu_seqlens[:first_tail]
+
+
 def get_packed_seq_params(batch: dict[str, torch.Tensor]) -> PackedSeqParams:
     """Build packed sequence parameters from a batch dictionary.
 
@@ -43,18 +57,12 @@ def get_packed_seq_params(batch: dict[str, torch.Tensor]) -> PackedSeqParams:
     cu_seqlens_argmin = batch.get("cu_seqlens_argmin")
     cu_seqlens_unpadded_argmin = batch.get("cu_seqlens_unpadded_argmin")
 
-    # note: if argmin is not pre-computed in the dataloader, torch.argmin here will incur a
-    # device-to-host synchronization, which can slow down training
-    if cu_seqlens_argmin is not None:
-        cu_seqlens_padded = cu_seqlens_padded[: cu_seqlens_argmin.item()]
-    else:
-        cu_seqlens_padded = cu_seqlens_padded[: torch.argmin(cu_seqlens_padded)]
+    # Note: if argmin metadata is absent, fallback tail detection can still incur
+    # device-to-host synchronization.
+    cu_seqlens_padded = _trim_cu_seqlens(cu_seqlens_padded, cu_seqlens_argmin)
 
     if cu_seqlens_unpadded is not None:
-        if cu_seqlens_unpadded_argmin is not None:
-            cu_seqlens_unpadded = cu_seqlens_unpadded[: cu_seqlens_unpadded_argmin.item()]
-        else:
-            cu_seqlens_unpadded = cu_seqlens_unpadded[: torch.argmin(cu_seqlens_unpadded)]
+        cu_seqlens_unpadded = _trim_cu_seqlens(cu_seqlens_unpadded, cu_seqlens_unpadded_argmin)
 
     max_seqlen = batch["max_seqlen"].squeeze() if "max_seqlen" in batch else None
 

From 970e2365b420283f41f9794e0911fa677734db0f Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Mon, 20 Apr 2026 09:12:45 +0800
Subject: [PATCH 41/51] fix: address reviewer feedback for qwen-vl packing path

Resolve the import-fallback logger initialization bug, clean up task encoder formatting, and add focused unit tests for greedy_knapsack edge cases called out in PR review.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../qwen_vl/data/energon/task_encoder.py      |  2 ++
 .../bridge/recipes/qwen_vl/qwen3_vl.py        |  4 ++--
 .../qwen_vl/data/energon/test_task_encoder.py | 19 +++++++++++++++++++
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py b/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
index 9c4cb3b211..a04afca933 100644
--- a/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
+++ b/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
@@ -102,6 +102,8 @@ def greedy_knapsack(item_sizes: List[int], samples: List, max_capacity: int) ->
             current_knapsack.append(sorted_samples.pop(idx))
         knapsacks.append(current_knapsack)
     return knapsacks
+
+
 def process_vision(
     processor, images, videos, fps=None, model_version: str = "qwen-vl", min_pixels=None, max_pixels=None
 ):
diff --git a/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py b/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py
index 5f3e165d49..03963a70f4 100644
--- a/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py
+++ b/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py
@@ -27,6 +27,8 @@
 from transformers import AutoTokenizer
 from typing_extensions import TypedDict, Unpack
 
+logger = logging.getLogger(__name__)
+
 try:
     from transformers import Qwen3VLProcessor
 except (ImportError, ModuleNotFoundError) as exc:
@@ -59,8 +61,6 @@
 from megatron.bridge.training.flex_dispatcher_backend import apply_flex_dispatcher_backend
 from megatron.bridge.training.mixed_precision import MixedPrecisionConfig
 
-logger = logging.getLogger(__name__)
-
 
 class Qwen3VLCommonKwargs(TypedDict, total=False):
     """Typed options accepted by Qwen3-VL pretrain recipe helper."""
diff --git a/tests/unit_tests/recipes/qwen_vl/data/energon/test_task_encoder.py b/tests/unit_tests/recipes/qwen_vl/data/energon/test_task_encoder.py
index bb4e83abbd..a3d7a9c5b6 100644
--- a/tests/unit_tests/recipes/qwen_vl/data/energon/test_task_encoder.py
+++ b/tests/unit_tests/recipes/qwen_vl/data/energon/test_task_encoder.py
@@ -31,6 +31,7 @@
     _resolve_hf_mm_token_ids,
     convert_to_qwenvl_content,
     find_pattern_indices,
+    greedy_knapsack,
     get_ltor_masks_and_position_ids,
     process_vision,
     videohandler,
@@ -87,6 +88,24 @@ def test_get_ltor_masks_and_position_ids(self):
         self.assertTrue(torch.all(loss_mask == 1.0))
 
 
+class TestGreedyKnapsack(unittest.TestCase):
+    def test_empty_input(self):
+        self.assertEqual(greedy_knapsack([], [], max_capacity=8), [])
+
+    def test_single_item(self):
+        bins = greedy_knapsack([4], ["a"], max_capacity=8)
+        self.assertEqual(bins, [["a"]])
+
+    def test_exact_fill(self):
+        bins = greedy_knapsack([5, 3], ["a", "b"], max_capacity=8)
+        self.assertEqual(len(bins), 1)
+        self.assertCountEqual(bins[0], ["a", "b"])
+
+    def test_item_exceeds_capacity(self):
+        with self.assertRaises(ValueError):
+            greedy_knapsack([9], ["a"], max_capacity=8)
+
+
 class TestResolveHfMmTokenIds(unittest.TestCase):
     def test_resolves_from_tokenizer_attributes(self):
         tokenizer = MagicMock()

From 0b71929bc30dad3abccc191e82a95c288269d935 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Mon, 20 Apr 2026 09:25:33 +0800
Subject: [PATCH 42/51] docs: clarify in-batch vs batch-level packing semantics

Rewrite EnergonProvider packing comments to avoid deprecation ambiguity and explicitly describe how in-batch packing differs from dataloader-side batch-level packing.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 src/megatron/bridge/data/energon/energon_provider.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/megatron/bridge/data/energon/energon_provider.py b/src/megatron/bridge/data/energon/energon_provider.py
index ca550acbed..7c17460cde 100644
--- a/src/megatron/bridge/data/energon/energon_provider.py
+++ b/src/megatron/bridge/data/energon/energon_provider.py
@@ -34,10 +34,14 @@ class EnergonProvider(DatasetProvider):
     num_workers: int_repr
     dataloader_type: str = "external"
     task_encoder: Optional[Any] = None
-    # Legacy in-batch packing switch. Keep for compatibility with existing recipes.
+    # Existing in-batch packing switch.
+    # Semantics: pack samples *within each already formed micro-batch*.
+    # This path keeps historical behavior used by existing recipes.
     pack_sequences_in_batch: bool = False
     # THD dataloader switch for batch-level online packing.
-    # This is independent from pack_sequences_in_batch.
+    # Semantics: pack samples *across a dataloader-side candidate buffer*
+    # before a micro-batch is formed.
+    # This mode is independent from pack_sequences_in_batch.
     batch_level_packing: bool = False
     packing_buffer_size: Optional[int] = None
     shuffle_buffer_size: int = 100

From fcd08547422677797c17af54432174fb094770e9 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Mon, 20 Apr 2026 09:37:53 +0800
Subject: [PATCH 43/51] refactor: move qwen3-vl diag helpers into existing
 utils

Address reviewer feedback by moving THD diagnostic/rank helpers into modelling_qwen3_vl/utils.py, reusing a shared is_rank_0 utility, and adding focused unit tests for switch and rank helper behavior.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../qwen_vl/modelling_qwen3_vl/model.py       | 36 ++++----------
 .../qwen_vl/modelling_qwen3_vl/utils.py       | 17 +++++++
 src/megatron/bridge/utils/common_utils.py     |  5 ++
 .../models/qwen_vl/test_diag_utils.py         | 48 +++++++++++++++++++
 tests/unit_tests/utils/test_common_utils.py   | 19 ++++++++
 5 files changed, 97 insertions(+), 28 deletions(-)
 create mode 100644 tests/unit_tests/models/qwen_vl/test_diag_utils.py

diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
index 51dd7ca74a..a985024bd1 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import logging
-import os
 from typing import Optional
 
 import torch
@@ -45,11 +44,15 @@
     get_dist_train_vision_dp_data,
     get_vision_cp_data,
     pack_dist_train_vision_module_output,
+    is_rank_0,
     preprocess_packed_seqs,
     qwen3vl_cp_split,
     reorganize_inputs,
     split_data_cp_rank,
     split_deepstack_embs,
+    thd_diag_align_enabled,
+    thd_diag_enabled,
+    thd_diag_mrope_enabled,
 )
 from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.vision_model import Qwen3VLVisionModel
 
@@ -57,24 +60,6 @@
 logger = logging.getLogger(__name__)
 
 
-def _thd_diag_enabled() -> bool:
-    return os.environ.get("THD_DIAG", "0") not in ("0", "", "false", "False")
-
-
-def _thd_diag_align_enabled() -> bool:
-    return os.environ.get("THD_DIAG_ALIGN", "0") not in ("0", "", "false", "False")
-
-
-def _thd_diag_mrope_enabled() -> bool:
-    return os.environ.get("THD_DIAG_MROPE", "0") not in ("0", "", "false", "False")
-
-
-def _rank0() -> bool:
-    if not torch.distributed.is_available() or not torch.distributed.is_initialized():
-        return True
-    return torch.distributed.get_rank() == 0
-
-
 class Qwen3VLModel(MegatronModule):
     """Qwen3VL multi-modal model.
 
@@ -555,12 +540,7 @@ def _mask_summary(mask: torch.Tensor, max_items: int = 8) -> tuple[int, str]:
                         pre_process=True,
                         pg_collection=self.pg_collection,
                     )[0].bool()
-                if (
-                    _thd_diag_align_enabled()
-                    and _rank0()
-                    and labels is not None
-                    and loss_mask is not None
-                ):
+                if thd_diag_align_enabled() and is_rank_0() and labels is not None and loss_mask is not None:
                     labels_thd = preprocess_packed_seqs(
                         labels,
                         attn_mask_bool,
@@ -734,7 +714,7 @@ def _mask_summary(mask: torch.Tensor, max_items: int = 8) -> tuple[int, str]:
                     sl_int = int(sl)
                     packed_pos[:, 0, start : start + sl_int] = position_ids[:, i, :sl_int]
 
-                if _thd_diag_mrope_enabled() and _rank0():
+                if thd_diag_mrope_enabled() and is_rank_0():
                     # Summarize per-subsequence MRoPE position behavior before THD remap.
                     # This is diagnostics-only and intentionally does not affect semantics.
                     show_n = 4
@@ -773,7 +753,7 @@ def _mask_summary(mask: torch.Tensor, max_items: int = 8) -> tuple[int, str]:
                     .permute(2, 0, 1)
                     .contiguous()
                 )
-                if _thd_diag_mrope_enabled() and _rank0():
+                if thd_diag_mrope_enabled() and is_rank_0():
                     pos0_post = position_ids[0, 0]
                     post_len = int(pos0_post.numel())
                     head_vals = pos0_post[:12].tolist()
@@ -858,7 +838,7 @@ def _mask_summary(mask: torch.Tensor, max_items: int = 8) -> tuple[int, str]:
                 # Fallback for old call sites that do not provide explicit packed padding mask.
                 padding_mask_for_moe = lm_input_ids.eq(0)
                 mask_source = "token_eq_0_fallback"
-            if _thd_diag_enabled() and _rank0():
+            if thd_diag_enabled() and is_rank_0():
                 input_zero_cnt = int(input_ids.eq(0).sum().item()) if input_ids is not None else -1
                 lm_zero_cnt = int(lm_input_ids.eq(0).sum().item())
                 pad_cnt = int(padding_mask_for_moe.sum().item())
diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py
index 0ae6abe92c..7366d865b7 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 
+import os
 from dataclasses import dataclass
 from typing import Optional, Sequence, Union
 
@@ -26,6 +27,22 @@
 from torch import nn
 
 from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.transformer_config import Qwen3VLTransformerConfig
+from megatron.bridge.utils.common_utils import is_rank_0
+
+
+def thd_diag_enabled() -> bool:
+    """Return whether THD diagnostics are enabled."""
+    return os.environ.get("THD_DIAG", "0") not in ("0", "", "false", "False")
+
+
+def thd_diag_align_enabled() -> bool:
+    """Return whether THD alignment diagnostics are enabled."""
+    return os.environ.get("THD_DIAG_ALIGN", "0") not in ("0", "", "false", "False")
+
+
+def thd_diag_mrope_enabled() -> bool:
+    """Return whether THD MRoPE diagnostics are enabled."""
+    return os.environ.get("THD_DIAG_MROPE", "0") not in ("0", "", "false", "False")
 
 
 # copied from https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py
diff --git a/src/megatron/bridge/utils/common_utils.py b/src/megatron/bridge/utils/common_utils.py
index 35a3d507bd..0357a964cc 100644
--- a/src/megatron/bridge/utils/common_utils.py
+++ b/src/megatron/bridge/utils/common_utils.py
@@ -66,6 +66,11 @@ def get_rank_safe() -> int:
     return 0
 
 
+def is_rank_0() -> bool:
+    """Return whether the current process is rank 0."""
+    return get_rank_safe() == 0
+
+
 def get_world_size_safe() -> int:
     """Get the distributed world size safely, even if torch.distributed is not initialized.
 
diff --git a/tests/unit_tests/models/qwen_vl/test_diag_utils.py b/tests/unit_tests/models/qwen_vl/test_diag_utils.py
new file mode 100644
index 0000000000..17a13f4379
--- /dev/null
+++ b/tests/unit_tests/models/qwen_vl/test_diag_utils.py
@@ -0,0 +1,48 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from unittest.mock import patch
+
+from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.utils import (
+    thd_diag_align_enabled,
+    thd_diag_enabled,
+    thd_diag_mrope_enabled,
+)
+
+
+class TestDiagSwitches:
+    @patch.dict(os.environ, {"THD_DIAG": "1"}, clear=False)
+    def test_thd_diag_enabled_true(self):
+        assert thd_diag_enabled() is True
+
+    @patch.dict(os.environ, {"THD_DIAG": "false"}, clear=False)
+    def test_thd_diag_enabled_false(self):
+        assert thd_diag_enabled() is False
+
+    @patch.dict(os.environ, {"THD_DIAG_ALIGN": "1"}, clear=False)
+    def test_thd_diag_align_enabled_true(self):
+        assert thd_diag_align_enabled() is True
+
+    @patch.dict(os.environ, {"THD_DIAG_ALIGN": "False"}, clear=False)
+    def test_thd_diag_align_enabled_false(self):
+        assert thd_diag_align_enabled() is False
+
+    @patch.dict(os.environ, {"THD_DIAG_MROPE": "1"}, clear=False)
+    def test_thd_diag_mrope_enabled_true(self):
+        assert thd_diag_mrope_enabled() is True
+
+    @patch.dict(os.environ, {"THD_DIAG_MROPE": ""}, clear=False)
+    def test_thd_diag_mrope_enabled_false(self):
+        assert thd_diag_mrope_enabled() is False
diff --git a/tests/unit_tests/utils/test_common_utils.py b/tests/unit_tests/utils/test_common_utils.py
index e75ff20ba8..dec838259f 100644
--- a/tests/unit_tests/utils/test_common_utils.py
+++ b/tests/unit_tests/utils/test_common_utils.py
@@ -26,6 +26,7 @@
     get_master_port_safe,
     get_rank_safe,
     get_world_size_safe,
+    is_rank_0,
     is_last_rank,
     print_rank_0,
     print_rank_last,
@@ -178,6 +179,24 @@ def test_no_print_on_non_zero_rank(self, mock_print, mock_get_rank_safe):
         mock_print.assert_not_called()
 
 
+class TestIsRank0:
+    """Test is_rank_0 function."""
+
+    @patch("megatron.bridge.utils.common_utils.get_rank_safe")
+    def test_rank_zero(self, mock_get_rank_safe):
+        """Test is_rank_0 returns True when rank is 0."""
+        mock_get_rank_safe.return_value = 0
+        assert is_rank_0() is True
+        mock_get_rank_safe.assert_called_once()
+
+    @patch("megatron.bridge.utils.common_utils.get_rank_safe")
+    def test_non_zero_rank(self, mock_get_rank_safe):
+        """Test is_rank_0 returns False when rank is non-zero."""
+        mock_get_rank_safe.return_value = 3
+        assert is_rank_0() is False
+        mock_get_rank_safe.assert_called_once()
+
+
 class TestIsLastRank:
     """Test is_last_rank function."""
 

From 34f8e5802be6d0afa5551de87acc84a7275e8a7d Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Mon, 20 Apr 2026 09:52:17 +0800
Subject: [PATCH 44/51] chore: align Megatron-LM submodule pointer with main

Reset the 3rdparty/Megatron-LM gitlink to origin/main to avoid carrying submodule pointer drift in this PR.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 3rdparty/Megatron-LM | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM
index ec7ee7928f..23265d2db3 160000
--- a/3rdparty/Megatron-LM
+++ b/3rdparty/Megatron-LM
@@ -1 +1 @@
-Subproject commit ec7ee7928f78db084d55c0fa1bf8b6ac7c2a19a5
+Subproject commit 23265d2db37efa883b3983bce321c631898822f3

From a85bce1c841f0bc94cc6527e06839f68bc03a634 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Wed, 22 Apr 2026 11:45:22 +0800
Subject: [PATCH 45/51] fix: enforce strict qwen-vl task encoder config
 sourcing

Require task encoder vision and multimodal token settings to come from model/processor configuration and remove silent fallback behavior to prevent mismatched packing semantics.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 .../qwen_vl/data/energon/task_encoder.py      | 43 ++++++++++++-------
 .../bridge/recipes/qwen_vl/qwen3_vl.py        | 24 ++++++++++-
 .../qwen_vl/data/energon/test_task_encoder.py | 33 +++++++++++---
 3 files changed, 77 insertions(+), 23 deletions(-)

diff --git a/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py b/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
index a04afca933..7237919606 100644
--- a/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
+++ b/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
@@ -136,19 +136,25 @@ def process_vision(
 
 
 def _resolve_hf_mm_token_ids(hf_tokenizer):
-    """Resolve HF tokenizer ids for <image> and <video> tokens without nemo constants."""
+    """Resolve HF tokenizer ids for ``<image>`` and ``<video>`` tokens strictly."""
 
-    def _get(token_str: str, default_id: int) -> int:
+    def _get(token_str: str) -> int:
         token_attr = getattr(hf_tokenizer, f"{token_str.strip('<>')}_token_id", None)
         if token_attr is not None:
             return int(token_attr)
         try:
-            return int(hf_tokenizer.convert_tokens_to_ids(token_str))
-        except Exception:
-            return default_id
-
-    image_id = _get("<image>", 151655)
-    video_id = _get("<video>", 151656)
+            token_id = hf_tokenizer.convert_tokens_to_ids(token_str)
+        except Exception as exc:
+            raise ValueError(f"Failed to resolve token id for {token_str!r}") from exc
+        if token_id is None:
+            raise ValueError(f"Tokenizer returned None for {token_str!r}")
+        token_id = int(token_id)
+        if token_id < 0:
+            raise ValueError(f"Tokenizer returned invalid id {token_id} for {token_str!r}")
+        return token_id
+
+    image_id = _get("<image>")
+    video_id = _get("<video>")
     return image_id, video_id
 
 
@@ -264,12 +270,14 @@ def __init__(
         self,
         tokenizer,
         image_processor,
-        temporal_patch_size: int = 2,
-        spatial_merge_size: int = 2,
-        patch_size: int = 14,
-        max_padding_length: int = 4096,
-        min_pixels: int = 200704,
-        max_pixels: int = 1003520,
+        temporal_patch_size: int,
+        spatial_merge_size: int,
+        patch_size: int,
+        image_token_id: int,
+        video_token_id: int,
+        max_padding_length: int,
+        min_pixels: int,
+        max_pixels: int,
     ):
         super().__init__()
 
@@ -284,7 +292,12 @@ def __init__(
         self.patch_size = patch_size
 
         self.seq_len = max_padding_length
-        self.image_token_id, self.video_token_id = _resolve_hf_mm_token_ids(self.hf_tokenizer)
+        self.image_token_id = int(image_token_id)
+        self.video_token_id = int(video_token_id)
+        if self.image_token_id < 0 or self.video_token_id < 0:
+            raise ValueError(
+                f"image_token_id/video_token_id must be non-negative, got {self.image_token_id}/{self.video_token_id}"
+            )
 
     def encode_sample(self, sample: ChatMLSample):
         """
diff --git a/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py b/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py
index 03963a70f4..2b6f8c2456 100644
--- a/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py
+++ b/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py
@@ -43,7 +43,7 @@
 from megatron.bridge.data.vlm_datasets import MockVLMConversationProvider
 from megatron.bridge.peft.base import PEFT
 from megatron.bridge.recipes.common import _peft_common_vlm, _sft_common_vlm
-from megatron.bridge.recipes.qwen_vl.data.energon.task_encoder import QwenVLTaskEncoder
+from megatron.bridge.recipes.qwen_vl.data.energon.task_encoder import QwenVLTaskEncoder, _resolve_hf_mm_token_ids
 from megatron.bridge.recipes.utils.finetune_utils import default_peft_config
 from megatron.bridge.recipes.utils.optimizer_utils import distributed_fused_adam_with_cosine_annealing
 from megatron.bridge.recipes.utils.tokenizer_utils import DEFAULT_NULL_TOKENIZER_VOCAB_SIZE
@@ -281,6 +281,8 @@ def _make_energon_dataset(
     hf_path: str, seq_length: int, micro_batch_size: int, global_batch_size: int
 ) -> EnergonProvider:
     """Create an EnergonProvider dataset config for Qwen3-VL recipes."""
+    bridge = AutoBridge.from_hf_pretrained(hf_path)
+    model_cfg = bridge.to_megatron_provider(load_weights=False)
     tokenizer = AutoTokenizer.from_pretrained(hf_path)
     # Prefer Qwen3VLProcessor on newer transformers; fall back to AutoProcessor
     # for older container images that do not yet expose that symbol or lack
@@ -298,10 +300,30 @@ def _make_energon_dataset(
         from transformers import AutoProcessor
 
         image_processor = AutoProcessor.from_pretrained(hf_path)
+    image_token_id, video_token_id = _resolve_hf_mm_token_ids(tokenizer)
+    if image_token_id != model_cfg.image_token_id or video_token_id != model_cfg.video_token_id:
+        raise ValueError(
+            "Tokenizer multimodal token ids do not match model config: "
+            f"tokenizer(image={image_token_id}, video={video_token_id}) vs "
+            f"model(image={model_cfg.image_token_id}, video={model_cfg.video_token_id})"
+        )
+
+    min_pixels = getattr(image_processor, "min_pixels", None)
+    max_pixels = getattr(image_processor, "max_pixels", None)
+    if min_pixels is None or max_pixels is None:
+        raise ValueError("Processor must expose both min_pixels and max_pixels for strict task encoder config")
+
     task_encoder = QwenVLTaskEncoder(
         tokenizer=tokenizer,
         image_processor=image_processor,
+        temporal_patch_size=model_cfg.temporal_patch_size,
+        spatial_merge_size=model_cfg.spatial_merge_size,
+        patch_size=model_cfg.patch_size,
+        image_token_id=model_cfg.image_token_id,
+        video_token_id=model_cfg.video_token_id,
         max_padding_length=seq_length,
+        min_pixels=int(min_pixels),
+        max_pixels=int(max_pixels),
     )
     return EnergonProvider(
         path="",  # Must be set via CLI override: dataset.path=<path>
diff --git a/tests/unit_tests/recipes/qwen_vl/data/energon/test_task_encoder.py b/tests/unit_tests/recipes/qwen_vl/data/energon/test_task_encoder.py
index a3d7a9c5b6..336550b503 100644
--- a/tests/unit_tests/recipes/qwen_vl/data/energon/test_task_encoder.py
+++ b/tests/unit_tests/recipes/qwen_vl/data/energon/test_task_encoder.py
@@ -43,6 +43,28 @@ def cleanup_local_folder():
     pass
 
 
+@pytest.fixture
+def qwenvl_encoder_config():
+    """Centralized encoder config for unit tests."""
+    return {
+        "temporal_patch_size": 2,
+        "patch_size": 14,
+        "spatial_merge_size": 2,
+        "image_token_id": 151655,
+        "video_token_id": 151656,
+        "max_padding_length": 128,
+        "min_pixels": 200704,
+        "max_pixels": 1003520,
+    }
+
+
+@pytest.fixture(autouse=True)
+def inject_qwenvl_encoder_config(request, qwenvl_encoder_config):
+    """Inject shared encoder config into unittest-style test cases."""
+    if request.instance is not None:
+        request.instance.qwenvl_encoder_config = qwenvl_encoder_config
+
+
 class TestHelperFunctions(unittest.TestCase):
     def test_find_pattern_indices(self):
         seq = np.array([1, 2, 3, 4, 5])
@@ -124,14 +146,13 @@ def test_falls_back_to_convert_tokens_to_ids(self):
         self.assertEqual(image_id, 300)
         self.assertEqual(video_id, 400)
 
-    def test_returns_defaults_when_all_fail(self):
+    def test_raises_when_all_resolution_paths_fail(self):
         tokenizer = MagicMock()
         tokenizer.image_token_id = None
         tokenizer.video_token_id = None
         tokenizer.convert_tokens_to_ids.side_effect = Exception("not found")
-        image_id, video_id = _resolve_hf_mm_token_ids(tokenizer)
-        self.assertEqual(image_id, 151655)
-        self.assertEqual(video_id, 151656)
+        with self.assertRaises(ValueError):
+            _resolve_hf_mm_token_ids(tokenizer)
 
 
 class TestVideoHandler(unittest.TestCase):
@@ -183,9 +204,7 @@ def setUp(self):
         self.encoder = QwenVLTaskEncoder(
             tokenizer=self.tokenizer,
             image_processor=self.image_processor,
-            max_padding_length=128,
-            patch_size=14,
-            spatial_merge_size=2,
+            **self.qwenvl_encoder_config,
         )
 
     def test_process_vision(self):

From e4a88ee990b7e810d93802a5c22f6c2e9f67fadf Mon Sep 17 00:00:00 2001
From: yaoyu-33 <yaoyu.094@gmail.com>
Date: Tue, 28 Apr 2026 01:49:39 -0700
Subject: [PATCH 46/51] chore: ruff isort fix on merge-resolved files

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
---
 src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py | 2 +-
 src/megatron/bridge/training/vlm_step.py                       | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
index e390dccda7..699551edc9 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
@@ -43,8 +43,8 @@
     collapse_thw,
     get_dist_train_vision_dp_data,
     get_vision_cp_data,
-    pack_dist_train_vision_module_output,
     is_rank_0,
+    pack_dist_train_vision_module_output,
     preprocess_packed_seqs,
     qwen3vl_cp_split,
     reorganize_inputs,
diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index 86f8d019d1..857140810e 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -13,8 +13,8 @@
 # limitations under the License.
 
 import logging
-import os
 import math
+import os
 from functools import partial
 from typing import Any, Iterable
 

From 50b55572e0f5ed11cffe812a682a80dc76beae0a Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Tue, 28 Apr 2026 17:16:25 +0800
Subject: [PATCH 47/51] fix: resolve lint blockers in qwen-vl checks

Align fallback and packing code with lint requirements so PR CI can proceed past lint gating.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 src/megatron/bridge/models/conversion/auto_bridge.py         | 2 ++
 .../bridge/recipes/qwen_vl/data/energon/task_encoder.py      | 5 ++---
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/megatron/bridge/models/conversion/auto_bridge.py b/src/megatron/bridge/models/conversion/auto_bridge.py
index 45e14876c2..ddb411142c 100644
--- a/src/megatron/bridge/models/conversion/auto_bridge.py
+++ b/src/megatron/bridge/models/conversion/auto_bridge.py
@@ -61,8 +61,10 @@
     )
 
     def is_quantized(_model: object) -> bool:
+        """Fallback quantization probe when ModelOpt is unavailable."""
         return False
 
+
 MegatronModelT = TypeVar("MegatronModelT", bound=MegatronModule)
 DataclassT = TypeVar("DataclassT")
 
diff --git a/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py b/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
index 7237919606..12ec6d9f15 100644
--- a/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
+++ b/src/megatron/bridge/recipes/qwen_vl/data/energon/task_encoder.py
@@ -482,7 +482,6 @@ def pack_selected_samples(self, samples: List[QwenVLTaskSample]) -> QwenVLTaskSa
         total_video_tokens = 0
         total_text_tokens = 0
         total_vit_patches = 0
-        merge_length = self.merge_size ** 2
 
         for s in samples:
             seq_len = len(s.text)
@@ -541,8 +540,8 @@ def pack_selected_samples(self, samples: List[QwenVLTaskSample]) -> QwenVLTaskSa
         packed_video_mask = torch.cat(all_video_masks, dim=0)
 
         cu_lengths_list = [0]
-        for l in sub_lengths:
-            cu_lengths_list.append(cu_lengths_list[-1] + l)
+        for length in sub_lengths:
+            cu_lengths_list.append(cu_lengths_list[-1] + length)
         cu_lengths = torch.tensor(cu_lengths_list, dtype=torch.int32)
 
         total_packed = int(cu_lengths[-1].item())

From d2b73cc1b8fcc639259837d27486d9709e0b5dbf Mon Sep 17 00:00:00 2001
From: yaoyu-33 <yaoyu.094@gmail.com>
Date: Wed, 29 Apr 2026 19:34:38 -0700
Subject: [PATCH 48/51] style: ruff-format vlm_step.py after merge

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
---
 src/megatron/bridge/training/vlm_step.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index 60ffc06fe5..24a8f036b6 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -408,7 +408,9 @@ def _ceil_to_mult(n: int, mult: int) -> int:
 
         if force_single_segment_cu and tokens_or_input is not None:
             seq_len = int(tokens_or_input.shape[-1])
-            device = energon_cu_seqlens.device if isinstance(energon_cu_seqlens, torch.Tensor) else tokens_or_input.device
+            device = (
+                energon_cu_seqlens.device if isinstance(energon_cu_seqlens, torch.Tensor) else tokens_or_input.device
+            )
             energon_cu_seqlens = torch.tensor([0, seq_len], dtype=torch.int32, device=device)
             energon_max_seqlen = torch.tensor(seq_len, dtype=torch.int32, device=device)
             energon_cu_argmin = torch.tensor(2, dtype=torch.int64)

From db737642525b9e8fc0d7c70573889d660c9262c6 Mon Sep 17 00:00:00 2001
From: yaoyu-33 <yaoyu.094@gmail.com>
Date: Wed, 29 Apr 2026 19:39:39 -0700
Subject: [PATCH 49/51] refactor(vlm_step): use shared is_rank_0() helper
 instead of local _rank0()

The local _rank0() in vlm_step.py duplicated the shared is_rank_0()
helper in megatron.bridge.utils.common_utils. Use the shared helper for
consistency with model.py (which was switched in 89cfeda1).

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
---
 src/megatron/bridge/training/vlm_step.py | 21 ++++++++-------------
 1 file changed, 8 insertions(+), 13 deletions(-)

diff --git a/src/megatron/bridge/training/vlm_step.py b/src/megatron/bridge/training/vlm_step.py
index 24a8f036b6..4b0b67e249 100644
--- a/src/megatron/bridge/training/vlm_step.py
+++ b/src/megatron/bridge/training/vlm_step.py
@@ -35,18 +35,13 @@
     pad_or_truncate_pos_to_len,
 )
 from megatron.bridge.training.utils.pg_utils import get_pg_collection
+from megatron.bridge.utils.common_utils import is_rank_0
 
 
 logger = logging.getLogger(__name__)
 _SWITCH_LOGGED: set[str] = set()
 
 
-def _rank0() -> bool:
-    if not torch.distributed.is_available() or not torch.distributed.is_initialized():
-        return True
-    return torch.distributed.get_rank() == 0
-
-
 def _thd_diag_enabled() -> bool:
     return _switch_enabled("THD_DIAG")
 
@@ -57,7 +52,7 @@ def _thd_diag_boundary_enabled() -> bool:
 
 def _switch_enabled(name: str) -> bool:
     enabled = os.environ.get(name, "0") not in ("0", "", "false", "False")
-    if enabled and _rank0() and name not in _SWITCH_LOGGED:
+    if enabled and is_rank_0() and name not in _SWITCH_LOGGED:
         logger.info("[THD_SWITCH] %s=1 enabled", name)
         _SWITCH_LOGGED.add(name)
     return enabled
@@ -333,7 +328,7 @@ def get_batch(data_iterator: Iterable, cfg: ConfigContainer, use_mtp: bool = Fal
     force_bshd = _thd_force_bshd_enabled()
     force_single_segment_cu = _thd_force_single_segment_cu_enabled()
 
-    if force_bshd and force_single_segment_cu and _rank0():
+    if force_bshd and force_single_segment_cu and is_rank_0():
         logger.warning(
             "[THD_SWITCH] THD_FORCE_BSHD=1 overrides THD_FORCE_SINGLE_SEGMENT_CU=1 (single-segment CU ignored)."
         )
@@ -414,7 +409,7 @@ def _ceil_to_mult(n: int, mult: int) -> int:
             energon_cu_seqlens = torch.tensor([0, seq_len], dtype=torch.int32, device=device)
             energon_max_seqlen = torch.tensor(seq_len, dtype=torch.int32, device=device)
             energon_cu_argmin = torch.tensor(2, dtype=torch.int64)
-            if _thd_diag_enabled() and _rank0():
+            if _thd_diag_enabled() and is_rank_0():
                 logger.info("[THD_SWITCH] force single-segment cu_seqlens=[0, %d]", seq_len)
 
         # Log detailed packed-iteration diagnostics when explicitly enabled.
@@ -596,7 +591,7 @@ def forward_step(
             cu_padded[-1] = physical_seq_len
             max_seqlen_out = (cu_padded[1:] - cu_padded[:-1]).max()
         else:
-            if last_boundary < physical_seq_len and _thd_diag_enabled() and _rank0():
+            if last_boundary < physical_seq_len and _thd_diag_enabled() and is_rank_0():
                 logger.info(
                     "[THD_SWITCH] skip packed-pos preprocess: preserve original cu_seqlens last_boundary=%d physical_seq_len=%d",
                     last_boundary,
@@ -632,7 +627,7 @@ def forward_step(
         # Pass unpadded boundaries only to Qwen model's MRoPE construction.
         forward_args["rope_cu_seqlens"] = cu_unpadded
         forward_args["moe_padding_mask"] = moe_padding_mask
-        if _thd_diag_enabled() and _rank0():
+        if _thd_diag_enabled() and is_rank_0():
             logger.info(
                 "[THD_DIAG][packed] physical_seq_len=%d cu_unpadded_last=%d cu_padded_last=%d implicit_pad=%d max_seqlen=%d moe_padding_tokens=%d",
                 int(physical_seq_len),
@@ -642,7 +637,7 @@ def forward_step(
                 int(max_seqlen_out.item()) if torch.is_tensor(max_seqlen_out) else int(max_seqlen_out),
                 int(moe_padding_mask.sum().item()),
             )
-        if _thd_diag_boundary_enabled() and _rank0():
+        if _thd_diag_boundary_enabled() and is_rank_0():
             cu_padded_cpu = cu_padded.detach().cpu()
             cu_unpadded_cpu = cu_unpadded.detach().cpu()
             seg_lens_padded = (cu_padded_cpu[1:] - cu_padded_cpu[:-1]).tolist()
@@ -714,7 +709,7 @@ def forward_step(
             else:
                 output_tensor = model_output
 
-    if _thd_diag_enabled() and _rank0() and loss_mask is not None:
+    if _thd_diag_enabled() and is_rank_0() and loss_mask is not None:
         loss_mask_flat = loss_mask.view(-1)
         valid_loss_tokens = int((loss_mask_flat > 0).sum().item())
         total_loss_tokens = int(loss_mask_flat.numel())

From b319cd165691e80aa02e08faac894a9e64766ccb Mon Sep 17 00:00:00 2001
From: yaoyu-33 <yaoyu.094@gmail.com>
Date: Fri, 1 May 2026 01:39:30 -0700
Subject: [PATCH 50/51] style: ruff lint and format fixes

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
---
 src/megatron/bridge/data/energon/energon_provider.py |  6 +++---
 .../models/qwen_vl/modelling_qwen3_vl/model.py       | 12 +++---------
 .../models/qwen_vl/modelling_qwen3_vl/utils.py       |  1 -
 src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py     |  1 -
 src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py      |  1 +
 .../qwen_vl/data/energon/test_task_encoder.py        |  2 +-
 tests/unit_tests/utils/test_common_utils.py          |  2 +-
 7 files changed, 9 insertions(+), 16 deletions(-)

diff --git a/src/megatron/bridge/data/energon/energon_provider.py b/src/megatron/bridge/data/energon/energon_provider.py
index 7c17460cde..ded5f8d308 100644
--- a/src/megatron/bridge/data/energon/energon_provider.py
+++ b/src/megatron/bridge/data/energon/energon_provider.py
@@ -54,9 +54,9 @@ class EnergonProvider(DatasetProvider):
     def build_datasets(self, context: DatasetBuildContext):
         resolved_path = self.path
         if self.cord_bin_id is not None and self.cord_bin_id != "":
-            assert (
-                self.cord_bins_root
-            ), "EnergonProvider.cord_bins_root must be set when dataset.cord_bin_id is provided."
+            assert self.cord_bins_root, (
+                "EnergonProvider.cord_bins_root must be set when dataset.cord_bin_id is provided."
+            )
             resolved_path = os.path.join(self.cord_bins_root, f"{self.cord_bin_prefix}{self.cord_bin_id}")
 
         assert resolved_path, "EnergonProvider.path must be set. Use CLI override: dataset.path=<path>"
diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
index 699551edc9..b753b05e51 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
@@ -684,12 +684,8 @@ def _mask_summary(mask: torch.Tensor, max_items: int = 8) -> tuple[int, str]:
                 total_len = input_ids.shape[1]
                 flat_ids = input_ids.squeeze(0)
 
-                sub_ids = torch.zeros(
-                    num_seqs, max_subseq_len, dtype=input_ids.dtype, device=input_ids.device
-                )
-                sub_mask = torch.zeros(
-                    num_seqs, max_subseq_len, dtype=torch.int32, device=input_ids.device
-                )
+                sub_ids = torch.zeros(num_seqs, max_subseq_len, dtype=input_ids.dtype, device=input_ids.device)
+                sub_mask = torch.zeros(num_seqs, max_subseq_len, dtype=torch.int32, device=input_ids.device)
                 for i, sl in enumerate(seq_lens):
                     start = int(cu[i].item())
                     sl_int = int(sl)
@@ -707,9 +703,7 @@ def _mask_summary(mask: torch.Tensor, max_items: int = 8) -> tuple[int, str]:
                     attention_mask=sub_mask,
                 )  # [3, num_seqs, max_subseq_len]
 
-                packed_pos = torch.zeros(
-                    3, 1, total_len, dtype=position_ids.dtype, device=position_ids.device
-                )
+                packed_pos = torch.zeros(3, 1, total_len, dtype=position_ids.dtype, device=position_ids.device)
                 for i, sl in enumerate(seq_lens):
                     start = int(cu[i].item())
                     sl_int = int(sl)
diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py
index 51fdc2538e..5d12c916cb 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py
@@ -27,7 +27,6 @@
 from torch import nn
 
 from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.transformer_config import Qwen3VLTransformerConfig
-from megatron.bridge.utils.common_utils import is_rank_0
 
 
 def thd_diag_enabled() -> bool:
diff --git a/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py b/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
index 2642e38bfb..a90d908a18 100644
--- a/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
+++ b/src/megatron/bridge/recipes/qwen_vl/qwen35_vl.py
@@ -26,7 +26,6 @@
 from typing_extensions import Unpack
 
 from megatron.bridge import AutoBridge
-from megatron.bridge.data.energon.energon_provider import EnergonProvider
 from megatron.bridge.peft.base import PEFT
 from megatron.bridge.recipes.common import _peft_common_vlm, _sft_common_vlm
 from megatron.bridge.recipes.qwen_vl.qwen3_vl import (
diff --git a/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py b/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py
index 2b6f8c2456..ae6b428f99 100644
--- a/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py
+++ b/src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py
@@ -27,6 +27,7 @@
 from transformers import AutoTokenizer
 from typing_extensions import TypedDict, Unpack
 
+
 logger = logging.getLogger(__name__)
 
 try:
diff --git a/tests/unit_tests/recipes/qwen_vl/data/energon/test_task_encoder.py b/tests/unit_tests/recipes/qwen_vl/data/energon/test_task_encoder.py
index 336550b503..163949c963 100644
--- a/tests/unit_tests/recipes/qwen_vl/data/energon/test_task_encoder.py
+++ b/tests/unit_tests/recipes/qwen_vl/data/energon/test_task_encoder.py
@@ -31,8 +31,8 @@
     _resolve_hf_mm_token_ids,
     convert_to_qwenvl_content,
     find_pattern_indices,
-    greedy_knapsack,
     get_ltor_masks_and_position_ids,
+    greedy_knapsack,
     process_vision,
     videohandler,
 )
diff --git a/tests/unit_tests/utils/test_common_utils.py b/tests/unit_tests/utils/test_common_utils.py
index dec838259f..6ff4b076a8 100644
--- a/tests/unit_tests/utils/test_common_utils.py
+++ b/tests/unit_tests/utils/test_common_utils.py
@@ -26,8 +26,8 @@
     get_master_port_safe,
     get_rank_safe,
     get_world_size_safe,
-    is_rank_0,
     is_last_rank,
+    is_rank_0,
     print_rank_0,
     print_rank_last,
 )

From 6c46f04feddf0492371d10846f4374614aa36e03 Mon Sep 17 00:00:00 2001
From: Daisy Gao <daisyg@nvidia.com>
Date: Wed, 6 May 2026 09:02:53 +0800
Subject: [PATCH 51/51] fix: restore qwen3-vl is_rank_0 re-export

Re-export is_rank_0 from qwen3-vl utils to keep downstream imports in unit-test conftest working.

Signed-off-by: Daisy Gao <daisyg@nvidia.com>
---
 src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py
index 5d12c916cb..9f87bb96a0 100644
--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py
+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py
@@ -27,6 +27,7 @@
 from torch import nn
 
 from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.transformer_config import Qwen3VLTransformerConfig
+from megatron.bridge.utils.common_utils import is_rank_0  # noqa: F401
 
 
 def thd_diag_enabled() -> bool: