Merge remote-tracking branch 'origin/main' into codex/pr5479-hybrid-model

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

# Conflicts:
#	src/megatron/bridge/training/model_load_save.py
#	tests/unit_tests/training/test_checkpointing.py

Author: yaoyu-33

.dev.commit

@@ -1 +1 @@
-b3b1e2d793c596c2fa61076d001ab93f294e4e87
+2f2f8ebaef90f4d62a18a43808398223b1fc5809

.github/scripts/test_docker_dependency_layers.sh

@@ -252,9 +252,18 @@ if ! grep -Fq \
   echo "Bridge's lock metadata must preserve the MCore TE extra" >&2
   exit 1
 fi
-if grep -q 'transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@' pyproject.toml || \
-  grep -q '^name = "transformer-engine"$' pyproject.toml; then
-  echo "Bridge must inherit the TransformerEngine source and metadata from the selected MCore ref" >&2
+if grep -q 'transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@' pyproject.toml; then
+  echo "Bridge must inherit the TransformerEngine source from the selected MCore ref" >&2
+  exit 1
+fi
+
+if ! grep -Fq '    NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD=1 \' "$dockerfile"; then
+  echo "CI builds must disable TransformerEngine's recursive build-time submodule fetch" >&2
+  exit 1
+fi
+if grep -R -qE 'git submodule update.*--recursive|git submodule update --init --recursive' \
+  "$dockerfile" docker .github/actions; then
+  echo "CI build surfaces must not fetch recursive submodules at build time" >&2
   exit 1
 fi
 

.main.commit

@@ -1 +1 @@
-2a75ac12c54ba6a42b34024756820bf819a2e25f
+731b791469004f8fdcb896e65d610d6b1b0ebb32

3rdparty/Megatron-LM

@@ -55,6 +55,7 @@ ENV PATH="/opt/venv/bin:/usr/local/bin:$PATH" \
     UV_VERSION="0.7.2" \
     UV_HTTP_TIMEOUT=120 \
     NVTE_BUILD_NUM_PHILOX_ROUNDS=3 \
+    NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD=1 \
     HYBRID_EP_MULTINODE=1 \
     RDMA_CORE_HOME=/opt/rdma-core/build \
     LD_LIBRARY_PATH="/usr/local/cuda/lib64/:$LD_LIBRARY_PATH"
@@ -192,7 +193,7 @@ COPY --chown=1000:1000 . /opt/Megatron-Bridge
 # prepared environment, and uses /nemo_run for runtime outputs. Keep those paths owned by the
 # existing unprivileged base-image user instead of relying on a root final user.
 RUN install -d -o 1000 -g 1000 -m 0755 /nemo_run /home/ubuntu/.cache && \
-    chown 1000:1000 /opt/Megatron-Bridge
+    chown -R 1000:1000 /opt/Megatron-Bridge
 
 ENV HOME=/home/ubuntu
 USER 1000:1000

docker/Dockerfile.ci

@@ -22,12 +22,12 @@ enablement through context parallelism:
 | Path | Use case | Key config |
 |---|---|---|
 | Offline packed SFT | Text-only finetuning | `enable_offline_packing=True` plus `offline_packing_specs` |
-| Direct-HF/VLM in-batch packing | Direct Hugging Face and supported VLM finetuning | `enable_in_batch_packing=True` |
+| Runtime in-batch packing | GPT-SFT JSONL, Direct Hugging Face, and supported VLM finetuning | `enable_in_batch_packing=True` |
 | Energon online packing | Qwen-VL data using the model-owned Energon collator | `packing_buffer_size=<candidate samples per worker>` |
 | Long-context (CP) | Pretrain / finetune at 16K-128K+ | `context_parallel_size > 1` |
 
 These are related but they are not the same knob. Offline packed SFT and
-Direct-HF/VLM in-batch packing solve padding waste; long-context training
+runtime in-batch packing solve padding waste; long-context training
 primarily addresses activation memory and communication tradeoffs at larger
 sequence lengths.
 
@@ -40,6 +40,15 @@ non-packed padding remains in `megatron.bridge.data.collators`. Use
 `scripts/training/prepare_gpt_sft_packed_data.py` when packed GPT SFT artifacts
 should be prepared before launching training.
 
+For `GPTSFTDatasetConfig`, in-batch packing works with both local mmap JSONL
+schemas: prompt/completion (`GPTSFTDataset`) and chat
+(`GPTSFTChatDataset`). Tokenization remains lazy: workers mmap the JSONL and
+read, parse, and tokenize only the rows selected for the current logical
+microbatch. Collation then concatenates those rows into one physical THD batch
+row; it does not materialize an offline dataset or load the full source into
+RAM. Use a microbatch-yielding `single` or `cyclic` dataloader; GPT-SFT
+in-batch packing does not support the global-batch `batch` dataloader.
+
 ## When to Use It
 
 Packed sequences are a good fit when all of the following are true:
@@ -103,13 +112,34 @@ when using CUDA graphs. CUDA graphs additionally require
 `pad_cu_seqlens=true` and packing metadata. Ordinary eager offline packing
 does not universally require fixed-width padding.
 
+## Choosing Runtime In-Batch Packing
+
+Use GPT-SFT in-batch packing when retaining the original JSONL is preferable
+to generating packed Parquet artifacts. Enable it directly on the dataset
+config and use a logical micro-batch larger than one:
+
+```text
+dataset.enable_in_batch_packing=true
+dataset.dataloader_type=single
+train.micro_batch_size=4
+```
+
+The collator preserves each sample's prompt/completion or chat loss mask and
+emits current MCore packed metadata (`cu_seqlens_q`, `cu_seqlens_kv`, and the
+corresponding padded boundaries when CP/SP alignment is required). The model
+sees one physical THD row. `enable_in_batch_packing` and
+`enable_offline_packing` are mutually exclusive. The `batch` dataloader is not
+supported for GPT-SFT in-batch packing; use `single` or `cyclic`.
+
 ## Stable Constraints
 
 The durable constraints for packed sequences in Bridge are:
 
 - offline packed SFT requires configured `micro_batch_size == 1`
-- Direct-HF/VLM in-batch packing requires configured `micro_batch_size > 1`;
+- GPT-SFT/Direct-HF/VLM in-batch packing requires configured `micro_batch_size > 1`;
   collation flattens those input rows into one physical THD batch row
+- GPT-SFT in-batch packing requires `dataloader_type="single"` or `"cyclic"`;
+  the global-batch `"batch"` dataloader is not supported
 - Energon online packing currently supports the eager Qwen-VL collator path,
   requires physical `micro_batch_size == 1`, the generic `vlm_step`, per-token loss, and
   `ddp.average_in_collective=False`
@@ -130,7 +160,7 @@ The durable constraints for packed sequences in Bridge are:
   overlap is disabled with a warning so training uses the non-overlapped path
 - when context parallelism is used, sequence length must satisfy the standard
   CP divisibility constraints
-- Direct-HF sequence length must also satisfy the LCM of the training and
+- GPT-SFT and Direct-HF sequence length must also satisfy the LCM of the training and
   evaluation CP constraints and `CP * TP` when sequence parallelism is enabled
 - for fine-tuning with CP enabled, per-token loss behavior and reduction
   settings matter

docs/training/packed-sequences.md

@@ -279,6 +279,7 @@ def __init__(
         pixel_values=None,
         image_grid_thw=None,
         inference_context=None,
+        mm_token_type_ids=None,
     ):
         self.batch = dict(
             tokens=input_ids,
@@ -292,6 +293,8 @@ def __init__(
             self.batch["pixel_values"] = pixel_values
         if image_grid_thw is not None:
             self.batch["image_grid_thw"] = image_grid_thw
+        if mm_token_type_ids is not None:
+            self.batch["mm_token_type_ids"] = mm_token_type_ids
 
         self._yielded = False
 
@@ -332,6 +335,8 @@ def vlm_forward_step(data_iterator, model, **kwargs) -> torch.Tensor:
         forward_args["pixel_values"] = batch["pixel_values"]
     if "image_grid_thw" in batch:
         forward_args["image_grid_thw"] = batch["image_grid_thw"]
+    if "mm_token_type_ids" in batch:
+        forward_args["mm_token_type_ids"] = batch["mm_token_type_ids"]
 
     def loss_func(x, **kwargs):
         return x
@@ -441,7 +446,8 @@ def process_inputs(tokenizer, processor, image_path: Optional[str], prompt: str,
         tp_size: Tensor parallel size for padding sequence length
 
     Returns:
-        Tuple of (input_ids, pixel_values, image_grid_thw, token_type_ids)
+        Tuple of (input_ids, pixel_values, image_grid_thw, token_type_ids,
+        mm_token_type_ids)
     """
     if is_vl_model and image_path:
         messages = [
@@ -465,11 +471,15 @@ def process_inputs(tokenizer, processor, image_path: Optional[str], prompt: str,
         token_type_ids = inputs.get("token_type_ids")
         if token_type_ids is not None:
             token_type_ids = pad_input_ids_to_tp_multiple(token_type_ids, tp_size, 0)
+        mm_token_type_ids = inputs.get("mm_token_type_ids")
+        if mm_token_type_ids is not None:
+            mm_token_type_ids = pad_input_ids_to_tp_multiple(mm_token_type_ids, tp_size, 0)
         return (
             input_ids,
             inputs.get("pixel_values"),
             inputs.get("image_grid_thw"),
             token_type_ids,
+            mm_token_type_ids,
         )
     else:
         # Text-only processing for both VL models without images and regular LLMs
@@ -480,7 +490,7 @@ def process_inputs(tokenizer, processor, image_path: Optional[str], prompt: str,
             # Use tokenizer for regular LLMs
             inputs = tokenizer(prompt, return_tensors="pt")
         input_ids = pad_input_ids_to_tp_multiple(inputs.input_ids, tp_size, tokenizer.pad_token_id or 0)
-        return input_ids, None, None, None
+        return input_ids, None, None, None, None
 
 
 def _load_hf_model(args, is_vl_model: bool):
@@ -573,7 +583,16 @@ def _get_hf_forward_model(hf_model, pixel_values):
     return hf_model
 
 
-def _run_hf_inference(hf_model, input_ids, pixel_values, image_grid_thw, tokenizer, *, token_type_ids=None):
+def _run_hf_inference(
+    hf_model,
+    input_ids,
+    pixel_values,
+    image_grid_thw,
+    tokenizer,
+    *,
+    token_type_ids=None,
+    mm_token_type_ids=None,
+):
     """Run HuggingFace model inference and return results.
 
     Args:
@@ -582,7 +601,8 @@ def _run_hf_inference(hf_model, input_ids, pixel_values, image_grid_thw, tokeniz
         pixel_values: Pixel values for vision models (optional).
         image_grid_thw: Image grid dimensions (optional).
         tokenizer: Tokenizer for decoding.
-        token_type_ids: Multimodal token type IDs (optional).
+        token_type_ids: Legacy multimodal token type IDs (optional).
+        mm_token_type_ids: Multimodal token type IDs used for M-RoPE (optional).
 
     Returns:
         Tuple of (hf_logits, hf_next_token, hf_logits_stats, hf_top5_info, logits_shape).
@@ -613,6 +633,8 @@ def _run_hf_inference(hf_model, input_ids, pixel_values, image_grid_thw, tokeniz
             hf_inputs["image_grid_thw"] = image_grid_thw.to(hf_device)
         if token_type_ids is not None:
             hf_inputs["token_type_ids"] = token_type_ids.to(hf_device)
+        if mm_token_type_ids is not None:
+            hf_inputs["mm_token_type_ids"] = mm_token_type_ids.to(hf_device)
 
         hf_output = hf_forward_model(**hf_inputs)
 
@@ -875,7 +897,7 @@ def compare_models_one_step(args) -> None:
 
     # Process inputs
     print_rank_0(f"Processing inputs - Prompt: '{args.prompt}', Image: {args.image_path}")
-    input_ids, pixel_values, image_grid_thw, token_type_ids = process_inputs(
+    input_ids, pixel_values, image_grid_thw, token_type_ids, mm_token_type_ids = process_inputs(
         tokenizer, processor, args.image_path, args.prompt, is_vl_model, args.tp
     )
 
@@ -887,6 +909,8 @@ def compare_models_one_step(args) -> None:
         image_grid_thw = image_grid_thw.cuda()
     if token_type_ids is not None:
         token_type_ids = token_type_ids.cuda()
+    if mm_token_type_ids is not None:
+        mm_token_type_ids = mm_token_type_ids.cuda()
 
     print_rank_0(f"Input shape: {input_ids.shape}")
     print_rank_0(f"Pixel values shape: {pixel_values.shape if pixel_values is not None else 'None'}")
@@ -904,6 +928,7 @@ def compare_models_one_step(args) -> None:
             image_grid_thw,
             tokenizer,
             token_type_ids=token_type_ids,
+            mm_token_type_ids=mm_token_type_ids,
         )
 
     del hf_model
@@ -949,6 +974,7 @@ def compare_models_one_step(args) -> None:
                 attention_mask,
                 pixel_values,
                 image_grid_thw,
+                mm_token_type_ids=mm_token_type_ids,
             )
             megatron_output = fwd_bwd_function(
                 forward_step_func=vlm_forward_step,