Merge branch 'main' into pipe_main_merge

Author: MayDomine

.github/workflows/build.yml

@@ -16,7 +16,7 @@ jobs:
     uses: OpenBMB/BMTrain/.github/workflows/build_whl.yml@main
     secrets: inherit
 
-  publish:
+  fake-publish:
     needs: build-archive-wheel
     runs-on: ubuntu-latest
     steps:
@@ -29,7 +29,7 @@ jobs:
         python-version: 3.9
 
     - name: Download distribution files
-      uses: actions/download-artifact@v2
+      uses: actions/download-artifact@v4
       with:
         name: dist
         path: dist

.github/workflows/build_whl.yml

@@ -44,12 +44,46 @@ jobs:
         version=${{ matrix.python-version }}
         docker run -e BUILD_DOCKER_ENV=1 -e CUDACXX=/usr/local/cuda-11.3/bin/nvcc -e PATH="/opt/rh/devtoolset-9/root/usr/bin:$PATH" -e LD_LIBRARY_PATH="/opt/rh/devtoolset-9/root/usr/lib64:/opt/rh/devtoolset-9/root/usr/lib:/usr/local/nvidia/lib:/usr/local/nvidia/lib64:$LD_LIBRARY_PATH" -v ${{ github.workspace }}:/workspace/BMTrain -i pytorch/manylinux-cuda113:latest /bin/bash -c "cd /workspace/BMTrain;/opt/python/cp${version}*/bin/pip install build; /opt/python/cp${version}*/bin/python -m build .;for file in dist/*-linux_x86_64.whl; do mv \"\$file\" \"\${file//-linux_x86_64/-manylinux2014_x86_64}\"; done"
 
+    - name: Upload wheels as artifacts
+      uses: actions/upload-artifact@v4
+      with:
+        name: wheels_py${{ matrix.python-version }}
+        path: dist/*.whl
+
+    - name: Upload source distribution (only once)
+      if: matrix.python-version == '37'  # Only upload source distribution once
+      uses: actions/upload-artifact@v4
+      with:
+        name: source_dist
+        path: dist/*.tar.gz
+    
+  archive:
+    runs-on: ubuntu-latest
+    needs: build
+    steps:
+    - name: Download all wheels
+      uses: actions/download-artifact@v4
+      with:
+        path: wheels
+        pattern: wheels_py*
+
+    - name: Download source distribution
+      uses: actions/download-artifact@v4
+      with:
+        path: source_dist
+        name: source_dist
+
+    - name: Combine all wheels into a single directory
+      run: |
+        mkdir -p dist
+        find wheels -name '*.whl' -exec mv {} dist/ \;
+        find source_dist -name '*.tar.gz' -exec mv {} dist/ \;
+
     - name: Archive distribution files
-      uses: actions/upload-artifact@v2
+      uses: actions/upload-artifact@v4
       with:
         name: dist
         path: |
           dist/*.tar.gz
           dist/*.whl
-
-    
+        overwrite: true

.github/workflows/publish.yaml

@@ -27,7 +27,7 @@ jobs:
       run: python -m pip install twine
 
     - name: Download distribution files
-      uses: actions/download-artifact@v2
+      uses: actions/download-artifact@v4
       with:
         name: dist
         path: dist

.github/workflows/release.yml

@@ -27,21 +27,18 @@ jobs:
         python-version: 3.9
 
     - name: Download distribution files
-      uses: actions/download-artifact@v2
+      uses: actions/download-artifact@v4
       with:
         name: dist
         path: dist
 
-    - name: Upload Distribution Files
+    - name: Upload Distribution Files to Existing Release
       uses: softprops/action-gh-release@v1
       with:
-        body_path: "Release.txt"
         files: |
           dist/*.tar.gz
           dist/*.whl
-        prerelease: false
-        name: "BMTrain v0.2.3"
+        tag_name: ${{ github.ref_name }}  # 使用当前触发工作流的 tag
         token: ${{ secrets.RELEASE_TOKEN }}
-        tag_name: ${{ steps.create_release.outputs.tag }}
       env:
         GITHUB_REPOSITORY: OpenBMB/BMTrain

bmtrain/block_layer.py

@@ -308,7 +308,10 @@ def post_hook(self, out):
         post_out = tuple(post_out)
         return post_out
 
-    def forward(self, *args):
+    def forward(self, *args, **kwargs):
+        signature = inspect.signature(self._module.forward)
+        bound_args = signature.bind(*args, **kwargs)
+        args = bound_args.args
         arg_list = self.pre_hook(*args)
         if self.all_input_no_grad and not self.all_param_no_grad:
             placeholder = torch.tensor([], requires_grad=torch.is_grad_enabled())

bmtrain/wrapper.py

@@ -15,11 +15,15 @@ def make_distributed(model: torch.nn.Module):
     for kw in list(model._buffers.keys()):
         if model._buffers[kw] is not None:
             model._buffers[kw] = model._buffers[kw].cuda()
-
+    is_module_list = isinstance(model, torch.nn.ModuleList)
+    pre_module = None
     for kw in list(model._modules.keys()):
-        if isinstance(model, torch.nn.ModuleList):
+        if is_module_list:
             if not isinstance(model._modules[kw], Block):
                 model._modules[kw] = Block(model_wrapper_dispatch(model._modules[kw]))
+                if pre_module is not None:
+                    model._modules[kw].set_pre_module(pre_module)
+                pre_module = model._modules[kw]
         else:
             model._modules[kw] = model_wrapper_dispatch(model._modules[kw])
 

setup.py

@@ -93,7 +93,7 @@ def build_extension(self, ext):
 ]
 setup(
     name='bmtrain',
-    version='1.0.0',
+    version='1.0.1',
     author="Guoyang Zeng",
     author_email="[email protected]",
     description="A toolkit for training big models",