feat: lora fine-tuning in FHE + gpt2 use case example (#823)

Co-authored-by: RomanBredehoft <[email protected]> Co-authored-by: jfrery <[email protected]> Co-authored-by: jfrery <[email protected]>
zama-ai · Sep 26, 2024 · 4d2f2e6 · 4d2f2e6
1 parent b37ceed
commit 4d2f2e6
Show file tree

Hide file tree

Showing 19 changed files with 2,881 additions and 13 deletions.
diff --git a/.github/workflows/refresh-one-notebook.yaml b/.github/workflows/refresh-one-notebook.yaml
@@ -21,6 +21,7 @@ on:
         - FullyConnectedNeuralNetwork \n
         - FullyConnectedNeuralNetworkOnMNIST \n
         - GLMComparison \n
+        - GPT2FineTuneHybrid \n
         - HealthCarePrediction \n
         - ImportingFromScikitLearn \n
         - KaggleTitanic \n
@@ -29,6 +30,7 @@ on:
         - LinearSVR \n
         - LogisticRegression \n
         - LogisticRegressionTraining \n
+        - LoraMLP \n
         - PerrorImpactOnFMNIST \n
         - PoissonRegression \n
         - QGPT2Evaluate \n
@@ -67,6 +69,7 @@ env:
   FullyConnectedNeuralNetwork: "docs/advanced_examples/FullyConnectedNeuralNetwork.ipynb" 
   FullyConnectedNeuralNetworkOnMNIST: "docs/advanced_examples/FullyConnectedNeuralNetworkOnMNIST.ipynb" 
   GLMComparison: "docs/advanced_examples/GLMComparison.ipynb" 
+  GPT2FineTuneHybrid: "use_case_examples/lora_finetuning/GPT2FineTuneHybrid.ipynb" 
   HealthCarePrediction: "use_case_examples/disease_prediction/HealthCarePrediction.ipynb" 
   ImportingFromScikitLearn: "docs/advanced_examples/ImportingFromScikitLearn.ipynb" 
   KaggleTitanic: "use_case_examples/titanic/KaggleTitanic.ipynb" 
@@ -75,6 +78,7 @@ env:
   LinearSVR: "docs/advanced_examples/LinearSVR.ipynb" 
   LogisticRegression: "docs/advanced_examples/LogisticRegression.ipynb" 
   LogisticRegressionTraining: "docs/advanced_examples/LogisticRegressionTraining.ipynb" 
+  LoraMLP: "docs/advanced_examples/LoraMLP.ipynb" 
   PerrorImpactOnFMNIST: "use_case_examples/cifar/cifar_brevitas_finetuning/PerrorImpactOnFMNIST.ipynb" 
   PoissonRegression: "docs/advanced_examples/PoissonRegression.ipynb" 
   QGPT2Evaluate: "use_case_examples/llm/QGPT2Evaluate.ipynb" 
@@ -195,11 +199,11 @@ jobs:
         with:
           token: ${{ secrets.BOT_TOKEN }}
           commit-message: "chore: refresh ${{ github.event.inputs.notebook }} notebook"
-          branch: "refresh-${{ github.event.inputs.notebook }}-notebook-for-${{ github.ref_name }}"
+          branch: "refresh-${{ github.event.inputs.notebook }}-notebook-for-branch-${{ github.ref_name }}"
           base: "${{ github.ref_name }}"
-          title: "Refresh ${{ github.event.inputs.notebook }} notebook for ${{ github.ref_name }}"
+          title: "Refresh ${{ github.event.inputs.notebook }} notebook for branch ${{ github.ref_name }}"
           body: "Automatic PR with notebook refresh of ${{ github.event.inputs.notebook }} \
-            for ${{ github.ref_name }}."
+            for branch ${{ github.ref_name }}."
           add-paths: |
             docs/**/*.ipynb
             use_case_examples/**/*.ipynb
@@ -211,6 +215,9 @@ jobs:
         with:
           commit_message: "chore: refresh ${{ github.event.inputs.notebook }} notebook"
           add_options: '-u'
+          file_pattern: |
+            docs/**/*.ipynb
+            use_case_examples/**/*.ipynb
 
 
   stop-runner-linux:

diff --git a/.github/workflows/run_one_use_cases_example.yaml b/.github/workflows/run_one_use_cases_example.yaml
@@ -19,10 +19,16 @@ on:
           - federated_learning
           - hybrid_model
           - llm
+          - lora_finetuning
           - resnet
           - sentiment_analysis_with_transformer
           - titanic
           # --- refresh_use_cases_list.py: refresh list of use cases currently available [END] ---
+      push_changes:
+        description: 'Push refreshed notebook(s)'
+        required: false
+        type: boolean
+        default: false
 
 concurrency:
   group: ${{ github.ref }}
@@ -102,6 +108,37 @@ jobs:
           USE_CASE=${{ github.event.inputs.use_case }}
           make run_one_use_case_example USE_CASE=$USE_CASE
 
+      # Pull the latest changes if there are some
+      - name: Pull latest changes
+        if: ${{ github.event.inputs.push_changes == 'true' }}
+        run: |
+          git pull -X theirs
+
+      # If the target branch is main or a release branch, a Pull Request is opened for everyone to 
+      # review.
+      - name: Open PR
+        if: ${{ github.event.inputs.push_changes == 'true' && (github.ref_name == 'main' || startsWith(github.ref_name , 'release/')) }} 
+        uses: peter-evans/create-pull-request@c5a7806660adbe173f04e3e038b0ccdcd758773c
+        with:
+          token: ${{ secrets.BOT_TOKEN }}
+          commit-message: "chore: refresh notebook(s) for use case ${{ github.event.inputs.use_case }}"
+          branch: "refresh-notebook(s)-for-use-case-${{ github.event.inputs.use_case }}-for-branch-${{ github.ref_name }}"
+          base: "${{ github.ref_name }}"
+          title: "Refresh notebook(s) for use case ${{ github.event.inputs.use_case }} for branch ${{ github.ref_name }}"
+          body: "Automatic PR with notebook(s) refresh of use case ${{ github.event.inputs.use_case }} \
+            for branch ${{ github.ref_name }}."
+          add-paths: |
+            use_case_examples/**/*.ipynb
+
+      # If the target branch is another branch, the current branch is automatically merged into it
+      - name: Push changes into the current branch
+        if: ${{ github.event.inputs.push_changes == 'true' && github.ref_name != 'main' && !(startsWith(github.ref_name , 'release/')) }}
+        uses: stefanzweifel/git-auto-commit-action@8621497c8c39c72f3e2a999a26b4ca1b5058a842 #v5.0.1
+        with:
+          commit_message: "chore: refresh notebook(s) for use case ${{ github.event.inputs.use_case }}"
+          add_options: '-u'
+          file_pattern: 'use_case_examples/**/*.ipynb'
+
   stop-runner-linux:
     name: Stop EC2 runner
     needs: [run-use-case-examples, start-runner-linux]
@@ -162,4 +199,4 @@ jobs:
             - run-use-case-examples: ${{ needs.run-use-case-examples.result || 'Did not run.' }}\n\n\
             - stop-runner-linux: ${{ needs.stop-runner-linux.result || 'Did not run.'}}"
           SLACK_USERNAME: ${{ secrets.BOT_USERNAME }}
-          SLACK_WEBHOOK: ${{ secrets.SLACK_WEBHOOK }}
+          SLACK_WEBHOOK: ${{ secrets.SLACK_WEBHOOK }}
diff --git a/deps_licenses/licenses_linux_user.txt.md5 b/deps_licenses/licenses_linux_user.txt.md5
@@ -1 +1 @@
-31249a607336424af0d790feb9db6252
+af423f91bb5313f1f1670ea72892d364
diff --git a/deps_licenses/licenses_mac_intel_user.txt.md5 b/deps_licenses/licenses_mac_intel_user.txt.md5
@@ -1 +1 @@
-31249a607336424af0d790feb9db6252
+af423f91bb5313f1f1670ea72892d364
diff --git a/deps_licenses/licenses_mac_silicon_user.txt.md5 b/deps_licenses/licenses_mac_silicon_user.txt.md5
@@ -1 +1 @@
-31249a607336424af0d790feb9db6252
+af423f91bb5313f1f1670ea72892d364
diff --git a/docs/advanced_examples/LoraMLP.ipynb b/docs/advanced_examples/LoraMLP.ipynb
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -124,6 +124,7 @@ torchvision = [
     {version = "0.17.2", markers = "platform_system=='Darwin' and platform_machine!='arm64'" },
     {version = "0.18.1", markers = "platform_system!='Darwin' or platform_machine=='arm64'" }
 ]
+peft = "^0.12.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

diff --git a/script/make_utils/run_use_case_examples.sh b/script/make_utils/run_use_case_examples.sh
@@ -41,7 +41,7 @@ install_requirements() {
     local example_dir=$1
     if [ -f "${example_dir}/requirements.txt" ]; then
         pushd "$example_dir"
-        if pip install -r requirements.txt; then
+        if pip install -r requirements.txt --extra-index-url https://pypi.zama.ai/cpu; then
             echo "Requirements installed successfully."
         else
             echo "Failed to install requirements."

diff --git a/src/concrete/ml/torch/hybrid_model.py b/src/concrete/ml/torch/hybrid_model.py
@@ -565,13 +565,22 @@ def save_and_clear_private_info(self, path: Path, via_mlir=True):
         """
         path = Path(path)
         path.mkdir(parents=True, exist_ok=True)
-        for name in self.module_names:
-            module = self._get_module_by_name(self.model, name)
-            # Remove private information
+
+        # Save the complete model (including private info) for the developer
+        complete_model_path = path / "complete_model.pth"
+        torch.save(self.model.state_dict(), complete_model_path.resolve())
+
+        def clear_private_info(module):
             for attr in ["private_module", "calibration_data", "private_q_module"]:
                 if hasattr(module, attr):
                     setattr(module, attr, None)
 
+            for child in module.children():
+                clear_private_info(child)
+
+        # Clear private info for the entire model
+        clear_private_info(self.model)
+
         # Save the model with a specific filename
         model_path = path / "model.pth"
Original file line number	Diff line number	Diff line change
		@@ -1 +1 @@
		31249a607336424af0d790feb9db6252
		af423f91bb5313f1f1670ea72892d364