feat: add lora fine tuning for llama 3.2 #958

jfrery · 2024-12-10T16:25:31Z

No description provided.

docs/advanced_examples/LoraMLP.ipynb

docs/deep-learning/lora_training.md

kcelia · 2024-12-11T14:30:04Z

docs/deep-learning/lora_training.md

 ```

 ### 3. Compile a hybrid FHE model for the LORA adapted PyTorch model

-Compile the hybrid FHE model to convert the selected outsourced layers to use FHE, while the rest will run on the client side. Note that the exchange of encrypted activations and gradients may require significant bandwidth.
+Before training in FHE, we need to compile the model. Compilation calibrates and converts the outsourced linear layers to their FHE equivalents. The compile method uses representative data for this step.


I suggest:

Before training in FHE, the model must first be compiled. This process calibrates and converts the outsourced linear layers into their FHE equivalents. The compilation step needs representative data to ensure accurate calibration.

I will use the passive voice yes good point.

docs/deep-learning/lora_training.md

kcelia · 2024-12-11T14:42:38Z

docs/deep-learning/lora_training.md


 <!--pytest-codeblocks:skip-->

 ```python
-hybrid_model.model.inference_model(x)
+peft_model(x)


much better

You may, precise the default mode of inference here

use_case_examples/lora_finetuning/data_finetune/raw_cml_1.7.0_examples.txt

use_case_examples/lora_finetuning/utils_lora.py

use_case_examples/lora_finetuning/GPT2FineTuneHybrid.ipynb

kcelia · 2024-12-11T16:03:08Z

docs/advanced_examples/LoraMLP.ipynb

-    "    loss_fn=nn.CrossEntropyLoss(),\n",
-    "    training_args={\"gradient_accumulation_steps\": 1},\n",
+    "# Set up LoRA training\n",
+    "lora_trainer = LoraTrainer(\n",


Maybe add a comment here, to say that the LoraTrainer uses the hybrid approach.

At this point, everything is perfectly encapsulated, the user doesn't see a hybrid model, but the title does.
So maybe add a comment here.

We should probably not mention the hybrid model at this point. It introduces complex topics. I don't think there is any mention to hybrid model?

The title mentioned the hybrid model: Setup FHE fine-tuning with LoraTraining and HybridFHEModel

Ah yes that's a miss. I will update to mention LoraTrainer instead.

use_case_examples/lora_finetuning/GPT2FineTuneHybrid.ipynb

use_case_examples/lora_finetuning/LLamaFineTuning.ipynb

src/concrete/ml/torch/lora.py

kcelia · 2024-12-11T16:33:30Z

src/concrete/ml/torch/lora.py

-            n_layers_to_skip (int): Number of layers to skip.
+            model (nn.Module): The model to replace layers in.
+            n_layers_to_skip_for_backprop (int): Number of initial linear layers to keep as standard
+                layers. Since the first layer doesn't need backpropagation (no previous layer to


maybe you should change the signature of the function, since you mentioned default to 1

n_layers_to_skip_for_backprop: int = 1

I suggest:

n_layers_to_skip_for_backprop (int): Determines how many of the first linear layers are excluded from backpropagation. This is typically set to 1 because the first layer only transforms the input data and does not depend on previous layers for gradient updates. By skipping this layer, we save unnecessary computations. Defaults to 1.

Maybe explain why we replace with custom linear layers.
(attach the forward_module and backward_module...)

I will remove the default to one here. It's default to one in th LoraTrainer

We will have to update the documentation. The definition if this variable is already quite complex

kcelia

Thanks for your PR.

Some comments:

if we want to go for LoRA, maybe we should add it in the forbidden list, I stopped spamming you with my LoRA comments lol
The new Lora API is very cool
GPT2 and LLAma notebooks follow the same logic and share same utility functions, maybe we can create a utils file for them.
In GPT2 notebook, I think you don't use the full potential of the new LoRA API, or maybe you wanted to highlight what's happening behind the scene and I did not get it
In the 3 notebooks, I think it's not clear for the reader, if we are using FHE only for the inference or for adapters as well, maybe you should explicitly specify it in the conclusion or the introduction.

jfrery · 2024-12-16T08:26:28Z

GPT2 and LLAma notebooks follow the same logic and share same utility functions, maybe we can create a utils file for them.

I think they share a few function already with the utils file. GPT2 uses the previous API version without the LoraTrainer so a bit more complicated but more flexible as well.

In GPT2 notebook, I think you don't use the full potential of the new LoRA API, or maybe you wanted to highlight what's happening behind the scene and I did not get it

Yes I kept GPT2 without LoraTrainer to show that one could use its own training method but it implies defining the hybrid model / remote layers and so on.

In the 3 notebooks, I think it's not clear for the reader, if we are using FHE only for the inference or for adapters as well, maybe you should explicitly specify it in the conclusion or the introduction.

I will add a sentence at the beginning to make sure what we do here is clear.

use_case_examples/lora_finetuning/GPT2FineTuneHybrid.ipynb

kcelia

Thanks for the changes.

It would be nice to specify if the weights are encrypted too.

kcelia · 2024-12-17T08:42:28Z

use_case_examples/lora_finetuning/README.md

+
+### LLaMA Results
+
+TBD


@jfrery, i think you forgot that part

No we don't have the result yet. it's a WIP

github-actions · 2024-12-17T10:31:11Z

⚠️ Known flaky tests have been rerun ⚠️

One or several tests initially failed but were identified as known flaky. tests. Therefore, they have been rerun and passed. See below for more details.

Failed tests details

Known flaky tests that initially failed:

tests/torch/test_compile_torch.py::test_compile_torch_or_onnx_conv_networks[True-True-CNN_conv1d-relu]- tests/torch/test_compile_torch.py::test_compile_torch_or_onnx_conv_networks[False-True-CNN_grouped-relu]

github-actions · 2024-12-17T10:31:14Z

Coverage passed ✅

Coverage details

---------- coverage: platform linux, python 3.8.18-final-0 -----------
Name    Stmts   Miss  Cover   Missing
-------------------------------------
TOTAL    8482      0   100%

63 files skipped due to complete coverage.

cla-bot bot added the cla-signed label Dec 10, 2024

jfrery marked this pull request as ready for review December 11, 2024 11:13

jfrery requested a review from a team as a code owner December 11, 2024 11:13