diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml index 9e74726..345b22a 100644 --- a/.github/workflows/python-publish.yml +++ b/.github/workflows/python-publish.yml @@ -11,6 +11,8 @@ name: Upload Python Package on: release: types: [published] + pull_request: + branches: [main] permissions: contents: read @@ -26,14 +28,16 @@ jobs: uses: actions/setup-python@v5 with: python-version: '3.x' + - name: Install the dependencies - run: | - python -m pip install --upgrade pip - pip install build twine + run: pip install build twine + - uses: gaurav-nelson/github-action-markdown-link-check@v1 - name: Build and publish + run: python -m build --wheel + + - name: Publish + if: github.event_name == 'release' env: TWINE_USERNAME: __token__ TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }} - run: | - python -m build --wheel - twine upload dist/* + run: twine upload dist/* diff --git a/README.md b/README.md index 0042dfc..c430eeb 100644 --- a/README.md +++ b/README.md @@ -27,10 +27,10 @@ In this work, we propose a new optimizer, **LO**w-Memory **O**ptimization (**LOM Our approach enables the full parameter fine-tuning of a 7B model on a single RTX 3090, or a 65B model on a single machine with 8×RTX 3090, each with 24GB memory. -![LOMO](assets/LOMO.png) +![LOMO](https://raw.githubusercontent.com/OpenLMLab/LOMO/main/assets/LOMO.png) ## Implementation -![Hook function](assets/hook_func.png) +![Hook function](https://raw.githubusercontent.com/OpenLMLab/LOMO/main/assets/hook_func.png) Our implementation relies on injecting hook functions into PyTorch's backward pass. As depicted in the figure, we register a customized hook function for each parameter. When the gradient of a parameter is computed (prior to writing it to the .grad attribute), its corresponding hook function is invoked. For more information about hook functions and the backward pass of the autograd graph, please refer to [PyTorch's documentation](https://pytorch.org/docs/stable/notes/autograd.html#backward-hooks-execution). In summary, during the backward pass, we go through a tensor and its grad_fn, write the gradient into the .grad attribute, and then pass to the next tensor. Our customized hook function scans all the parameters, updating a parameter if its .grad attribute is not empty, and then clears and frees the .grad attribute. Since the hook function for a parameter is called before its .grad attribute is set, the .grad attribute of the last parameter in the autograd graph is not ready when the last hook function is invoked. Therefore, we perform an additional scan to update the last parameter. @@ -42,7 +42,7 @@ The code for LOMO is in [lomo](lomo) folder. In this work, we examined the distinctions between the LOMO and Adam optimization techniques and introduce AdaLomo, which provides an adaptive learning rate for each parameter and utilizes grouped update normalization while maintaining memory efficiency. AdaLomo achieves results comparable to AdamW in both instruction-tuning and further pre-training with less memory footprint. -![AdaLomo](assets/adalomo_algorithm.png) +![AdaLomo](https://raw.githubusercontent.com/OpenLMLab/LOMO/main/assets/adalomo_algorithm.png) The code for AdaLomo is in [adalomo](adalomo) folder. diff --git a/lomo_optim/__init__.py b/lomo_optim/__init__.py index d5be4ce..d814e5a 100644 --- a/lomo_optim/__init__.py +++ b/lomo_optim/__init__.py @@ -1,5 +1,5 @@ from .adalomo import AdaLomo from .lomo import Lomo -__version__ = "0.1.0" +__version__ = "0.1.1" __all__ = ["Lomo", "AdaLomo"]