[Feature] Use uv instead of pip in CI/CD

.github/workflows/3d_parallelism_unit_tests.yaml

@@ -39,20 +39,22 @@ jobs:
 
     - name: Install nanotron's dependencies
       run: |
+        python -m venv $GITHUB_WORKSPACE/venv
+        . $GITHUB_WORKSPACE/venv/bin/activate
+        export PATH="$PATH:$(python -c 'import site; print(site.USER_BASE + "/bin")')"
         python -m pip install --upgrade pip
-        pip install packaging
-        pip install wheel
-        pip install -e .
-        pip install -e .[dev]
-        pip install -e .[test]
-
-    - name: Show installed libraries and their versions
-      run: pip freeze | tee installed.txt
+        python -m pip install --upgrade uv
+        uv pip install packaging
+        uv pip install wheel
+        uv pip install -e .
+        uv pip install -e .[dev]
+        uv pip install -e .[test]
 
     - name: Run nanotron tests
       # NOTE: -m "not fa2" will run all the unit tests that don't have the mark
       # "fa2" (these are FA2-related tests, we can't run it on T4)
       run: |
+        . $GITHUB_WORKSPACE/venv/bin/activate
         pytest \
         -m "not fa2" \
         --color=yes \
@@ -61,6 +63,7 @@ jobs:
         --ignore tests/fp8 \
         --verbose \
         tests/
+
     # NOTE: T4 can't run FA2, DoReMi's LLaMa needs FÀ
     # - name: Run DoReMi tests
     #   # NOTE: -m "not fa2" will run all the unit tests that don't have the mark

.github/workflows/fa2_unit_tests.yaml

@@ -39,20 +39,29 @@ jobs:
         python -c "import torch; print('torch:', torch.__version__, torch)"
         python -c "import torch; print('CUDA available:', torch.cuda.is_available())"
 
-    - name: Instal nanotron
+    - name: Install nanotron
+      # NOTE: uv can't build flash-attn
       run: |
+        python -m venv $GITHUB_WORKSPACE/venv
+        . $GITHUB_WORKSPACE/venv/bin/activate
         python -m pip install --upgrade pip
-        pip install packaging
-        pip install wheel
+        python -m pip install --upgrade uv
+        uv pip install packaging
+        uv pip install wheel
+        uv pip install torch
         pip install "flash-attn>=2.5.0" --no-build-isolation
-        pip install -e .
-        pip install -e .[dev]
-        pip install -e .[test]
+        uv pip install -e .
+        uv pip install -e .[dev]
+        uv pip install -e .[test]
 
     - name: Show installed libraries and their versions
-      run: pip freeze | tee installed.txt
+      run: |
+        . $GITHUB_WORKSPACE/venv/bin/activate
+        uv pip freeze | tee installed.txt
 
     - name: Run tests
       # NOTE: -m fa2 will only run the unit tests that have the mark
       # "fa2" (these are FA2-related tests)
-      run: pytest -m fa2 --color=yes --durations=0 --ignore tests/fp8 --verbose tests/
+      run: |
+        . $GITHUB_WORKSPACE/venv/bin/activate
+        pytest -m fa2 --color=yes --durations=0 --ignore tests/fp8 --verbose tests/

tests/test_checkpointing.py

@@ -16,7 +16,7 @@ def __init__(self, is_checkpointed: bool = False):
         self.fwd_counter = 0
 
     @checkpoint_method("is_checkpointed")
-    def forward(self, x: Union[torch.Tensor, TensorPointer]):
+    def forward(self, x: Union[torch.Tensor, TensorPointer]) -> torch.Tensor:
         x = self.dense1(x)
         if self.is_checkpointed and self.fwd_counter == 0:
             assert not x.requires_grad, "x should not require grad when checkpointed, because fwd runs in no_grad mode"