test_3d_parallelism.yaml

name: Run unit tests

on:
  workflow_dispatch:
  push:
    branches: [ main ]
    # paths:
    #   - "src/**/*.py"
    #   - "examples/**/*.py"
    #   - "tests/**/*.py"

  pull_request:
    branches: [ main ]
    #paths:
    #  - "src/**/*.py"
    #  - "examples/**/*.py"
    #  - "tests/**/*.py"

jobs:
  tests:
    runs-on: [multi-gpu, nvidia-gpu, 8-t4, ci]
    container:
      image: nvcr.io/nvidia/pytorch:23.03-py3
      ports:
        - 80
      options: --gpus all --shm-size "8G"
    steps:
    - uses: actions/checkout@v3
    - name: Set up Python 3.10
      uses: actions/setup-python@v2
      with:
        python-version: '3.10'

    - name: Python environment
      run: |
        which python
        python --version

    - name: Check Pytorch version
      run: |
        nvidia-smi
        python -c "import torch; print('torch:', torch.__version__, torch)"
        python -c "import torch; print('CUDA available:', torch.cuda.is_available())"

    - name: Instal nanotron
      run: |
        python -m pip install --upgrade pip
        pip install packaging
        pip install wheel
        pip install "flash-attn>=2.4.2"  --no-build-isolation
        git clone https://github.com/huggingface/nanotron.git
        cd nanotron
        pip install -e .
        pip install -e [dev]
        pip install -e [test]

    - name: Install test dependencies
      run: |
        pip install pytest
        pip install pytest-cov

    - name: Show installed libraries and their versions
      run: pip freeze | tee installed.txt

    - name: Run tests
      run: pytest --color=yes --durations=0 --verbose tests/