merge: update branch to latest main

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

Author: yaoyu-33

docs/versions1.json

@@ -4,14 +4,15 @@
     "url": "https://docs.nvidia.com/nemo/megatron-bridge/nightly/"
   },
   {
+    "name": "0.6.0 (latest) · 26.08",
     "version": "0.6.0",
-    "url": "https://docs.nvidia.com/nemo/megatron-bridge/0.6.0/"
+    "url": "https://docs.nvidia.com/nemo/megatron-bridge/0.6.0/",
+    "preferred": true
   },
   {
     "name": "0.5.1 (latest) · 26.06.01",
     "version": "0.5.1",
-    "url": "https://docs.nvidia.com/nemo/megatron-bridge/0.5.1/",
-    "preferred": true
+    "url": "https://docs.nvidia.com/nemo/megatron-bridge/0.5.1/"
   },
   {
     "name": "0.5.0 · 26.06",

src/megatron/bridge/data/datasets/gpt_sft.py

@@ -408,7 +408,7 @@ def _multiple_truncation(self, template_ids: list[list[int]], template_ids_keys:
 
         if total_ids > self.max_seq_length:
             truncation_length_total = total_ids - self.max_seq_length
-            num_fields = len(self.truncation_fields)
+            num_fields = sum(key in self.truncation_fields for key in template_ids_keys)
             if num_fields > 0:
                 # sorted equal divide length to each field
                 # examples:

src/megatron/bridge/training/setup.py

@@ -596,6 +596,7 @@ def _build_distributed_model(cfg: ConfigContainer, pg_collection: ProcessGroupCo
             data_parallel_random_init=cfg.rng.data_parallel_random_init,
         )
     else:
+        model_config.finalize()
         return model_config.provide_distributed_model(
             ddp_config=cfg.ddp,
             use_megatron_fsdp=cfg.dist.use_megatron_fsdp,

tests/unit_tests/data/datasets/test_gpt_sft.py

@@ -219,6 +219,25 @@ def test_multiple_truncation(self, tmp_path):
         assert context_ids == [101, 102, 103, 104, 201, 202, 203]
         assert label_ids == [301, 302]
 
+    def test_repeated_truncation_field_placeholder_handles_overflow(self, tmp_path):
+        dataset_path = tmp_path / "repeated_prompt.jsonl"
+        dataset_path.write_text(json.dumps({"input": "one two three four five", "output": "answer"}) + "\n")
+        dataset = GPTSFTDataset(
+            file_path=str(dataset_path),
+            tokenizer=create_mock_tokenizer(),
+            max_seq_length=8,
+            max_num_samples=None,
+            label_key="output",
+            prompt_template="{input} Again: {input} {output}",
+            truncation_field="input",
+            memmap_workers=1,
+        )
+
+        processed = dataset[0]
+
+        assert len(processed["input_ids"]) <= dataset.max_seq_length
+        assert processed["answer_ids"]
+
     def test_utils_func(self, tmp_path):
         dataset, _ = get_gpt_sft(tmp_path)
 

tests/unit_tests/training/test_setup.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 
+import json
 from types import SimpleNamespace
 from unittest.mock import MagicMock, Mock, patch
 
@@ -25,6 +26,7 @@
 from megatron.bridge.models.gpt_provider import GPTModelProvider
 from megatron.bridge.models.hybrid.hybrid_builder import HybridModelConfig
 from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider
+from megatron.bridge.models.llama_nemotron.llama_nemotron_provider import LlamaNemotronHeterogeneousProvider
 from megatron.bridge.models.transformer_config import TransformerConfig
 from megatron.bridge.training.callbacks import CallbackManager
 from megatron.bridge.training.checkpointing import load_checkpoint
@@ -556,9 +558,42 @@ def test_build_with_provider(self):
 
         result = _build_distributed_model(cfg, pg_collection=MagicMock())
 
+        mock_provider.finalize.assert_called_once_with()
         mock_provider.provide_distributed_model.assert_called_once()
         assert result == mock_dist_model
 
+    def test_finalizes_heterogeneous_provider_before_entering_distributed_model_build(self):
+        """Finalize deferred heterogeneous fields before entering provider model construction."""
+        block = {
+            "attention": {"no_op": False, "replace_with_linear": False, "num_query_groups": 4},
+            "mlp": {"no_op": False, "replace_with_linear": False, "ffn_hidden_size": 256},
+        }
+        provider = LlamaNemotronHeterogeneousProvider(
+            num_layers=1,
+            hidden_size=64,
+            num_attention_heads=4,
+            heterogeneous_layers_config_encoded_json=json.dumps({"block_configs": [block]}),
+        )
+        expected_model = [MagicMock()]
+
+        def provide_distributed_model(**_kwargs):
+            assert len(provider.per_block_parameters) == 1
+            return expected_model
+
+        provider.provide_distributed_model = Mock(side_effect=provide_distributed_model)
+        cfg = SimpleNamespace(
+            model=provider,
+            ddp=MagicMock(),
+            optimizer=SimpleNamespace(overlap_param_gather_with_optimizer_step=False),
+            dist=SimpleNamespace(use_megatron_fsdp=False, use_torch_fsdp2=False),
+            rng=SimpleNamespace(data_parallel_random_init=False),
+        )
+
+        result = _build_distributed_model(cfg, pg_collection=MagicMock())
+
+        provider.provide_distributed_model.assert_called_once()
+        assert result == expected_model
+
 
 def test_restart_rebinds_overlap_callbacks_to_rebuilt_model():
     """A restart must replace callbacks bound to the discarded model."""