[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/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,9 @@ artifacts/
 transformer_engine/plugin/core/_build_config.py
 # Mac OS
 .DS_Store
-.claude/
+# Integration test outputs
+qa/L1_pytorch_mcore_integration/output/
+*.distcp
+.coverage
+.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,14 +178,60 @@ 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"
+
+
+# ==============================================================================
+# New Step: Plugin Core backend
+# ==============================================================================
+
+# Step: Backend flagos =========================================================
+run_test_step "pytest_test_backend_flagos.xml" "$TE_PATH/tests/pytorch/test_backend_flagos.py" \
+"python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_backend_flagos.xml $TE_PATH/tests/pytorch/test_backend_flagos.py" "test_backend_flagos.py"
+
+# Step: Backend impl fused adam
+run_test_step "pytest_test_fused_adam.xml" "$TE_PATH/tests/pytorch/test_fused_adam.py" \
+"python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_adam.xml $TE_PATH/tests/pytorch/test_fused_adam.py" "test_fused_adam.py"
+
+# Step: Backend impl gemm
+run_test_step "pytest_test_gemm.xml" "$TE_PATH/tests/pytorch/test_gemm.py" \
+"python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gemm.xml $TE_PATH/tests/pytorch/test_gemm.py" "test_gemm.py"
+
+# Step: Backend impl multi_tensor
+run_test_step "pytest_test_multi_tensor_unit.xml" "$TE_PATH/tests/pytorch/test_multi_tensor_unit.py" \
+"python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_multi_tensor_unit.xml $TE_PATH/tests/pytorch/test_multi_tensor_unit.py" "test_multi_tensor_unit.py"
+
+# Step: Backend impl rmsnorm
+run_test_step "pytest_test_rmsnorm.xml" "$TE_PATH/tests/pytorch/test_rmsnorm.py" \
+"python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_rmsnorm.xml $TE_PATH/tests/pytorch/test_rmsnorm.py" "test_rmsnorm.py"
+
+# Step: Backend impl softmax
+run_test_step "pytest_test_softmax.xml" "$TE_PATH/tests/pytorch/test_softmax.py" \
+"python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_softmax.xml $TE_PATH/tests/pytorch/test_softmax.py" "test_softmax.py"
+
+
+# Step: Backend reference =========================================================
+run_test_step "pytest_reference.xml" "$TE_PATH/tests/pytorch/test_reference.py" \
+"python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_reference.xml $TE_PATH/tests/pytorch/test_reference.py" "test_reference.py"
+
+run_test_step "pytest_reference_activation.xml" "$TE_PATH/tests/pytorch/test_reference_activation.py" \
+"python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_reference_activation.xml $TE_PATH/tests/pytorch/test_reference_activation.py" "test_reference_activation.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"
+run_test_step "pytest_reference_dropout.xml" "$TE_PATH/tests/pytorch/test_reference_dropout.py" \
+"python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_reference_dropout.xml $TE_PATH/tests/pytorch/test_reference_dropout.py" "test_reference_dropout.py"
 
+run_test_step "pytest_reference_gemm.xml" "$TE_PATH/tests/pytorch/test_reference_gemm.py" \
+"python3 -m pytest -s -v --tb=auto --junitxml=$XML_LOG_DIR/pytest_reference_gemm.xml $TE_PATH/tests/pytorch/test_reference_gemm.py" "test_reference_gemm.py"
 
 if [ "$FAIL" -ne 0 ]; then
     echo "Some tests failed."

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}