Merge commit 'b010cef540ccae9d207aed2dbacb1b12ecc5f654' into chore/bump-mcore-main-260813

Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>

Author: svcnemo-autobot

.github/actions/test-template/action.yml

@@ -59,6 +59,14 @@ inputs:
   runner:
     description: "Runner to use for test"
     required: true
+  shard-id:
+    description: "Zero-based pytest shard index"
+    required: false
+    default: ""
+  num-shards:
+    description: "Total number of pytest shards"
+    required: false
+    default: ""
 
 runs:
   using: "composite"
@@ -141,6 +149,8 @@ runs:
           --env NEMO_HOME=$NEMO_HOME \
           --env RUN_ID=${{ github.run_id }} \
           --env UV_NO_SYNC=1 \
+          --env SHARD_ID=${{ inputs.shard-id }} \
+          --env NUM_SHARDS=${{ inputs.num-shards }} \
           --workdir /opt/Megatron-Bridge \
           $VOLUME_ARGS \
           ${{ inputs.container-image }} \

.github/scripts/test_docker_dependency_layers.sh

@@ -236,7 +236,12 @@ if ! grep -A1 '^te = \[$' pyproject.toml | grep -Fxq '    "megatron-core[te]",';
   echo "Bridge's TE extra must enable the selected MCore ref's TE extra" >&2
   exit 1
 fi
-grep -Fq '{ name = "megatron-core", extra = ["te"] }' uv.lock
+if ! grep -Fq \
+  '{ name = "megatron-core", extras = ["te"], marker = "extra == '\''te'\''", editable = "3rdparty/Megatron-LM" }' \
+  uv.lock; then
+  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

.github/workflows/cicd-main.yml

@@ -567,6 +567,10 @@ jobs:
               --skip-file /opt/import-check/import_check_skip.txt
 
   cicd-unit-tests-core:
+    strategy:
+      fail-fast: false
+      matrix:
+        shard-id: [0, 1, 2]
     if: |
       (
         success()
@@ -579,7 +583,7 @@ jobs:
       && (github.event.inputs.test_suite == '' || github.event.inputs.test_suite == 'all' || github.event.inputs.test_suite == 'unit-only' || contains('L0 L1 L2', github.event.inputs.test_suite))
     needs: [pre-flight, configure, cicd-container-build]
     runs-on: ${{ needs.pre-flight.outputs.runner_prefix }}
-    name: Launch_Unit_Tests_Core
+    name: Launch_Unit_Tests_Core_${{ matrix.shard-id }}
     env:
       TRANSFORMERS_OFFLINE: "1"
       HF_HUB_OFFLINE: "1"
@@ -603,6 +607,8 @@ jobs:
           script: Launch_Unit_Tests_Core
           timeout: 30
           is_unit_test: "true"
+          shard-id: ${{ matrix.shard-id }}
+          num-shards: 3
           github-token: ${{ steps.app-token.outputs.token }}
           container-image: ${{ needs.pre-flight.outputs.registry }}/megatron-bridge:${{ needs.configure.outputs.merge_sha }}
           test-data-path: ${{ needs.pre-flight.outputs.test_data_path }}

.github/workflows/claude-review.yml

@@ -24,13 +24,15 @@ jobs:
   auto-review:
     name: Claude Review (auto)
     needs: [pre-flight]
-    # Auto-review non-draft PRs from NVIDIA authors on open / reopen / ready-for-review.
+    # Auto-review same-repository, non-draft PRs from NVIDIA authors on open / reopen / ready-for-review.
+    # Fork pull_request workflows cannot access the secrets or OIDC token required by the review action.
     # Subsequent commits do not re-trigger; re-request a review by commenting
     # "/claude review" (handled by the manual-review job below).
     # Skip bot-generated PRs: dependency bumps (`bump-ci-container-*`) and
     # auto cherry-picks (`cherry-pick-*`).
     if: |
       github.event_name == 'pull_request' &&
+      github.event.pull_request.head.repo.full_name == github.repository &&
       needs.pre-flight.outputs.is_member == 'true' &&
       github.event.pull_request.draft == false &&
       !startsWith(github.event.pull_request.head.ref, 'bump-ci-container-') &&

docs/models/deepseek/deepseek-v4.md

@@ -4,6 +4,30 @@
 
 DeepSeek V4 models are supported via the Bridge system with auto-detected configuration and weight mapping.
 
+## Chat SFT preprocessing
+
+`DirectHFSFTDatasetConfig` automatically selects the DeepSeek V4 chat collator
+from the tokenizer's `name_or_path`. Because the official V4 tokenizer does not
+provide a Jinja chat template, each training row must select the rendering mode
+explicitly with either `enable_thinking` or `thinking_mode`:
+
+```json
+{
+  "messages": [
+    {"role": "user", "content": "Solve 1 + 1."},
+    {"role": "assistant", "reasoning_content": "Add the terms.", "content": "2"}
+  ],
+  "enable_thinking": true,
+  "truncate_history_thinking": false
+}
+```
+
+Use `enable_thinking: false` for chat mode. Historical reasoning is truncated by
+default in thinking mode; set `truncate_history_thinking: false` to retain it.
+OpenAI-format tool definitions can be
+provided in the row's top-level `tools` field, and tool calls/results remain in
+the structured `messages` list.
+
 ## Model Architecture Features
 
 - **Hybrid Attention (DSv4HybridSelfAttention)**: Per-layer mix of dense MLA and Compressed Sparse Attention selected by `compress_ratios`

pyproject.toml

@@ -201,6 +201,7 @@ test = [
     "pytest-timeout>=2.4.0",
     "pytest-mock>=3.14.0",
     "pytest-runner>=6.0.1",
+    "pytest-shard>=0.1.2",
     "pygithub",
     "click",
 ]

scripts/inference/async_text_generation.py

@@ -116,6 +116,10 @@ async def _generate(
     ) as llm:
         if llm.is_primary_rank:
             results = await asyncio.gather(*(llm.generate(prompt, sampling_params) for prompt in prompts))
+            failed_results = [result for result in results if result.failed()]
+            if failed_results:
+                details = ", ".join(f"request {result.request_id}={result.status.name}" for result in failed_results)
+                raise RuntimeError(f"Async inference failed: {details}")
             print_rank_0("======== ASYNC GENERATED TEXT OUTPUT ========")
             for idx, result in enumerate(results):
                 print_rank_0(f"[{idx}] Prompt: {prompts[idx]}")

src/megatron/bridge/data/builders/direct_hf_sft.py

@@ -51,6 +51,9 @@
 logger = logging.getLogger(__name__)
 
 CollateFunction = Callable[..., dict[str, torch.Tensor]]
+_DEEPSEEK_V4_TOKENIZER_SIZE = 129280
+_DEEPSEEK_V4_ASSISTANT_TOKEN = "<｜Assistant｜>"
+_DEEPSEEK_V4_ASSISTANT_TOKEN_ID = 128804
 
 
 def _as_source(entry: "HFDatasetSourceConfig | Mapping[str, Any]") -> HFDatasetSourceConfig:
@@ -257,6 +260,26 @@ def select_direct_hf_sft_collate(
     raise ValueError("Prompt-completion preprocessing supports text-only examples.")
 
 
+def _model_collate_key(processor: Any) -> str:
+    tokenizer = get_processor_tokenizer(processor)
+    name_or_path = getattr(tokenizer, "name_or_path", "")
+    convert_tokens_to_ids = getattr(tokenizer, "convert_tokens_to_ids", None)
+    try:
+        has_deepseek_v4_fingerprint = len(tokenizer) == _DEEPSEEK_V4_TOKENIZER_SIZE and (
+            callable(convert_tokens_to_ids)
+            and convert_tokens_to_ids(_DEEPSEEK_V4_ASSISTANT_TOKEN) == _DEEPSEEK_V4_ASSISTANT_TOKEN_ID
+        )
+    except TypeError:
+        has_deepseek_v4_fingerprint = False
+    if (
+        isinstance(name_or_path, str)
+        and "deepseek-v4" in name_or_path.lower().replace("_", "-")
+        and has_deepseek_v4_fingerprint
+    ):
+        return "deepseek-v4"
+    return type(processor).__name__
+
+
 def load_direct_hf_sft_train_examples(config: DirectHFSFTDatasetConfig) -> list[dict[str, Any]]:
     """Load the training rows, blending them when the config names several sources."""
     sources = config.training_sources
@@ -278,17 +301,24 @@ def build_direct_hf_sft_split(
     """Build one requested direct-HF SFT split."""
     if target_length <= 0:
         return None
-    from megatron.bridge.data.collators.registry import model_collate_required_for_all_examples
+    from megatron.bridge.data.collators.registry import always_use_model_collate, resolve_model_collate
 
     if examples is None:
         examples = load_direct_hf_sft_examples(source, config.preprocessing)
-    if collate_impl is None and model_collate_required_for_all_examples(type(processor).__name__):
+    collate_key = _model_collate_key(processor)
+    if collate_impl is None and always_use_model_collate(collate_key):
         if not isinstance(config.preprocessing, ChatSFTPreprocessingConfig):
             raise ValueError(
                 f"Processor type '{type(processor).__name__}' requires chat preprocessing through its "
                 "model-owned collator."
             )
-        selected_collate = None
+        if collate_key == "deepseek-v4":
+            selected_collate = partial(
+                resolve_model_collate(collate_key),
+                loss_mode=config.preprocessing.loss_mode,
+            )
+        else:
+            selected_collate = None
     else:
         selected_collate = select_direct_hf_sft_collate(examples, config.preprocessing, collate_impl)
     return DirectSFTDataset(

src/megatron/bridge/data/collators/registry.py

@@ -25,7 +25,7 @@
 class _ModelCollateSpec:
     module_name: str
     symbol_name: str
-    required_for_all_examples: bool = False
+    always_use_model_collate: bool = False
 
 
 _MODEL_COLLATE_SPECS = {
@@ -42,7 +42,7 @@ class _ModelCollateSpec:
     "NemotronH_Nano_Omni_Reasoning_V3Processor": _ModelCollateSpec(
         "megatron.bridge.models.nemotron_omni.data.collate_fn",
         "nemotron_omni_expanded_collate_fn",
-        required_for_all_examples=True,
+        always_use_model_collate=True,
     ),
     "PixtralProcessor": _ModelCollateSpec(
         "megatron.bridge.models.ministral3.data.collate_fn", "ministral3_collate_fn"
@@ -54,13 +54,18 @@ class _ModelCollateSpec:
     ),
     "Glm4vProcessor": _ModelCollateSpec("megatron.bridge.models.glm_vl.data.collate_fn", "glm4v_collate_fn"),
     "KimiK25Processor": _ModelCollateSpec("megatron.bridge.models.kimi_vl.data.collate_fn", "kimi_k25_vl_collate_fn"),
+    "deepseek-v4": _ModelCollateSpec(
+        "megatron.bridge.models.deepseek.data.collate_fn",
+        "deepseek_v4_collate_fn",
+        always_use_model_collate=True,
+    ),
 }
 
 
-def model_collate_required_for_all_examples(processor_type: str) -> bool:
+def always_use_model_collate(processor_type: str) -> bool:
     """Return whether a processor must always use its model-owned collator."""
     spec = _MODEL_COLLATE_SPECS.get(processor_type)
-    return spec is not None and spec.required_for_all_examples
+    return spec is not None and spec.always_use_model_collate
 
 
 @lru_cache(maxsize=None)

src/megatron/bridge/data/collators/sft.py

@@ -16,13 +16,14 @@
 
 from __future__ import annotations
 
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping
 from dataclasses import dataclass
 from typing import Any, Literal
 
 import torch
 
 from megatron.bridge.data.conversation_processing import (
+    TokenizedConversation,
     build_shifted_labels_and_loss_mask,
     ensure_position_ids,
     get_processor_tokenizer,
@@ -180,6 +181,7 @@ def text_chat_collate_fn(
     ignore_index: int = IGNORE_INDEX,
     enable_in_batch_packing: bool = False,
     in_batch_packing_pad_to_multiple_of: int = 1,
+    tokenize_impl: Callable[..., TokenizedConversation] | None = None,
     **kwargs: Any,
 ) -> dict[str, Any]:
     """Collate text-only HF chat examples using the shared assistant-mask path.
@@ -204,6 +206,9 @@ def text_chat_collate_fn(
             packed-sequence metadata for GPT-style training steps.
         in_batch_packing_pad_to_multiple_of: Optional per-sequence length multiple
             used when ``enable_in_batch_packing`` inserts padding for CP/SP constraints.
+        tokenize_impl: Optional model-owned chat tokenizer. Models without a
+            Jinja chat template may provide an equivalent renderer while reusing
+            the shared padding, shifting, and packing implementation.
         **kwargs: Additional common collate kwargs accepted for parity with
             VLM collate functions and ignored by the text-only path.
 
@@ -218,9 +223,10 @@ def text_chat_collate_fn(
     tokenizer = get_processor_tokenizer(processor)
     boundary_config = infer_assistant_mask_boundary_config(processor)
     skipped_tokens = extract_skipped_token_ids(processor)
+    tokenize = tokenize_impl or tokenize_chat_example
     tokenized_rows = []
     for example in examples:
-        tokenized = tokenize_chat_example(
+        tokenized = tokenize(
             example,
             processor,
             max_length=max_length,