Add initial muon support (#3)

Add muon and tests.

Also set everything up.

Signed-off-by: Hao Wu <[email protected]>

Author: skyw

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

@@ -36,7 +36,7 @@ inputs:
   image:
     description: "Image to use for test"
     required: false
-    default: "llm_shower"
+    default: "emerging_optimizers"
   cpu-only:
     description: "Run tests on CPU only"
     required: false
@@ -140,7 +140,7 @@ runs:
     - name: Checkout repository
       uses: actions/checkout@v2
       with:
-        path: LLM-Shower
+        path: Emerging-Optimizers
 
     - name: Start container
       shell: bash
@@ -164,7 +164,7 @@ runs:
           --env HYDRA_FULL_ERROR=1 \
           --env HF_HOME=/home/TestData/HF_HOME \
           --env RUN_ID=${{ github.run_id }} \
-          --volume $(pwd)/LLM-Shower:/workspace \
+          --volume $(pwd)/Emerging-Optimizers:/workspace \
           --volume $MNT_PATH/TestData:/home/TestData \
           nemoci.azurecr.io/${{ inputs.image }}:${{ github.run_id }} \
           bash -c "sleep $(( ${{ inputs.timeout }} * 60 + 60))"
@@ -189,7 +189,7 @@ runs:
 
           docker exec -t nemo_container_${{ github.run_id }} bash -c '\
             uv pip install --no-deps -e . && \
-            bash tests/${{ inputs.is_unit_test == 'true' && 'unit_tests' || 'functional_tests' }}/${{ inputs.script }}.sh && \
+            bash tests/${{ inputs.is_unit_test == 'true' && 'ci' }}/${{ inputs.script }}.sh && \
             echo "Finished successfully." || echo "Did not finish."'
         ) 2>&1 | tee err.log
 

.github/workflows/_build_container.yml

@@ -51,7 +51,7 @@ jobs:
         run: |
           LAST_PRS=$(gh api graphql -f query='
             query {
-              repository(owner: "NVIDIA-NeMo", name: "LLM-Shower") {
+              repository(owner: "NVIDIA-NeMo", name: "Emerging-Optimizers") {
                 pullRequests(states: MERGED, first: 100, orderBy: {field: UPDATED_AT, direction: DESC}) {
                   nodes {
                     number

.github/workflows/build-test-publish-wheel.yml

@@ -30,7 +30,7 @@ jobs:
     if: ${{ vars.BUILD_TEST_PUBLISH_WHEEL == 'true' }}
     with:
       dry-run: true
-      python-package: llm_shower
+      python-package: emerging_optimizers
       python-version: "3.10"
       packaging: uv
     secrets:

.github/workflows/cicd-main.yml

@@ -42,7 +42,7 @@ jobs:
     uses: ./.github/workflows/_build_container.yml
     needs: cicd-wait-in-queue
     with:
-      image-name: llm_shower
+      image-name: emerging_optimizers
       dockerfile: docker/Dockerfile.ci
       runner: self-hosted-nemo
     secrets:
@@ -55,7 +55,7 @@ jobs:
       fail-fast: false
       matrix:
         include:
-          - script: L0_Unit_Tests_GPU
+          - script: L0_Tests_GPU
             runner: self-hosted-nemo
             timeout: 30
     needs: [cicd-container-build]
@@ -72,14 +72,44 @@ jobs:
           script: ${{ matrix.script }}
           timeout: ${{ matrix.timeout || 10 }}
           is_unit_test: "true"
-          image: llm_shower
+          image: emerging_optimizers
           cpu-only: ${{ matrix.cpu-only || false }}
           has-azure-credentials: "true"
           azure-client-id: ${{ secrets.AZURE_CLIENT_ID }}
           azure-tenant-id: ${{ secrets.AZURE_TENANT_ID }}
           azure-subscription-id: ${{ secrets.AZURE_SUBSCRIPTION_ID }}
 
+  cicd-unit-tests-cpu:
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - script: L0_Tests_CPU
+            runner: linux-amd64-cpu16
+            timeout: 30
+    needs: [cicd-container-build]
+    runs-on: ${{ matrix.runner }}
+    name: ${{ matrix.script }}
+    environment: nemo-ci
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: main
+        uses: ./.github/actions/test-template
+        with:
+          runner: ${{ runner.name }}
+          script: ${{ matrix.script }}
+          timeout: ${{ matrix.timeout || 10 }}
+          is_unit_test: "true"
+          image: emerging_optimizers
+          cpu-only: ${{ matrix.cpu-only || true }}
+          has-azure-credentials: "true"
+          azure-client-id: ${{ secrets.AZURE_CLIENT_ID }}
+          azure-tenant-id: ${{ secrets.AZURE_TENANT_ID }}
+          azure-subscription-id: ${{ secrets.AZURE_SUBSCRIPTION_ID }}
+
   cicd-e2e-tests:
+    if: false
     strategy:
       fail-fast: false
       matrix:
@@ -100,7 +130,7 @@ jobs:
           runner: ${{ runner.name }}
           script: ${{ matrix.script }}
           timeout: ${{ matrix.timeout || 10 }}
-          image: llm_shower
+          image: emerging_optimizers
           cpu-only: ${{ matrix.cpu-only || false }}
           has-azure-credentials: "true"
           azure-client-id: ${{ secrets.AZURE_CLIENT_ID }}

.github/workflows/release-freeze.yml

@@ -36,8 +36,8 @@ jobs:
   code-freeze:
     uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/[email protected]
     with:
-      library-name: LLM-Shower
-      python-package: llm_shower
+      library-name: Emerging-Optimizers
+      python-package: emerging_optimizers
       release-type: ${{ inputs.release-type }}
       freeze-commit: ${{ inputs.freeze-commit }}
       dry-run: ${{ inputs.dry-run }}

.github/workflows/release.yaml

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-name: "Release LLM-Shower"
+name: "Release Emerging-Optimizers"
 
 on:
   workflow_dispatch:
@@ -40,8 +40,8 @@ jobs:
     uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/[email protected]
     with:
       release-ref: ${{ inputs.release-ref }}
-      python-package: llm_shower
-      library-name: LLM-Shower
+      python-package: emerging_optimizers
+      library-name: Emerging-Optimizers
       dry-run: ${{ inputs.dry-run }}
       version-bump-branch: ${{ inputs.version-bump-branch }}
       create-gh-release: ${{ inputs.create-gh-release }}

.pre-commit-config.yaml

@@ -32,6 +32,12 @@ repos:
       args: ["check", "--select", "I", "--fix"]
     - id: ruff-format
 
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.14.0
+    hooks:
+      - id: mypy
+        exclude: ^docs|^tests|^benchmarks|^docker
+
   - repo: local
     hooks:
       - id: no-underscore-md

CONTRIBUTING.md

@@ -1,4 +1,4 @@
-# Contributing To LLM-Shower
+# Contributing To Emerging-Optimizers
 
 ## Signing Your Work
 

README.md

@@ -1 +1 @@
-# LLM Shower
+# Emerging Optimizers

docker/Dockerfile.ci

@@ -25,8 +25,8 @@ ENV PATH="$UV_PROJECT_ENVIRONMENT/bin:$PATH"
 
 WORKDIR /workspace
 RUN --mount=type=bind,source=pyproject.toml,target=/workspace/pyproject.toml \
-    --mount=type=bind,source=llm_shower/__init__.py,target=/workspace/llm_shower/__init__.py \
-    --mount=type=bind,source=llm_shower/package_info.py,target=/workspace/llm_shower/package_info.py \
+    --mount=type=bind,source=emerging_optimizers/__init__.py,target=/workspace/emerging_optimizers/__init__.py \
+    --mount=type=bind,source=emerging_optimizers/package_info.py,target=/workspace/emerging_optimizers/package_info.py \
     --mount=type=bind,source=uv.lock,target=/workspace/uv.lock bash -exu <<"EOF"
 
     # Use the container's torch installation rather than reinstall it