TransformerEngine

1
# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
#
3
# See LICENSE for license information.
4

5
# A workflow to trigger TE build on GitHub
6
name: 'Build'
7
on:
8
  pull_request:
9
  workflow_dispatch:
10
jobs:
11
  pytorch:
12
    name: 'PyTorch'
13
    runs-on: ubuntu-latest
14
    if: false  # NGC PyTorch container does not fit on GitHub runner
15
    container:
16
      image: nvcr.io/nvidia/pytorch:23.03-py3
17
      options: --user root
18
    steps:
19
      - name: 'Checkout'
20
        uses: actions/checkout@v3
21
        with:
22
          submodules: recursive
23
      - name: 'Build'
24
        run: pip install . -v --no-deps
25
        env:
26
          NVTE_FRAMEWORK: pytorch
27
          MAX_JOBS: 1
28
      - name: 'Sanity check'
29
        run: python tests/pytorch/test_sanity_import.py
30
  jax:
31
    name: 'JAX'
32
    runs-on: ubuntu-latest
33
    container:
34
      image: ghcr.io/nvidia/jax:jax
35
      options: --user root
36
    steps:
37
      - name: 'Checkout'
38
        uses: actions/checkout@v3
39
        with:
40
          submodules: recursive
41
      - name: 'Build'
42
        run: pip install . -v
43
        env:
44
          NVTE_FRAMEWORK: jax
45
      - name: 'Sanity check'
46
        run: python tests/jax/test_sanity_import.py
47