Add pylint to Lint action (#19)

* Add pylint to Lint action
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Test Ubuntu 20.04
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Pylint inside the container
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Update transformer_engine/pytorch/distributed.py
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

.github/workflows/lint.yml

@@ -16,7 +16,23 @@ jobs:
         uses: actions/checkout@v3
       - name: Lint
         run: |
-          cp qa/L0_lint/CPPLINT.cfg . && \
-          export CPP_ONLY=1 && \
-          export TE_PATH=. && \
+          cp qa/L0_lint/CPPLINT.cfg .
+          export CPP_ONLY=1
+          export TE_PATH=.
           bash ./qa/L0_lint/test.sh
+  pylint:
+    name: Python
+    runs-on: ubuntu-latest
+    container:
+      image: nvcr.io/nvidia/pytorch:22.09-py3
+      options: --user root
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+      - name: Lint
+        run: |
+          cp qa/L0_lint/pylintrc .
+          export PYTHON_ONLY=1
+          export TE_PATH=.
+          bash ./qa/L0_lint/test.sh
+

qa/L0_lint/test.sh

@@ -7,12 +7,15 @@ set -e
 : "${TE_PATH:=/opt/transformerengine}"
 
 pip install cpplint==1.6.0 pylint==2.13.5
-echo "Checking common API headers"
-cd $TE_PATH && \
-cpplint --root transformer_engine/common/include --recursive transformer_engine/common/include
-echo "Checking C++ files"
-cd $TE_PATH && \
-cpplint --recursive --exclude=transformer_engine/common/include transformer_engine
+if [ -z "${PYTHON_ONLY}" ]
+then
+  echo "Checking common API headers"
+  cd $TE_PATH && \
+  cpplint --root transformer_engine/common/include --recursive transformer_engine/common/include
+  echo "Checking C++ files"
+  cd $TE_PATH && \
+  cpplint --recursive --exclude=transformer_engine/common/include transformer_engine
+fi
 if [ -z "${CPP_ONLY}" ]
 then
   echo "Checking Python files"

transformer_engine/pytorch/distributed.py

@@ -181,6 +181,8 @@ class CheckpointFunction(torch.autograd.Function):
         kwargs: Dict[str, Any],
         *args: Tuple[torch.Tensor, ...],
     ) -> Tuple[torch.Tensor, ...]:
+        """Call forward function while saving state to be able to
+        redo the computation later."""
         ctx.run_function = run_function
         ctx.distribute_saved_activations = distribute_saved_activations
 
@@ -216,8 +218,9 @@ class CheckpointFunction(torch.autograd.Function):
 
     @staticmethod
     def backward(
-        ctx, *args: Tuple[torch.Tensor, ...]
+        ctx, *args: Tuple[Union[torch.Tensor, None], ...]
     ) -> Tuple[Union[torch.Tensor, None], ...]:
+        """Call backward function with activation recomputation."""
         if not torch.autograd._is_checkpoint_valid():
             raise RuntimeError(
                 "Checkpointing is not compatible with .grad(), "