Fix issues w. python 3.6 + add py-version checks to CI (#2589)

Author: jeffra

.github/workflows/python.yml

@@ -0,0 +1,39 @@
+name: python
+
+on:
+  push:
+    branches:
+      - 'master'
+      - 'staging**'
+  pull_request:
+    branches:
+      '**'
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  version-check:
+    strategy:
+      matrix:
+        pyVersion: ["3.6", "3.7", "3.8", "3.9", "3.10"]
+      fail-fast: false
+
+    runs-on: ubuntu-20.04
+    container:
+      image: deepspeed/gh-builder:py${{ matrix.pyVersion }}
+
+    steps:
+        - uses: actions/checkout@v2
+
+        - name: environment
+          run: |
+            which python
+            python --version
+        - name: Install deepspeed
+          run: |
+            pip3 install .
+        - name: DS Report
+          run: |
+             ds_report

deepspeed/model_implementations/diffusers/unet.py

@@ -12,6 +12,7 @@ def __init__(self, unet, enable_cuda_graph=True):
         self.in_channels = unet.in_channels
         self.device = self.unet.device
         self.dtype = self.unet.dtype
+        self.config = self.unet.config
         self.fwd_count = 0
         self.unet.requires_grad_(requires_grad=False)
         self.unet.to(memory_format=torch.channels_last)

deepspeed/module_inject/replace_module.py

@@ -3,9 +3,12 @@
 import tqdm
 import deepspeed
 import deepspeed.ops.transformer as transformer_inference
+from deepspeed.ops.transformer.inference.diffusers_attention import DeepSpeedDiffusersAttention
+from deepspeed.ops.transformer.inference.diffusers_transformer_block import DeepSpeedDiffusersTransformerBlock
+from deepspeed.ops.transformer.inference.diffusers_2d_transformer import Diffusers2DTransformerConfig
 from .replace_policy import HFBertLayerPolicy, HFGPT2LayerPolicy, BLOOMLayerPolicy
 from .replace_policy import replace_policies, generic_policies
-#from ..runtime.weight_quantizer import WeightQuantization
+
 from deepspeed import comm as dist
 from torch import nn
 
@@ -211,7 +214,7 @@ def replace_attn(child, policy):
             triangular_masking=False,
             max_out_tokens=4096,
         )
-        attn_module = transformer_inference.DeepSpeedDiffusersAttention(config)
+        attn_module = DeepSpeedDiffusersAttention(config)
 
         def transpose(data):
             data = data.contiguous()
@@ -234,8 +237,8 @@ def transpose(data):
         return attn_module
 
     def replace_attn_block(child, policy):
-        config = transformer_inference.Diffusers2DTransformerConfig()
-        return transformer_inference.DeepSpeedDiffusersTransformerBlock(child, config)
+        config = Diffusers2DTransformerConfig()
+        return DeepSpeedDiffusersTransformerBlock(child, config)
 
     if isinstance(module, torch.nn.Module):
         pass

deepspeed/ops/transformer/__init__.py

@@ -2,6 +2,3 @@
 from .inference.config import DeepSpeedInferenceConfig
 from ...model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference
 from .inference.moe_inference import DeepSpeedMoEInferenceConfig, DeepSpeedMoEInference
-from .inference.diffusers_attention import DeepSpeedDiffusersAttention
-from .inference.diffusers_transformer_block import DeepSpeedDiffusersTransformerBlock
-from .inference.diffusers_2d_transformer import Diffusers2DTransformerConfig

deepspeed/ops/transformer/inference/__init__.py

@@ -1,6 +1,3 @@
 from .config import DeepSpeedInferenceConfig
 from ....model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference
 from .moe_inference import DeepSpeedMoEInferenceConfig, DeepSpeedMoEInference
-from .diffusers_attention import DeepSpeedDiffusersAttention
-from .diffusers_transformer_block import DeepSpeedDiffusersTransformerBlock
-from .diffusers_2d_transformer import Diffusers2DTransformerConfig

deepspeed/ops/transformer/inference/diffusers_transformer_block.py

@@ -5,8 +5,8 @@
 import torch
 import torch.nn as nn
 from ... import op_builder
-from ....module_inject import GroupQuantizer
 
+from deepspeed import module_inject
 from .diffusers_attention import DeepSpeedDiffusersAttention
 from .bias_add import nhwc_bias_add
 from .diffusers_2d_transformer import Diffusers2DTransformerConfig
@@ -35,7 +35,7 @@ def __init__(self,
                  equivalent_module: nn.Module,
                  config: Diffusers2DTransformerConfig):
         super(DeepSpeedDiffusersTransformerBlock, self).__init__()
-        self.quantizer = GroupQuantizer(q_int8=config.int8_quantization)
+        self.quantizer = module_inject.GroupQuantizer(q_int8=config.int8_quantization)
         # Ensure ops are built by the time we start running
         self.config = config
 

setup.py

@@ -305,7 +305,8 @@ def create_dir_symlink(src, dest):
           'Programming Language :: Python :: 3.6',
           'Programming Language :: Python :: 3.7',
           'Programming Language :: Python :: 3.8',
-          'Programming Language :: Python :: 3.9'
+          'Programming Language :: Python :: 3.9',
+          'Programming Language :: Python :: 3.10'
       ],
       license='MIT',
       ext_modules=ext_modules,