TransformerEngine/.github/workflows/build.yml at main · augmentcode/TransformerEngine · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.

# A workflow to trigger TE build on GitHub
name: 'Build'
on:
  pull_request:
  workflow_dispatch:
jobs:
  core:
    name: 'Core'
    runs-on: ubuntu-latest
    container:
      image: nvcr.io/nvidia/cuda:12.1.0-devel-ubuntu22.04
      options: --user root
    steps:
      - name: 'Dependencies'
        run: |
          apt-get update
          apt-get install -y git python3.9 pip cudnn9-cuda-12
          pip install cmake==3.21.0 pybind11[global] ninja
      - name: 'Checkout'
        uses: actions/checkout@v3
        with:
          submodules: recursive
      - name: 'Build'
        run: pip install --no-build-isolation . -v
        env:
          NVTE_FRAMEWORK: none
          MAX_JOBS: 1
      - name: 'Sanity check'
        run: python3 -c "import transformer_engine"
        working-directory: /
  pytorch:
    name: 'PyTorch'
    runs-on: ubuntu-latest
    container:
      image: nvcr.io/nvidia/cuda:12.8.0-devel-ubuntu22.04
      options: --user root
    steps:
      - name: 'Dependencies'
        run: |
          apt-get update
          apt-get install -y git python3.9 pip cudnn9-cuda-12
          pip install cmake torch ninja pydantic importlib-metadata>=1.0 packaging pybind11 numpy einops onnxscript
      - name: 'Checkout'
        uses: actions/checkout@v3
        with:
          submodules: recursive
      - name: 'Build'
        run: pip install --no-build-isolation . -v --no-deps
        env:
          NVTE_FRAMEWORK: pytorch
          MAX_JOBS: 1
      - name: 'Sanity check'
        run: python3 tests/pytorch/test_sanity_import.py
  jax:
    name: 'JAX'
    runs-on: ubuntu-latest
    container:
      image: ghcr.io/nvidia/jax:jax
      options: --user root
    steps:
      - name: 'Dependencies'
        run: pip install pybind11[global]
      - name: 'Checkout'
        uses: actions/checkout@v3
        with:
          submodules: recursive
      - name: 'Build'
        run: pip install --no-build-isolation . -v
        env:
          NVTE_FRAMEWORK: jax
          MAX_JOBS: 1
      - name: 'Sanity check'
        run: python3 tests/jax/test_sanity_import.py
  all:
    name: 'All'
    runs-on: ubuntu-latest
    container:
      image: ghcr.io/nvidia/jax:jax
      options: --user root
    steps:
      - name: 'Dependencies'
        run: |
          pip install pybind11[global] einops onnxscript
          pip install torch --index-url https://download.pytorch.org/whl/cu130
      - name: 'Checkout'
        uses: actions/checkout@v3
        with:
          submodules: recursive
      - name: 'Build'
        run: pip install --no-build-isolation . -v --no-deps
        env:
          NVTE_FRAMEWORK: all
          MAX_JOBS: 1
      - name: 'Sanity check'
        run: python3 tests/pytorch/test_sanity_import.py && python3 tests/jax/test_sanity_import.py