gpu github runner (#19)

Author: wenxie-amd

.github/workflows/ci.yaml

@@ -49,3 +49,45 @@ jobs:
         if: always()
         with:
           options: "--check --diff --color --verbose --line-length=110"
+
+  run-unittest:
+    needs: code-lint
+    runs-on: run-ut
+    steps:
+      - run: echo "🎉 Begin Primus Unit Test."
+      - name: Change owner
+        run: |
+          echo "change the owner of all primus files, as some files are generated by root"
+          echo "GITHUB_WORKSPACE: ${GITHUB_WORKSPACE}"
+          sudo chown -R $(id -u):$(id -g) ${GITHUB_WORKSPACE}
+      - name: Checkout
+        uses: actions/checkout@v4
+        with:
+          submodules: 'recursive'
+      - name: Print dir
+        run: |
+          echo "GITHUB_WORKSPACE: ${GITHUB_WORKSPACE}"
+          echo "current dir: $(pwd)"
+      - name: Set ut log path
+        run: |
+          if [[ "${{ github.event_name }}" == "pull_request" ]]; then
+            echo "UT_LOG_PATH=$GITHUB_WORKSPACE/ut_out/pr-${{ github.event.pull_request.number }}" >> $GITHUB_ENV
+          elif [[ "${{ github.event_name }}" == "push" && "${{ github.ref }}" == "refs/heads/main" ]]; then
+            echo "UT_LOG_PATH=$GITHUB_WORKSPACE/ut_out/latest" >> $GITHUB_ENV
+          elif [[ "${{ github.event_name }}" == "release" ]]; then
+            TAG_NAME=${{ github.ref }}  # refs/tags/v1.0.0
+            TAG="${TAG_NAME#refs/tags/}" # v1.0.0
+            echo "UT_LOG_PATH=$GITHUB_WORKSPACE/ut_out/$TAG" >> $GITHUB_ENV
+          else
+            echo "UT_LOG_PATH=$GITHUB_WORKSPACE/ut_out/others" >> $GITHUB_ENV
+          fi
+      - name: Run unit tests
+        run: |
+          echo "set UT_LOG_PATH as: ${{ env.UT_LOG_PATH }}"
+          if [ -d "${{ env.UT_LOG_PATH }}" ]; then
+            echo "UT_LOG_PATH dir exists. Deleting..."
+            rm -rf "${{ env.UT_LOG_PATH }}"
+          fi
+          docker exec -e UT_LOG_PATH="${{ env.UT_LOG_PATH }}" \
+                      primus_github_runner bash -c \
+                      "cd $GITHUB_WORKSPACE && python ./tests/run_unit_tests.py"

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "third_party/Megatron-LM-20250324"]
+	path = third_party/Megatron-LM-20250324
+	url = https://github.com/NVIDIA/Megatron-LM.git

README.md

@@ -31,10 +31,10 @@ Use the following command to clone the repo:
 - [ ] Set Megatron-LM as a submodule repo
 ```bash
 mkdir workspace && cd workspace
-git clone [email protected]:AMD-AIG-AIMA/Primus.git
-git clone [email protected]:NVIDIA/Megatron-LM.git
-# version 20250324
-cd Megatron-LM && git checkout d61821b7174bac690afbad9134bcb4983521052f
+# If you are cloning the repository for the first time:
+git clone --recurse-submodules [email protected]:AMD-AIG-AIMA/Primus.git
+# If you've already cloned primus without submodules, run the following commands:
+git submodule update --init --recursive
 ```
 
 ## Setup Primus

examples/megatron/run_pretrain.sh

@@ -12,9 +12,9 @@ export MODEL_CONFIG=${MODEL_CONFIG:-deepseek_v2_lite}
 # framework path
 PRIMUS_PATH=$(realpath "$(dirname "$0")/../..")
 export PRIMUS_PATH
-export MEGATRON_PATH=${MEGATRON_PATH:-${PRIMUS_PATH}/../Megatron-LM}
-[[ ! -d "${MEGATRON_PATH}" ]] && {
-    echo "Error: MEGATRON_PATH (${MEGATRON_PATH}) does not exist"
+export MEGATRON_PATH=${MEGATRON_PATH:-${PRIMUS_PATH}/third_party/Megatron-LM-20250324}
+[[ ! -d "${MEGATRON_PATH}" || -z "$(ls -A "${MEGATRON_PATH}")" ]] && {
+    echo "Error: MEGATRON_PATH (${MEGATRON_PATH}) does not exist or is empty"
     exit 1
 }
 

tests/configs/test_config.py

@@ -28,9 +28,9 @@ def parse_config(self, cli_args: argparse.Namespace):
         exp_config = self.config_parser.parse(cli_args)
         return exp_config
 
-    def test_exp_sft(self):
+    def test_exp_configs(self):
         exps = [
-            "examples/deepseek_v3/exp_pretrain.yaml",
+            "examples/megatron/exp_pretrain.yaml",
         ]
 
         for exp in exps:

tests/run_unit_tests.py

@@ -17,7 +17,7 @@
 def get_all_unit_tests():
     global DISTRIBUTED_UNIT_TESTS
 
-    cur_dir = "./"
+    cur_dir = "./tests"
     unit_tests = {}
 
     for root, dirs, files in os.walk(cur_dir):
@@ -37,7 +37,7 @@ def launch_unit_test(ut_path, nproc_per_node):
     global UNIT_TEST_PASS
 
     if nproc_per_node == 1:
-        cmd = f"pytest {ut_path}"
+        cmd = f"pytest {ut_path} -s"
     else:
         cmd = f"torchrun --nnodes 1 --nproc-per-node {nproc_per_node} {ut_path}"
 

tests/trainer/__init__.py

@@ -0,0 +1,56 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+#################################################################################
+
+
+import os
+import subprocess
+import sys
+import time
+import unittest
+
+from primus.core.utils import logger
+from tests.utils import PrimusUT
+
+
+class TestMegatronTrainer(PrimusUT):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def setUp(self):
+        pass
+
+    def tearDown(self):
+        pass
+
+    def test_pretrain(self):
+        shell_entry = "examples/megatron/run_pretrain.sh"
+        do_print_at_runtime = False
+        run_stdout = subprocess.PIPE if not do_print_at_runtime else sys.stdout
+        run_stderr = subprocess.PIPE if not do_print_at_runtime else sys.stderr
+        try:
+            logger.info(f"Begin run {shell_entry}...")
+            start = time.time()
+            result = subprocess.run(
+                ["bash", f"{shell_entry}"],
+                check=True,
+                stdout=run_stdout,
+                stderr=run_stderr,
+                text=True,
+            )
+            logger.info(f"End run {shell_entry}, time={time.time()-start:.3f} s")
+            if not do_print_at_runtime:
+                ut_log_path = os.environ.get("UT_LOG_PATH", "ut_out")
+                logger.info(f"Training log path: {ut_log_path}/logs/UT-{self.__class__.__name__}")
+
+            logger.debug(f"Standard Output:\n {result.stdout}")
+            logger.debug(f"Standard Error:\n {result.stderr}")
+        except subprocess.CalledProcessError as e:
+            os.environ["SCRIPT_ERROR"] = e.stderr.strip()
+            assert False, f"Shell script failed: {os.environ['SCRIPT_ERROR']}"
+
+
+if __name__ == "__main__":
+    unittest.main(buffer=False)

tests/trainer/test_megatron_trainer.py

@@ -17,7 +17,7 @@ def __init__(self, *args, **kwargs):
 
     @classmethod
     def setUpClass(cls):
-        ut_log_path = os.environ.get("UT_LOG_PATH", "output/ut")
+        ut_log_path = os.environ.get("UT_LOG_PATH", "ut_out")
         logger_cfg = logger.LoggerConfig(
             exp_root_path=ut_log_path,
             work_group="develop",