Merge remote-tracking branch 'upstream/main' into ci/shard-core-unit-tests

Author: svcnemo-autobot

.dev.commit

@@ -1 +1 @@
-cae55c1f0991abd3e391d0029b15bcb35723d3f6
+12eed0a5ac3f1ced6a3a4f76727db530bdbf7391

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

@@ -122,9 +122,14 @@ runs:
         echo "::group::Start test container"
         if [[ "$MOUNT_FS" == "true" ]]; then
           VOLUME_ARGS="--volume ${{ inputs.test-data-path }}:/home/TestData --env HF_HUB_OFFLINE=1"
+          HF_HOME=/home/TestData/HF_HOME
+          NEMO_HOME=/home/TestData/nemo_home
         else
           VOLUME_ARGS=""
+          HF_HOME=/home/ubuntu/.cache/huggingface
+          NEMO_HOME=/home/ubuntu/.cache/nemo
         fi
+        HF_MODULES_CACHE=/home/ubuntu/.cache/huggingface/modules
 
         cmd=$(cat <<RUN_TEST_EOF
         #!/bin/bash
@@ -139,8 +144,9 @@ runs:
           --cpus=40 \
           --env GHA_RUNNER=${{ inputs.runner }} \
           --env HYDRA_FULL_ERROR=1 \
-          --env HF_HOME=/home/TestData/HF_HOME \
-          --env NEMO_HOME=/home/TestData/nemo_home \
+          --env HF_HOME=$HF_HOME \
+          --env HF_MODULES_CACHE=$HF_MODULES_CACHE \
+          --env NEMO_HOME=$NEMO_HOME \
           --env RUN_ID=${{ github.run_id }} \
           --env UV_NO_SYNC=1 \
           --env SHARD_ID=${{ inputs.shard-id }} \

.github/scripts/test_docker_dependency_layers.sh

@@ -31,7 +31,7 @@ if grep -q -- '--mount=type=secret,id=GH_TOKEN' "$dockerfile"; then
 fi
 mcore_reinstall_line=$(grep -n 'uv pip install --no-deps --reinstall -e 3rdparty/Megatron-LM' "$dockerfile" | cut -d: -f1)
 helper_assertion_line=$(grep -n "find 3rdparty/Megatron-LM/megatron/core/datasets -maxdepth 1 -name 'helpers_cpp\*\.so'" "$dockerfile" | cut -d: -f1)
-final_copy_line=$(grep -n '^COPY --chmod=644 \. /opt/Megatron-Bridge$' "$dockerfile" | cut -d: -f1)
+final_copy_line=$(grep -nE '^COPY (--chmod=644|--chown=1000:1000) \. /opt/Megatron-Bridge$' "$dockerfile" | cut -d: -f1)
 [[ -n "$mcore_reinstall_line" ]]
 [[ -n "$helper_assertion_line" ]]
 [[ -n "$final_copy_line" ]]
@@ -243,6 +243,14 @@ if grep -qE 'mcore_(commit|ref)|MCORE_COMMIT|uv sync --all-extras --all-groups'
   echo "Test template must use the already-validated MCore in the built image" >&2
   exit 1
 fi
+grep -Fq 'VOLUME_ARGS="--volume ${{ inputs.test-data-path }}:/home/TestData --env HF_HUB_OFFLINE=1"' "$composite_action"
+grep -q 'MOUNT_FS: ${{ inputs.is_unit_test == '\''false'\'' }}' "$composite_action"
+grep -q 'HF_HOME=/home/TestData/HF_HOME' "$composite_action"
+grep -q 'NEMO_HOME=/home/TestData/nemo_home' "$composite_action"
+grep -q 'HF_HOME=/home/ubuntu/.cache/huggingface' "$composite_action"
+grep -q 'NEMO_HOME=/home/ubuntu/.cache/nemo' "$composite_action"
+grep -q 'HF_MODULES_CACHE=/home/ubuntu/.cache/huggingface/modules' "$composite_action"
+grep -q -- '--env HF_MODULES_CACHE=\$HF_MODULES_CACHE' "$composite_action"
 
 # The baseline dependency layer must be structurally independent of the mutable
 # dispatched checkout. CI validates the ordering statically so this regression

.github/workflows/cicd-main.yml

@@ -605,7 +605,7 @@ jobs:
         uses: ./.github/actions/test-template
         with:
           script: Launch_Unit_Tests_Core
-          timeout: 18
+          timeout: 30
           is_unit_test: "true"
           shard-id: ${{ matrix.shard-id }}
           num-shards: 3

.main.commit

@@ -1 +1 @@
-24bad8e677d22625d86ef2a54c9506b6e4992c93
+14346b65a2d0790e451919858f7771078105c5f0

3rdparty/Megatron-LM

@@ -102,7 +102,8 @@ RUN --mount=type=bind,source=docker/patches/deepep.patch,target=/opt/deepep.patc
         apt-get autoremove -y && \
         rm -rf /var/lib/apt/lists/* && \
     popd; \
-    fi
+    fi && \
+    chown -R 1000:1000 ${UV_PROJECT_ENVIRONMENT}
 
 COPY pyproject.toml uv.lock /opt/Megatron-Bridge/
 COPY src/megatron/bridge/__init__.py src/megatron/bridge/package_info.py /opt/Megatron-Bridge/src/megatron/bridge/
@@ -171,7 +172,8 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     # Create symlink for tilelang's libcudart_stub.so to the actual libcudart.so
     # Otherwise, the stub will be called in some cases and fail
     ln -sf "$(ldconfig -p | awk '/libcudart\.so\.[0-9]+ /{print $NF; exit}')" /opt/venv/lib/python3.12/site-packages/tilelang/lib/libcudart_stub.so && \
-    uv cache prune ${UV_CACHE_PRUNE_ARGS}
+    uv cache prune ${UV_CACHE_PRUNE_ARGS} && \
+    chown -R 1000:1000 ${UV_PROJECT_ENVIRONMENT}
 
 COPY scripts/install_diffusion_deps.sh scripts/diffusion-deps.lock /opt/Megatron-Bridge/scripts/
 
@@ -181,6 +183,16 @@ ARG INSTALL_DIFFUSION_DEPS=false
 RUN --mount=type=cache,target=/root/.cache/uv \
     if [ "$INSTALL_DIFFUSION_DEPS" = "true" ]; then \
         bash scripts/install_diffusion_deps.sh; \
-    fi
+    fi && \
+    chown -R 1000:1000 ${UV_PROJECT_ENVIRONMENT}
+
+COPY --chown=1000:1000 . /opt/Megatron-Bridge
+
+# CI writes coverage beside the checked-out source, installs that source editable into the
+# 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
 
-COPY --chmod=644 . /opt/Megatron-Bridge
+ENV HOME=/home/ubuntu
+USER 1000:1000

docker/Dockerfile.ci

@@ -294,7 +294,7 @@ def megatron_generate(
         torch.distributed.broadcast_object_list(
             generated_logits,
             group=parallel_state.get_pipeline_model_parallel_group(),
-            group_src=parallel_state.get_pipeline_model_parallel_last_rank(),
+            src=parallel_state.get_pipeline_model_parallel_last_rank(),
         )
 
     generated_text = tokenizer.detokenize(generated_ids[0])

examples/conversion/compare_text_generation.py

@@ -504,71 +504,75 @@ def _hf_layer_prefix(self) -> str:
         """Text-only CausalLM: weights at ``model.*``; override in VL subclass."""
         return "model."
 
-    def _moe_mapping_registry(self) -> MegatronMappingRegistry:
+    def _moe_mapping_registry(self, megatron_prefix: str = "") -> MegatronMappingRegistry:
         """Parameter mappings for the MoE variant."""
+        mp = megatron_prefix
+        hp = self._hf_layer_prefix()
         param_mappings = {
-            "embedding.word_embeddings.weight": "model.embed_tokens.weight",
-            "decoder.final_layernorm.weight": "model.norm.weight",
-            "decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": "model.layers.*.input_layernorm.weight",
-            "decoder.layers.*.self_attention.q_layernorm.weight": "model.layers.*.self_attn.q_norm.weight",
-            "decoder.layers.*.self_attention.k_layernorm.weight": "model.layers.*.self_attn.k_norm.weight",
-            "decoder.layers.*.self_attention.linear_proj.weight": "model.layers.*.self_attn.o_proj.weight",
-            "decoder.layers.*.self_attention.linear_proj.post_layernorm.weight": (
-                "model.layers.*.post_attention_layernorm.weight"
+            f"{mp}embedding.word_embeddings.weight": f"{hp}embed_tokens.weight",
+            f"{mp}decoder.final_layernorm.weight": f"{hp}norm.weight",
+            f"{mp}decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": (
+                f"{hp}layers.*.input_layernorm.weight"
+            ),
+            f"{mp}decoder.layers.*.self_attention.q_layernorm.weight": f"{hp}layers.*.self_attn.q_norm.weight",
+            f"{mp}decoder.layers.*.self_attention.k_layernorm.weight": f"{hp}layers.*.self_attn.k_norm.weight",
+            f"{mp}decoder.layers.*.self_attention.linear_proj.weight": f"{hp}layers.*.self_attn.o_proj.weight",
+            f"{mp}decoder.layers.*.self_attention.linear_proj.post_layernorm.weight": (
+                f"{hp}layers.*.post_attention_layernorm.weight"
             ),
-            "decoder.layers.*.pre_mlp_layernorm.weight": "model.layers.*.pre_feedforward_layernorm_2.weight",
-            "decoder.layers.*.mlp.shared_experts.linear_fc2.weight": "model.layers.*.mlp.down_proj.weight",
-            "decoder.layers.*.mlp.post_shared_expert_layernorm.weight": (
-                "model.layers.*.post_feedforward_layernorm_1.weight"
+            f"{mp}decoder.layers.*.pre_mlp_layernorm.weight": f"{hp}layers.*.pre_feedforward_layernorm_2.weight",
+            f"{mp}decoder.layers.*.mlp.shared_experts.linear_fc2.weight": f"{hp}layers.*.mlp.down_proj.weight",
+            f"{mp}decoder.layers.*.mlp.post_shared_expert_layernorm.weight": (
+                f"{hp}layers.*.post_feedforward_layernorm_1.weight"
             ),
-            "decoder.layers.*.mlp.router.weight": "model.layers.*.router.proj.weight",
+            f"{mp}decoder.layers.*.mlp.router.weight": f"{hp}layers.*.router.proj.weight",
         }
 
         mapping_list = [AutoMapping(megatron_param=m, hf_param=h) for m, h in param_mappings.items()]
         mapping_list.extend(
             [
                 _Gemma4QKVMapping(
-                    megatron_param="decoder.layers.*.self_attention.linear_qkv.weight",
-                    q="model.layers.*.self_attn.q_proj.weight",
-                    k="model.layers.*.self_attn.k_proj.weight",
-                    v="model.layers.*.self_attn.v_proj.weight",
+                    megatron_param=f"{mp}decoder.layers.*.self_attention.linear_qkv.weight",
+                    q=f"{hp}layers.*.self_attn.q_proj.weight",
+                    k=f"{hp}layers.*.self_attn.k_proj.weight",
+                    v=f"{hp}layers.*.self_attn.v_proj.weight",
                 ),
                 GatedMLPMapping(
-                    megatron_param="decoder.layers.*.mlp.shared_experts.linear_fc1.weight",
-                    gate="model.layers.*.mlp.gate_proj.weight",
-                    up="model.layers.*.mlp.up_proj.weight",
+                    megatron_param=f"{mp}decoder.layers.*.mlp.shared_experts.linear_fc1.weight",
+                    gate=f"{hp}layers.*.mlp.gate_proj.weight",
+                    up=f"{hp}layers.*.mlp.up_proj.weight",
                 ),
                 FusedGatedExpertMapping(
-                    megatron_param="decoder.layers.*.mlp.experts.linear_fc1.weight*",
-                    hf_param="model.layers.*.experts.gate_up_proj",
+                    megatron_param=f"{mp}decoder.layers.*.mlp.experts.linear_fc1.weight*",
+                    hf_param=f"{hp}layers.*.experts.gate_up_proj",
                 ),
                 FusedExpertMapping(
-                    megatron_param="decoder.layers.*.mlp.experts.linear_fc2.weight*",
-                    hf_param="model.layers.*.experts.down_proj",
+                    megatron_param=f"{mp}decoder.layers.*.mlp.experts.linear_fc2.weight*",
+                    hf_param=f"{hp}layers.*.experts.down_proj",
                 ),
                 ReplicatedMapping(
-                    megatron_param="decoder.layers.*.layer_scalar",
-                    hf_param="model.layers.*.layer_scalar",
+                    megatron_param=f"{mp}decoder.layers.*.layer_scalar",
+                    hf_param=f"{hp}layers.*.layer_scalar",
                 ),
                 ReplicatedMapping(
-                    megatron_param="decoder.layers.*.mlp.router.per_expert_scale",
-                    hf_param="model.layers.*.router.per_expert_scale",
+                    megatron_param=f"{mp}decoder.layers.*.mlp.router.per_expert_scale",
+                    hf_param=f"{hp}layers.*.router.per_expert_scale",
                 ),
                 ReplicatedMapping(
-                    megatron_param="decoder.layers.*.mlp.router.scale",
-                    hf_param="model.layers.*.router.scale",
+                    megatron_param=f"{mp}decoder.layers.*.mlp.router.scale",
+                    hf_param=f"{hp}layers.*.router.scale",
                 ),
                 ReplicatedMapping(
-                    megatron_param="decoder.layers.*.pre_shared_expert_layernorm.weight",
-                    hf_param="model.layers.*.pre_feedforward_layernorm.weight",
+                    megatron_param=f"{mp}decoder.layers.*.pre_shared_expert_layernorm.weight",
+                    hf_param=f"{hp}layers.*.pre_feedforward_layernorm.weight",
                 ),
                 ReplicatedMapping(
-                    megatron_param="decoder.layers.*.mlp.post_moe_layernorm.weight",
-                    hf_param="model.layers.*.post_feedforward_layernorm_2.weight",
+                    megatron_param=f"{mp}decoder.layers.*.mlp.post_moe_layernorm.weight",
+                    hf_param=f"{hp}layers.*.post_feedforward_layernorm_2.weight",
                 ),
                 ReplicatedMapping(
-                    megatron_param="decoder.layers.*.post_ffn_layernorm.weight",
-                    hf_param="model.layers.*.post_feedforward_layernorm.weight",
+                    megatron_param=f"{mp}decoder.layers.*.post_ffn_layernorm.weight",
+                    hf_param=f"{hp}layers.*.post_feedforward_layernorm.weight",
                 ),
             ]
         )

src/megatron/bridge/models/gemma/gemma4_bridge.py

@@ -21,10 +21,14 @@
 
   AutoBridge.from_hf_pretrained("google/gemma-4-E4B-it")
     └─ Gemma4VLBridge (registered for Gemma4ForConditionalGeneration)
-         ├─ provider_bridge()  text mode → Gemma4DenseProvider (pretraining)
-         │                     auto/vl   → Gemma4DenseVLProvider (full VL)
-         └─ mapping_registry()  Dense → _dense_vl_mapping_registry()
-                                 MoE   → _moe_vl_mapping_registry()
+         ├─ provider_bridge()  Dense: text mode → Gemma4DenseProvider (pretraining)
+         │                            auto/vl   → Gemma4DenseVLProvider (full VL)
+         │                     MoE:   text mode → Gemma4ModelProvider (text-only training)
+         │                            auto/vl   → Gemma4VLModelProvider (full VL)
+         └─ mapping_registry()  Dense: text → _dense_mapping_registry("")
+                                        vl   → _dense_vl_mapping_registry()
+                                 MoE:   text → _moe_mapping_registry()
+                                        vl   → _moe_vl_mapping_registry()
 """
 
 import os
@@ -91,6 +95,8 @@ def provider_bridge(
             return self._build_dense_vl_provider(hf_config, text_config, vision_config)
 
         self._is_dense = False
+        if self._conversion_mode() == "text":
+            return self._build_moe_provider(text_config)
 
         provider_kwargs = self.hf_config_to_provider_kwargs(text_config)
         provider = Gemma4VLModelProvider(**provider_kwargs)
@@ -226,6 +232,8 @@ def mapping_registry(self) -> MegatronMappingRegistry:
             if self._conversion_mode() == "text":
                 return self._dense_mapping_registry(megatron_prefix="")
             return self._dense_vl_mapping_registry()
+        if self._conversion_mode() == "text":
+            return self._moe_mapping_registry()
         return self._moe_vl_mapping_registry()
 
     def _dense_vl_mapping_registry(self) -> MegatronMappingRegistry: