[CICD] Improve CUDA unit test coverage

.github/configs/metax.yml

@@ -29,14 +29,14 @@ container_volumes:
 
 # Container options
 container_options: >-
-  --uts=host
+  --hostname=te_cicd
   --ipc=host
   --privileged=true
-  --group-add video
   --shm-size=100gb
   --ulimit memlock=-1
   --user root
   --ulimit nofile=65535:65535
+  --group-add video
   -e PLATFORM=metax
   -e TORCH_DISTRIBUTED_BACKEND=mccl
   -e LD_LIBRARY_PATH=/opt/maca/lib:/usr/local/lib:$LD_LIBRARY_PATH

.github/configs/musa.yml

@@ -0,0 +1,56 @@
+# Moore Threads (MUSA) Hardware Configuration for TE-FL
+# This file defines CI/CD settings for Moore Threads MUSA-based testing
+# This file defines environment variables, volumes, and test filters for TE tests.
+
+hardware_name: musa
+display_name: 'Moore Threads MUSA Tests'
+
+# CI image for Moore Threads env (Updated to the 0206 verified image)
+ci_image: registry.mthreads.com/presale/devtech/flagsclae:0206
+
+# Target precise runner labels for the 140 Moore Threads node
+runner_labels:
+  - mt-8g-cicd-te
+
+# Container volumes
+container_volumes:
+  - /data:/data
+
+# Container options for MUSA execution
+container_options: >-
+  --uts=host
+  --ipc=host
+  --privileged=true
+  --shm-size=32gb
+  --user root
+  --ulimit nofile=65535:65535
+  --device=/dev/mtgpu
+  -e PLATFORM=musa
+  -e ACCELERATOR_BACKEND=musa
+  -e MUSA_VISIBLE_DEVICES=0
+  -e TE_FL_PREFER=reference
+  -e LD_LIBRARY_PATH=/usr/local/musa/lib:/usr/local/lib:$LD_LIBRARY_PATH
+  -e http_proxy="http://flagcicd-yzli:A7fK9mQ2Zx4P@10.1.12.38:3128"
+  -e https_proxy="http://flagcicd-yzli:A7fK9mQ2Zx4P@10.1.12.38:3128"
+
+# Platform-specific environment setup script
+setup_script: .github/scripts/setup_musa.sh
+
+# Build environment variables & Proxy Settings for CI Runner
+build_env:
+  TE_FL_SKIP_CUDA: '1'
+  MUSA_HOME: /usr/local/musa
+  http_proxy: "http://flagcicd-yzli:A7fK9mQ2Zx4P@10.1.12.38:3128"
+  https_proxy: "http://flagcicd-yzli:A7fK9mQ2Zx4P@10.1.12.38:3128"
+
+# Device types to run tests on
+device_types:
+  - mt-gpu
+
+# Test matrix configuration
+test_matrix:
+  unit:
+    devices:
+      - mt-gpu
+    ignored_tests:
+      # - tests/pytorch/test_numerics.py

.github/workflows/all_tests_musa.yml

@@ -0,0 +1,34 @@
+name: musa_tests
+
+on:
+  push:
+    branches: ["main"]
+  pull_request:
+    branches: ["main"]
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-${{ github.actor }}
+  cancel-in-progress: true
+
+jobs:
+  run_tests:
+    # Package manager and environment settings are read from .github/configs/musa.yml
+    uses: ./.github/workflows/all_tests_common.yml
+    with:
+      platform: musa
+      run_unit_tests: true
+      run_integration_tests: true
+
+  all_tests:
+    needs: run_tests
+    runs-on: ubuntu-latest
+    if: always()
+    steps:
+      - name: Verify workflow status
+        run: |
+          if [ "${{ needs.run_tests.result }}" != "success" ]; then
+            echo "❌ MUSA Tests workflow failed"
+            exit 1
+          fi
+          echo "✅ All MUSA tests passed!"

.github/workflows/unit_tests_common.yml

@@ -51,6 +51,9 @@ jobs:
           - name: pytorch_distributed_unittest
             path: "qa/L1_pytorch_distributed_unittest/test.sh"
             test_type: "unittest"
+          - name: pytorch_onnx_unittest
+            path: "qa/L1_pytorch_onnx_unittest/test.sh"
+            test_type: "unittest"
     name: unit-${{ inputs.device }}-${{ matrix.test_group.name }}
     container:
       image: ${{ inputs.image }}
@@ -155,7 +158,7 @@ jobs:
           # Coverage setup: install once + configure collection via PYTEST_ADDOPTS
           COVERAGE_ENABLED=false
           if pip3 install coverage pytest-cov --quiet 2>/dev/null; then
-            export PYTEST_ADDOPTS="--cov=transformer_engine --cov-append --cov-report="
+            export PYTEST_ADDOPTS="--cov=transformer_engine/pytorch --cov=transformer_engine/debug --cov=transformer_engine/plugin --cov-append --cov-report="
             COVERAGE_ENABLED=true
           else
             echo "WARNING: Failed to install coverage/pytest-cov, coverage collection disabled"
@@ -175,27 +178,109 @@ jobs:
             python3 -m coverage combine --keep 2>/dev/null || true
             python3 -m coverage json \
               -o "coverage-${{ inputs.platform }}-${{ inputs.device }}-${{ matrix.test_group.name }}.json" \
-              --include="transformer_engine/*" 2>/dev/null \
+              --include="transformer_engine/pytorch/*,transformer_engine/debug/*,transformer_engine/plugin/*" \
+              --omit="*/setup.py,*/transformer_engine/plugin/core/_build_config.py" \
+              -i 2>/dev/null \
               || echo "WARNING: No coverage data found"
+
+            if [ -f .coverage ]; then
+              cp .coverage "coverage-${{ inputs.platform }}-${{ inputs.device }}-${{ matrix.test_group.name }}.coverage"
+            fi
           fi
           exit $exit_code
         timeout-minutes: 60
 
-      - name: Upload Coverage Report
+      - name: Upload Coverage Artifact
         uses: actions/upload-artifact@v4
         continue-on-error: true
         with:
           name: coverage-${{ inputs.platform }}-${{ inputs.device }}-${{ matrix.test_group.name }}
           path: |
             coverage-${{ inputs.platform }}-${{ inputs.device }}-${{ matrix.test_group.name }}.json
+            coverage-${{ inputs.platform }}-${{ inputs.device }}-${{ matrix.test_group.name }}.coverage
+
+  aggregate_coverage:
+    name: unit-${{ inputs.device }}-pytorch-unittest-coverage
+    needs: unit_test
+    if: always()
+    defaults:
+      run:
+        shell: bash
+    runs-on: ${{ fromJson(inputs.runs_on) }}
+    container:
+      image: ${{ inputs.image }}
+      volumes: ${{ fromJson(inputs.container_volumes) }}
+      options: --pull never ${{ inputs.container_options }}
+
+    steps:
+      - name: Checkout Source Code
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+          set-safe-directory: true
+
+      - name: Download Coverage Artifacts
+        uses: actions/download-artifact@v4
+        continue-on-error: true
+        with:
+          pattern: coverage-${{ inputs.platform }}-${{ inputs.device }}-pytorch_*
+          path: coverage-artifacts
+          merge-multiple: true
+
+      - name: Aggregate PyTorch Coverage
+        id: aggregate
+        run: |
+          set -euo pipefail
+
+          if ${{inputs.platform == 'metax'}}; then
+            source /opt/conda/etc/profile.d/conda.sh
+            conda activate base
+          else
+            source /opt/miniconda3/etc/profile.d/conda.sh
+            conda activate flagscale-train
+          fi
+
+          python3 -m pip install coverage --quiet
+
+          mkdir -p coverage-raw
+          coverage_count=0
+          for coverage_file in coverage-artifacts/*.coverage; do
+            [ -f "$coverage_file" ] || continue
+            coverage_count=$((coverage_count + 1))
+            cp "$coverage_file" "coverage-raw/.coverage.$coverage_count"
+          done
+
+          if [ "$coverage_count" -eq 0 ]; then
+            echo "WARNING: No raw coverage files found for aggregation"
+            echo "coverage_aggregated=false" >> "$GITHUB_OUTPUT"
+            exit 0
+          fi
+
+          python3 -m coverage combine coverage-raw
+          python3 -m coverage json \
+            -o "coverage-${{ inputs.platform }}-${{ inputs.device }}-pytorch-unittest.json" \
+            --include="transformer_engine/pytorch/*,transformer_engine/debug/*,transformer_engine/plugin/*" \
+            --omit="*/setup.py,*/transformer_engine/plugin/core/_build_config.py" \
+            -i
+
+          echo "coverage_aggregated=true" >> "$GITHUB_OUTPUT"
+
+      - name: Upload Aggregated Coverage Artifact
+        if: steps.aggregate.outputs.coverage_aggregated == 'true'
+        uses: actions/upload-artifact@v4
+        continue-on-error: true
+        with:
+          name: coverage-${{ inputs.platform }}-${{ inputs.device }}-pytorch-unittest
+          path: coverage-${{ inputs.platform }}-${{ inputs.device }}-pytorch-unittest.json
 
       - name: Upload Coverage Report to FlagCICD
+        if: steps.aggregate.outputs.coverage_aggregated == 'true'
         uses: flagos-ai/FlagOps/actions/post-pytest-report@v2
         continue-on-error: true
         env:
           NO_PROXY: "flagcicd-inner.flagos.net"
         with:
           backend_url: 'http://flagcicd-inner.flagos.net:8000/metrics/'
           user_id: '000000000000000000'
-          report_path: 'coverage-${{ inputs.platform }}-${{ inputs.device }}-${{ matrix.test_group.name }}.json'
-          fail_on_error: 'false'
+          report_path: 'coverage-${{ inputs.platform }}-${{ inputs.device }}-pytorch-unittest.json'
+          fail_on_error: 'false'

.gitignore

@@ -45,4 +45,7 @@ artifacts/
 transformer_engine/plugin/core/_build_config.py
 # Mac OS
 .DS_Store
-.claude/
+qa/L1_pytorch_mcore_integration/output/
+*.distcp
+.coverage
+.coverage.*

qa/L0_pytorch_debug_unittest/test.sh

@@ -78,7 +78,7 @@ run_test_step "test_log.xml" "$TE_PATH/tests/pytorch/debug/test_log.py" \
 
 # Step 5: API Features
 run_test_step "test_api_features.xml" "$TE_PATH/tests/pytorch/debug/test_api_features.py" \
-"NVTE_TORCH_COMPILE=0 pytest -v -s --junitxml=$XML_LOG_DIR/test_api_features.xml $TE_PATH/tests/pytorch/debug/test_api_features.py -k \"not (test_per_tensor_scaling or test_fake_quant or test_statistics_collection or test_statistics_multi_run)\" --no-header --feature_dirs=$NVTE_TEST_NVINSPECT_FEATURE_DIRS --configs_dir=$NVTE_TEST_NVINSPECT_CONFIGS_DIR"
+"NVTE_TORCH_COMPILE=0 pytest -v -s --junitxml=$XML_LOG_DIR/test_api_features.xml $TE_PATH/tests/pytorch/debug/test_api_features.py --no-header --feature_dirs=$NVTE_TEST_NVINSPECT_FEATURE_DIRS --configs_dir=$NVTE_TEST_NVINSPECT_CONFIGS_DIR"
 
 # Step 6: Performance
 run_test_step "test_perf.xml" "$TE_PATH/tests/pytorch/debug/test_perf.py" \
@@ -88,11 +88,11 @@ run_test_step "test_perf.xml" "$TE_PATH/tests/pytorch/debug/test_perf.py" \
 # Step 7: Sanity 2
 run_test_step "test_sanity_2.xml" "$TE_PATH/tests/pytorch/test_sanity.py" \
 "NVTE_TEST_NVINSPECT_ENABLED=1 NVTE_TEST_NVINSPECT_CONFIG_FILE=$NVTE_TEST_NVINSPECT_DUMMY_CONFIG_FILE NVTE_TEST_NVINSPECT_FEATURE_DIRS=$NVTE_TEST_NVINSPECT_FEATURE_DIRS PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 \
-pytest -v -s --junitxml=$XML_LOG_DIR/test_sanity_2.xml $TE_PATH/tests/pytorch/test_sanity.py -k \"not (test_sanity_grouped_linear or test_inference_mode)\" --no-header"
+pytest -v -s --junitxml=$XML_LOG_DIR/test_sanity_2.xml $TE_PATH/tests/pytorch/test_sanity.py --no-header"
 
 # Step 8: Numerics 2
 run_test_step "test_numerics_2.xml" "$TE_PATH/tests/pytorch/test_numerics.py" \
 "NVTE_TEST_NVINSPECT_ENABLED=1 NVTE_TEST_NVINSPECT_CONFIG_FILE=$NVTE_TEST_NVINSPECT_DUMMY_CONFIG_FILE NVTE_TEST_NVINSPECT_FEATURE_DIRS=$NVTE_TEST_NVINSPECT_FEATURE_DIRS PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED_ATTN=0 \
-pytest -v -s --junitxml=$XML_LOG_DIR/test_numerics_2.xml $TE_PATH/tests/pytorch/test_numerics.py -k \"not (test_linear_accuracy or test_layernorm_linear_accuracy or test_layernorm_mlp_accuracy or test_grouped_linear_accuracy or test_transformer_layer_hidden_states_format or test_grouped_gemm)\" --no-header"
+pytest -v -s --junitxml=$XML_LOG_DIR/test_numerics_2.xml $TE_PATH/tests/pytorch/test_numerics.py -k \"not (test_linear_accuracy or test_layernorm_linear_accuracy or test_layernorm_mlp_accuracy or test_transformer_layer_hidden_states_format)\" --no-header"
 
 exit $FAIL

qa/L0_pytorch_unittest/test.sh

@@ -6,6 +6,8 @@
 mkdir -p "$XML_LOG_DIR"
 
 pip install pytest==8.2.1
+# solve test_fused_optimizer import error
+pip install expecttest
 FAIL=0
 
 IS_CUDA_BACKEND=$(python3 -c "import torch; print('cuda' if torch.cuda.is_available() else 'cpu')" 2>/dev/null)
@@ -66,8 +68,13 @@ run_test_step() {
 
 
 # Step: Sanity
-run_test_step "pytest_test_sanity.xml" "$TE_PATH/tests/pytorch/test_sanity.py" \
-"python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/test_sanity.py -k \"not (test_sanity_layernorm_mlp or test_sanity_gpt or test_sanity_bert or test_sanity_T5 or test_sanity_amp_and_nvfuser or test_sanity_drop_path or test_sanity_fused_qkv_params or test_sanity_gradient_accumulation_fusion or test_inference_mode or test_sanity_normalization_amp or test_sanity_layernorm_linear or test_sanity_linear_with_zero_tokens or test_sanity_grouped_linear)\" --no-header" "test_sanity.py"
+if [ "$PLATFORM" = "metax" ]; then
+    SANITY_CMD="python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/test_sanity.py -k \"not (test_sanity_layernorm_mlp or test_sanity_gpt or test_sanity_bert or test_sanity_T5 or test_sanity_amp_and_nvfuser or test_sanity_drop_path or test_sanity_fused_qkv_params or test_sanity_gradient_accumulation_fusion or test_inference_mode or test_sanity_normalization_amp or test_sanity_layernorm_linear or test_sanity_linear_with_zero_tokens or test_sanity_grouped_linear)\" --no-header"
+else
+    SANITY_CMD="python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/test_sanity.py --no-header"
+fi
+run_test_step "pytest_test_sanity.xml" "$TE_PATH/tests/pytorch/test_sanity.py" "$SANITY_CMD" "test_sanity.py"
+
 
 # Step: Recipe
 run_test_step "pytest_test_recipe.xml" "$TE_PATH/tests/pytorch/test_recipe.py" \
@@ -78,16 +85,26 @@ run_test_step "pytest_test_deferred_init.xml" "$TE_PATH/tests/pytorch/test_defer
 "python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_deferred_init.xml $TE_PATH/tests/pytorch/test_deferred_init.py" "test_deferred_init.py"
 
 # Step: Numerics
-run_test_step "pytest_test_numerics.xml" "$TE_PATH/tests/pytorch/test_numerics.py" \
-"PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED_ATTN=0 python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/test_numerics.py -k \"not (test_layernorm_mlp_accuracy or test_grouped_linear_accuracy or test_gpt_cuda_graph or test_transformer_layer_hidden_states_format or test_grouped_gemm or test_noncontiguous or test_gpt_checkpointing or test_gpt_accuracy or test_mha_accuracy or test_linear_accuracy or test_linear_accuracy_delay_wgrad_compute or test_rmsnorm_accuracy or test_layernorm_accuracy or test_layernorm_linear_accuracy)\" --no-header" "test_numerics.py"
+if [ "$PLATFORM" = "metax" ]; then
+    NUMERICS_CMD="PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED_ATTN=0 python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/test_numerics.py -k \"not (test_layernorm_mlp_accuracy or test_grouped_linear_accuracy or test_gpt_cuda_graph or test_transformer_layer_hidden_states_format or test_grouped_gemm or test_noncontiguous or test_gpt_checkpointing or test_gpt_accuracy or test_mha_accuracy or test_linear_accuracy or test_linear_accuracy_delay_wgrad_compute or test_rmsnorm_accuracy or test_layernorm_accuracy or test_layernorm_linear_accuracy)\" --no-header"
+else
+    # CUDA
+    NUMERICS_CMD="PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED_ATTN=0 python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/test_numerics.py -k \"not (test_linear_accuracy or test_layernorm_linear_accuracy or test_layernorm_mlp_accuracy or test_transformer_layer_hidden_states_format)\" --no-header"
+fi
+run_test_step "pytest_test_numerics.xml" "$TE_PATH/tests/pytorch/test_numerics.py" "$NUMERICS_CMD" "test_numerics.py"
 
 # Step: CUDA Graphs
 run_test_step "pytest_test_cuda_graphs.xml" "$TE_PATH/tests/pytorch/test_cuda_graphs.py" \
 "PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED_ATTN=0 python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cuda_graphs.xml $TE_PATH/tests/pytorch/test_cuda_graphs.py" "test_cuda_graphs.py"
 
 # Step: JIT
-run_test_step "pytest_test_jit.xml" "$TE_PATH/tests/pytorch/test_jit.py" \
-"python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_jit.xml $TE_PATH/tests/pytorch/test_jit.py -k \"not (test_torch_dynamo)\"" "test_jit.py"
+if [ "$PLATFORM" = "metax" ]; then
+    JIT_CMD="python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_jit.xml $TE_PATH/tests/pytorch/test_jit.py -k \"not (test_torch_dynamo)\""
+else
+    JIT_CMD="python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_jit.xml $TE_PATH/tests/pytorch/test_jit.py --no-header"
+fi
+run_test_step "pytest_test_jit.xml" "$TE_PATH/tests/pytorch/test_jit.py" "$JIT_CMD" "test_jit.py"
+
 
 # Step: Fused Rope
 run_test_step "pytest_test_fused_rope.xml" "$TE_PATH/tests/pytorch/test_fused_rope.py" \
@@ -161,13 +178,20 @@ run_test_step "pytest_test_hf_integration.xml" "$TE_PATH/tests/pytorch/test_hf_i
 run_test_step "pytest_test_checkpoint.xml" "$TE_PATH/tests/pytorch/test_checkpoint.py" \
 "NVTE_TEST_CHECKPOINT_ARTIFACT_PATH=$TE_PATH/artifacts/tests/pytorch/test_checkpoint python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_checkpoint.xml $TE_PATH/tests/pytorch/test_checkpoint.py" "test_checkpoint.py"
 
-# Step: Fused Router
-run_test_step "pytest_test_fused_router.xml" "$TE_PATH/tests/pytorch/test_fused_router.py" \
-"python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_router.xml $TE_PATH/tests/pytorch/test_fused_router.py" "test_fused_router.py"
+# ==============================================================================
+# New Step: Plugin Core
+# ==============================================================================
+
+# Step: Plugin Policy
+run_test_step "pytest_test_plugin_policy.xml" "$TE_PATH/tests/pytorch/test_plugin_policy.py" \
+"python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_plugin_policy.xml $TE_PATH/tests/pytorch/test_plugin_policy.py" "test_plugin_policy.py"
+
+# Step: Plugin manager
+run_test_step "pytest_test_plugin_manager.xml" "$TE_PATH/tests/pytorch/test_plugin_manager.py" \
+"python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_plugin_manager.xml $TE_PATH/tests/pytorch/test_plugin_manager.py" "test_plugin_manager.py"
+
 
-# Step: Partial Cast
-run_test_step "pytest_test_partial_cast.xml" "$TE_PATH/tests/pytorch/test_partial_cast.py" \
-"python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_partial_cast.xml $TE_PATH/tests/pytorch/test_partial_cast.py" "test_partial_cast.py"
+# ==============================================================================
 
 
 if [ "$FAIL" -ne 0 ]; then

qa/L1_pytorch_mcore_integration/output/checkpoints/iter_0000004/metadata.json

@@ -0,0 +1 @@
+{"sharded_backend": "torch_dist", "sharded_backend_version": 1, "common_backend": "torch", "common_backend_version": 1}

qa/L1_pytorch_mcore_integration/output/checkpoints/latest_checkpointed_iteration.txt

@@ -0,0 +1 @@
+4

qa/L1_pytorch_onnx_unittest/test.sh

@@ -2,10 +2,9 @@
 #
 # See LICENSE for license information.
 
-function error_exit() {
-    echo "Error: $1"
-    exit 1
-}
+
+pip3 install onnxruntime
+pip3 install onnxruntime_extensions
 
 : ${TE_PATH:=/opt/transformerengine}
 : ${XML_LOG_DIR:=/logs}