fix(ci): use writable unit test caches

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

Author: svcnemo-autobot

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

@@ -114,8 +114,12 @@ 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
 
         cmd=$(cat <<RUN_TEST_EOF
@@ -131,8 +135,8 @@ 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 NEMO_HOME=$NEMO_HOME \
           --env RUN_ID=${{ github.run_id }} \
           --env UV_NO_SYNC=1 \
           --workdir /opt/Megatron-Bridge \

.github/scripts/test_docker_dependency_layers.sh

@@ -243,6 +243,12 @@ 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"
 
 # The baseline dependency layer must be structurally independent of the mutable
 # dispatched checkout. CI validates the ordering statically so this regression

src/megatron/bridge/training/gpt_step.py

@@ -131,7 +131,10 @@ def _patch_mcore_schedule_plan_padding_mask() -> None:
     TODO: Remove this compatibility patch after MCore dev's schedule-plan
     padding-mask forwarding reaches the pinned main commit.
     """
-    from megatron.core.models.common import fine_grained_callables
+    try:
+        from megatron.core.models.common import fine_grained_callables
+    except ImportError:
+        from megatron.core.models.gpt import fine_grained_callables
 
     current_builder = fine_grained_callables.build_transformer_layer_callables
     if getattr(current_builder, _MCORE_SCHEDULE_PADDING_MASK_PATCHED, False):