Merge branch 'main' into super-v3-pr

Author: liding-nv

.github/workflows/_update_dependencies.yml

@@ -115,7 +115,7 @@ jobs:
           commit-message: ${{ env.title }}
           signoff: true
           committer: "github-actions[bot] <github-actions[bot]@users.noreply.github.com>"
-          labels: needs-more-tests
+          labels: full-test-suite
 
       - name: Post /ok to test comment
         env:

.github/workflows/cicd-main.yml

@@ -298,6 +298,7 @@ jobs:
     needs: [pre-flight, cicd-container-build]
     outputs:
       needs_more_tests: ${{ steps.check.outputs.needs_more_tests }}
+      full_test_suite: ${{ steps.check.outputs.full_test_suite }}
     steps:
       - name: Check PR labels
         id: check
@@ -312,8 +313,14 @@ jobs:
             else
               echo "needs_more_tests=false" >> $GITHUB_OUTPUT
             fi
+            if echo "$LABELS" | grep -q "full-test-suite"; then
+              echo "full_test_suite=true" >> $GITHUB_OUTPUT
+            else
+              echo "full_test_suite=false" >> $GITHUB_OUTPUT
+            fi
           else
             echo "needs_more_tests=false" >> $GITHUB_OUTPUT
+            echo "full_test_suite=false" >> $GITHUB_OUTPUT
           fi
 
   cicd-unit-tests-core:
@@ -558,7 +565,7 @@ jobs:
         || needs.pre-flight.outputs.force_run_all == 'true'
       )
       && !cancelled()
-      && (github.ref == 'refs/heads/main' || github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' || needs.check-labels.outputs.needs_more_tests == 'true')
+      && (github.ref == 'refs/heads/main' || github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' || needs.check-labels.outputs.needs_more_tests == 'true' || needs.check-labels.outputs.full_test_suite == 'true')
       && (github.event.inputs.test_suite == '' || github.event.inputs.test_suite == 'all' || github.event.inputs.test_suite == 'functional-only')
     name: ${{ matrix.script }}
     environment: nemo-ci
@@ -617,7 +624,7 @@ jobs:
       fail-fast: false
       max-parallel: 16
       matrix: ${{ fromJSON(needs.generate-test-matrix.outputs.matrix_l2) }}
-    needs: [pre-flight, generate-test-matrix, cicd-unit-tests-core, cicd-unit-tests-diffusion]
+    needs: [pre-flight, generate-test-matrix, cicd-unit-tests-core, cicd-unit-tests-diffusion, check-labels]
     runs-on: ${{ needs.pre-flight.outputs.runner_prefix }}-gpu-x2
     if: |
       (
@@ -626,7 +633,7 @@ jobs:
         || needs.pre-flight.outputs.force_run_all == 'true'
       )
       && !cancelled()
-      && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch')
+      && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' || needs.check-labels.outputs.full_test_suite == 'true')
       && (github.event.inputs.test_suite == '' || github.event.inputs.test_suite == 'all' || github.event.inputs.test_suite == 'functional-only')
     name: ${{ matrix.script }}
     environment: nemo-ci

docs/performance-summary.md

@@ -448,6 +448,7 @@ def provider_bridge(self, hf_pretrained: HFPreTrained) -> ModelProviderTarget:
             ModelProviderTarget: A configured model provider instance
         """
         from megatron.bridge.models.gpt_provider import GPTModelProvider
+        from megatron.bridge.models.mla_provider import MLAModelProvider
 
         hf_config = hf_pretrained.config
 
@@ -458,22 +459,25 @@ def provider_bridge(self, hf_pretrained: HFPreTrained) -> ModelProviderTarget:
 
         # Use specified provider class, defaulting to GPTModelProvider
         provider_class = self.PROVIDER_CLASS if self.PROVIDER_CLASS is not None else GPTModelProvider
+        is_mla_provider = issubclass(provider_class, MLAModelProvider)
         # Filter kwargs to only fields the provider dataclass accepts, so that MLA-only None
         # values (q_lora_rank, kv_lora_rank, …) are silently dropped for non-MLA providers
         # while still being passed through for MLA providers that declare them as fields.
         valid_fields = provider_class.__dataclass_fields__
         provider = provider_class(**{k: v for k, v in provider_kwargs.items() if k in valid_fields})
 
-        # All models that flow through the base provider_bridge use RoPE.
-        # YARN is already handled above (hf_config_to_provider_kwargs sets
-        # position_embedding_type="yarn"), so we only need to guard against
-        # overwriting it here.  Every other rope_type value (None, "default",
-        # "llama3", "longrope", …) means plain RoPE.
+        # Determine position_embedding_type from HF rope_scaling.
+        # For GPT providers: rope_type=="yarn" → "yarn"; everything else → "rope".
+        # For MLA providers: always "rope" — YaRN scaling parameters are applied
+        # separately via mla_rope_params (rotary_scaling_factor, mscale, etc.) and
+        # position_embedding_type="yarn" is not a valid MLA config value.
         hf_rope_scaling = getattr(hf_config, "rope_scaling", None)
         rope_type = None
         if hf_rope_scaling:
             rope_type = hf_rope_scaling.get("type") or hf_rope_scaling.get("rope_type")
-        if rope_type != "yarn":
+        if rope_type == "yarn" and not is_mla_provider:
+            provider.position_embedding_type = "yarn"
+        else:
             provider.position_embedding_type = "rope"
 
         # Apply MLA rope params via setattr (for MLA models like DeepSeek, Kimi)

src/megatron/bridge/models/conversion/model_bridge.py

@@ -62,6 +62,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider
         provider.moe_permute_fusion = True
 
         provider.apply_rope_fusion = False
+        provider.gradient_accumulation_fusion = True
         provider.bias_activation_fusion = True
         provider.bias_dropout_fusion = True
         provider.cross_entropy_fusion_impl = "te"

src/megatron/bridge/models/deepseek/deepseek_v2_bridge.py

@@ -69,6 +69,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider
         provider.moe_aux_loss_coeff = 0.0001
 
         provider.apply_rope_fusion = False
+        provider.gradient_accumulation_fusion = True
         provider.bias_activation_fusion = True
         provider.bias_dropout_fusion = True
         provider.cross_entropy_fusion_impl = "te"