Basic support of jax.lax.gather in TessellateIPU.
#107
Workflow file for this run
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: tests_internal | |
| env: | |
| GIT_MAIN_BRANCH: "main" | |
| # Controls when the workflow will run. | |
| on: | |
| push: | |
| branches: [ "main" ] | |
| pull_request: | |
| branches: [ "main" ] | |
| # Allows you to run this workflow manually from the Actions tab. | |
| workflow_dispatch: | |
| jobs: | |
| lint_and_typecheck: | |
| if: github.repository != 'graphcore-research/tessellate-ipu' | |
| runs-on: [self-hosted, Linux, X64, 20.04, Ubuntu, intel] | |
| timeout-minutes: 10 | |
| steps: | |
| - name: Cancel previous | |
| uses: styfle/[email protected] | |
| with: | |
| access_token: ${{ github.token }} | |
| if: ${{github.ref != 'refs/head/main'}} | |
| - uses: actions/checkout@v3 | |
| - name: Set up Python 3.8 | |
| uses: actions/setup-python@v4 | |
| with: | |
| python-version: 3.8 | |
| - uses: pre-commit/[email protected] | |
| # Tessellate unit tests using IPU model | |
| tessellate_unit_tests_ipu_model: | |
| if: github.repository != 'graphcore-research/tessellate-ipu' | |
| runs-on: [self-hosted, Linux, X64, 20.04, Ubuntu, amd] | |
| container: graphcore/pytorch:3.2.0-ubuntu-20.04 | |
| timeout-minutes: 20 | |
| steps: | |
| - name: Cancel previous | |
| uses: styfle/[email protected] | |
| with: | |
| access_token: ${{ github.token }} | |
| if: ${{github.ref != 'refs/head/main'}} | |
| - uses: actions/checkout@v3 | |
| - name: Update pip | |
| id: pip-cache | |
| run: | | |
| python3 -m pip install --upgrade build cmake nanobind ninja pip scikit_build_core[pyproject] | |
| # Install JAX for IPU. | |
| - name: Install JAX for IPU | |
| run: | | |
| pip3 install --default-timeout=100 -U numpy==1.23.5 scipy etils pytest | |
| pip3 install jax==0.3.16+ipu jaxlib==0.3.15+ipu.sdk320 -f https://graphcore-research.github.io/jax-experimental/wheels.html | |
| # Install other dependencies. | |
| - name: Install dependencies | |
| run: | | |
| pip3 install --default-timeout=100 -r ./test-requirements.txt | |
| # Build & install sdist package. | |
| - name: Install TessellateIPU | |
| run: | | |
| python3 -m build --no-isolation | |
| pip3 install dist/*.tar.gz | |
| # Run repository unit tests. | |
| - name: Run TessellateIPU unit tests (model) | |
| run: | | |
| cd ./tests/ | |
| JAX_IPU_USE_MODEL=true JAX_IPU_MODEL_NUM_TILES=16 pytest --tb=short -v --log-cli-level=INFO ./ | |
| # Dockerized workflow known to create issues with self-hosted servers. | |
| # Solution is to fully cleanup the workspace for the next action. | |
| # See: https://stackoverflow.com/questions/70483902/how-to-actually-clean-up-the-repository-on-self-hosted-runner-after-github-actio | |
| - name: Cleanup GITHUB_WORKSPACE folder | |
| run: | | |
| ls -la ./ | |
| rm -rf ./* || true | |
| rm -rf ./.??* || true | |
| ls -la ./ | |
| # Tessellate unit tests using IPU hardware | |
| tessellate_unit_tests_ipu_hardware: | |
| if: github.repository != 'graphcore-research/tessellate-ipu' | |
| runs-on: [self-hosted, Linux, X64, 20.04, Ubuntu, pod16] | |
| container: | |
| image: graphcore/pytorch:3.2.0-ubuntu-20.04 | |
| options: --ulimit memlock=-1:-1 --cap-add=SYS_PTRACE --cap-add=IPC_LOCK --device=/dev/infiniband/ -e IPUOF_VIPU_API_HOST -e IPUOF_VIPU_API_PARTITION_ID --shm-size=128G | |
| timeout-minutes: 20 | |
| steps: | |
| - name: Cancel previous | |
| uses: styfle/[email protected] | |
| with: | |
| access_token: ${{ github.token }} | |
| if: ${{github.ref != 'refs/head/main'}} | |
| - uses: actions/checkout@v3 | |
| - name: Instal Python base dependencies | |
| id: pip-cache | |
| run: | | |
| python3 -m pip install --upgrade build cmake docker nanobind ninja pip scikit_build_core[pyproject] | |
| # Necessary for attaching IPU hardware. | |
| - name: Attach RDMA Network | |
| run: | | |
| python3 -c "import docker; client=docker.from_env(); client.networks.get('macvlan_rdma_swarm').connect(client.containers.get('${{ job.container.id }}'))" | |
| gc-monitor | |
| # Install JAX for IPU. | |
| - name: Install JAX for IPU | |
| run: | | |
| pip3 install --default-timeout=100 -U numpy==1.23.5 scipy etils pytest | |
| pip3 install jax==0.3.16+ipu jaxlib==0.3.15+ipu.sdk320 -f https://graphcore-research.github.io/jax-experimental/wheels.html | |
| POPLAR_LOG_LEVEL=DEBUG TF_CPP_MIN_LOG_LEVEL=0 JAX_IPU_DEVICE_COUNT=1 python .github/scripts/check_ipu_hardware_device.py | |
| - name: Install repository dependencies | |
| run: | | |
| pip3 install --default-timeout=100 -r ./test-requirements.txt | |
| - name: Install TessellateIPU | |
| run: | | |
| python3 -m build --no-isolation | |
| pip3 install dist/*.whl | |
| # Run IPU hardware unit tests. | |
| - name: Run TessellateIPU unit tests (hardware) | |
| run: | | |
| cd ./tests/ | |
| JAX_IPU_DEVICE_COUNT=2 pytest --tb=short -v --log-cli-level=INFO -m ipu_hardware ./ |