fix lint issue

.github/workflows/unit-tests.yaml

@@ -16,23 +16,23 @@ jobs:
         # 1.14.0a0+410ce96
         - torch: "1.14"
           nvcr: 22.12-py3
+          dir: torch1
         # 2.1.0a0+fe05266f
         - torch: "2.1"
           nvcr: 23.10-py3
+          dir: torch2
     container:
       image: nvcr.io/nvidia/pytorch:${{ matrix.nvcr }}
       options: --privileged --ipc=host --gpus=all
     steps:
-    - name: Clean submodules
-      run: |
-        rm -rf third_party
     - name: Checkout msamp
       uses: actions/checkout@v2
       with:
           submodules: true
+          path: ${{ matrix.dir }}
     - name: Install MSCCL
       run: |
-        cd third_party/msccl
+        cd ${{ matrix.dir }}/third_party/msccl
         make -j src.build NVCC_GENCODE="\
           -gencode=arch=compute_70,code=sm_70 \
           -gencode=arch=compute_80,code=sm_80 \
@@ -42,15 +42,18 @@ jobs:
       run: |
         export LD_LIBRARY_PATH="/usr/local/lib:$LD_LIBRARY_PATH"
         python3 -m pip install --upgrade pip
+        cd ${{ matrix.dir }}/
         python3 -m pip install .[test]
         make postinstall
     - name: Run code lint
       run: |
+        cd ${{ matrix.dir }}/
         python3 setup.py lint
     - name: Run unit tests
       run: |
         export LD_LIBRARY_PATH="/usr/local/lib:$LD_LIBRARY_PATH"
         export LD_PRELOAD="/usr/local/lib/libmsamp_dist.so:/usr/local/lib/libnccl.so:${LD_PRELOAD}"
+        cd ${{ matrix.dir }}/
         python3 setup.py test
     # - name: Report coverage results
     #   run: |

msamp/fsdp/_runtime_utils.py

@@ -24,4 +24,3 @@ def _fp8_post_backward_hook(state, handle, *unused):
     setattr(state, comm_hook_attr, state._get_fp8_comm_hook())
     old_post_backward_hook(state, handle, *unused)
     setattr(state, comm_hook_attr, old_communication_hook)
-

msamp/fsdp/flat_param.py

@@ -10,7 +10,8 @@
 class FP8FlatParamHandle(FlatParamHandle):
     """A handle for a flat parameter which may have fp32 and fp8."""
     def __init__(self, *args, **kwargs):
-        super().__init__( *args, **kwargs)
+        """Constructor."""
+        super().__init__(*args, **kwargs)
         self._init_fp8_meta()
 
     def _init_fp8_meta(self):

msamp/fsdp/fully_sharded_data_parallel.py

@@ -5,7 +5,6 @@
 
 import torch
 from torch.distributed.fsdp import FullyShardedDataParallel
-from torch.distributed.algorithms._comm_hooks import default_hooks
 from torch.distributed.fsdp._init_utils import _get_default_comm_hook
 
 from msamp.fsdp.flat_param import FP8FlatParamHandle
@@ -31,10 +30,12 @@ def _fp8_allreduce_hook(state, grad, output):
                     from msamp.operators.dist_op import DistOp
                     dtype = Dtypes.get_dtype_from_qtype(meta.qtype)
                     DistOp.enable_fp8(meta.qtype)
-                    torch.distributed.all_reduce(grad[start:end].view(dtype), group=state.process_group if state else None)
+                    torch.distributed.all_reduce(
+                        grad[start:end].view(dtype), group=state.process_group if state else None
+                    )
                     DistOp.disable_fp8()
                 else:
-                   torch.distributed.all_reduce(grad[start:end], group=state.process_group if state else None)
+                    torch.distributed.all_reduce(grad[start:end], group=state.process_group if state else None)
 
             start = self.rank * output.numel()
             end = (self.rank + 1) * output.numel()

msamp/te/modules.py

@@ -66,8 +66,7 @@ def set_fp8_weights(self):
                     return
 
                 setattr(
-                    self,
-                    weight_cast_attr,
+                    self, weight_cast_attr,
                     Float8Tensor(
                         data=torch.empty(
                             shape,
@@ -80,8 +79,7 @@ def set_fp8_weights(self):
                 )
 
                 setattr(
-                    self,
-                    weight_transpose_attr,
+                    self, weight_transpose_attr,
                     Float8Tensor(
                         data=torch.empty(
                             shape[1],