feat: use composition for non-interactive encrypted training #660

RomanBredehoft · 2024-04-29T16:49:46Z

it's still a WIP

closes https://github.com/zama-ai/concrete-ml-internal/issues/4374
closes https://github.com/zama-ai/concrete-ml-internal/issues/4485 - analysis of the new training approach

src/concrete/ml/quantization/quantized_module.py

src/concrete/ml/sklearn/linear_model.py

src/concrete/ml/quantization/quantized_module.py

src/concrete/ml/sklearn/linear_model.py

andrei-stoian-zama · 2024-05-31T16:09:46Z

src/concrete/ml/quantization/quantized_module.py

+                for output_i, input_i in self._composition_mapping.items()
+            )
+
+            if len(q_results) == 1:


can you assert on the shape here ? the input/output shapes should match

I've added the checks in the new _add_requant_for_composition method (name to be confirmed)

RomanBredehoft · 2024-06-03T16:50:06Z

src/concrete/ml/quantization/quantized_module.py

@@ -290,6 +293,61 @@ def _set_output_quantizers(self) -> List[UniformQuantizer]:
        )
        return output_quantizers

+    # Remove this once we handle the re-quantization step in post-training only
+    # FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/4472
+    def _add_requant_for_composition(self, composition_mapping: Optional[Dict]):


new (private) method for quantized module (avoids adding a param to the init and thus keep thing really internal)

RomanBredehoft · 2024-06-03T16:51:09Z

src/concrete/ml/quantization/quantized_module.py

+        max_output_pos = len(self.output_quantizers) - 1
+        max_input_pos = len(self.input_quantizers) - 1
+
+        for output_position, input_position in composition_mapping.items():


make sure the mapping is of the form {0:1, 3:2}

RomanBredehoft · 2024-06-03T16:53:50Z

src/concrete/ml/quantization/quantized_module.py

+
+        # Ignore [arg-type] check from mypy as it is not able to see that the input to `quant`
+        # cannot be None
+        q_x = tuple(


these are needed to match how CP works with encrypt, ie encrypt(None, x) = None, x_enc, since we do not encrypt all inputs at the same time with composition

RomanBredehoft · 2024-06-03T16:56:01Z

src/concrete/ml/sklearn/linear_model.py

+
+        # Similarly, we only quantize the weight and bias values using the third and fourth
+        # position parameter
+        _, _, q_weights, q_bias = self.training_quantized_module.quantize_input(


see https://github.com/zama-ai/concrete-ml/pull/660/files#r1624776488

RomanBredehoft · 2024-06-03T16:56:50Z

src/concrete/ml/torch/compile.py

@@ -181,16 +184,32 @@ def _compile_torch_or_onnx_model(
            for each input. By default all arguments will be encrypted.
        reduce_sum_copy (bool): if the inputs of QuantizedReduceSum should be copied to avoid
            bit-width propagation
+        composition_mapping (Optional[Dict]): Dictionary that maps output positions with input


adding this new parameter to the private funct _compile_torch_or_onnx_model instead of the other public ones to keep things internal

RomanBredehoft · 2024-06-03T16:57:26Z

src/concrete/ml/torch/compile.py

+    # If a mapping between input and output quantizers is set, add a re-quantization step at the
+    # end of the forward call. This is only useful for composable circuits in order to make sure
+    # that input and output quantizers match
+    if composition_mapping is not None:


this is where we decide to add the requant step or not

RomanBredehoft · 2024-06-04T12:27:38Z

src/concrete/ml/sklearn/linear_model.py

            # Additionally, there is no point in computing the following in case of a partial fit,
            # as it only represents a single iteration
            if self.early_stopping and not is_partial_fit:
+                weights_float, bias_float = self._decrypt_dequantize_training_output(


we keep early stopping possible with composition by adding this decrypt/dequant step here (since this is only for development, we believe it's not really an issue to do that)

RomanBredehoft · 2024-06-05T14:20:52Z

src/concrete/ml/deployment/fhe_client_server.py

+    # FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/4477
+    # We should also rename the input arguments to remove the `serialized` part, as we now accept
+    # both serialized and deserialized input values
+    # FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/4476
    def run(


we now allow serialized and deserialized inputs (avoids having to deser + ser at each server call with composition)

RomanBredehoft · 2024-06-05T14:21:44Z

src/concrete/ml/deployment/fhe_client_server.py

@@ -357,97 +388,78 @@ def get_serialized_evaluation_keys(self) -> bytes:
        return self.client.evaluation_keys.serialize()

    def quantize_encrypt_serialize(
-        self, x: Union[numpy.ndarray, Tuple[numpy.ndarray, ...]]
-    ) -> Union[bytes, Tuple[bytes, ...]]:
+        self, *x: Optional[numpy.ndarray]


we now allow unpacking. This is not a breaking change since allow tuples has been added by @jfrery only recently

this is mainly to make things more coherent with other methods + concrete

RomanBredehoft · 2024-06-05T14:21:51Z

src/concrete/ml/deployment/fhe_client_server.py

    def deserialize_decrypt(
-        self, serialized_encrypted_quantized_result: Union[bytes, Tuple[bytes, ...]]
+        self, *serialized_encrypted_quantized_result: Optional[bytes]


RomanBredehoft · 2024-06-05T14:22:20Z

src/concrete/ml/deployment/fhe_client_server.py

    def deserialize_decrypt_dequantize(
-        self, serialized_encrypted_quantized_result: Union[bytes, Tuple[bytes, ...]]
-    ) -> numpy.ndarray:
+        self, *serialized_encrypted_quantized_result: Optional[bytes]


jfrery · 2024-06-07T15:53:00Z

If so are we sure the new PBS dequant requant isn't impacting the convergence?

good question, i don't know. The requant part makes sure that values are the same with respect to what we were doing in interactive training before, so this part is ok. But I guess you are more worried of the PBS part and the fact that we are using rounding, right ? In that case yes not sure how we could easily assess that it does not impact the convergence. All I can say for now is that tests are fine and the notebook looks good 😅 @jfrery

I did the analysis -> #660 (comment).

All in all, looks like the convergence isn't impacted. Good to go!

github-actions · 2024-06-11T10:16:32Z

Coverage passed ✅

Coverage details

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

60 files skipped due to complete coverage.

andrei-stoian-zama

Very well done! thanks

jfrery

All good thanks!

cla-bot bot added the cla-signed label Apr 29, 2024

RomanBredehoft force-pushed the feat/use_composition_encrypted_training_4374 branch from 62e5ad5 to ce9fb57 Compare April 30, 2024 08:46

RomanBredehoft force-pushed the feat/use_composition_encrypted_training_4374 branch 2 times, most recently from a420bef to 2ffb370 Compare May 23, 2024 09:58

RomanBredehoft changed the title ~~feat: use composition for non-interactive encrypted training~~ feat: use composition for non-interactive encrypted training [BLOCKED BY CP] May 27, 2024

RomanBredehoft force-pushed the feat/use_composition_encrypted_training_4374 branch 5 times, most recently from 332ee71 to afca943 Compare May 29, 2024 13:30

RomanBredehoft commented May 30, 2024

View reviewed changes

src/concrete/ml/quantization/quantized_module.py Show resolved Hide resolved

RomanBredehoft commented May 30, 2024

View reviewed changes

src/concrete/ml/sklearn/linear_model.py Show resolved Hide resolved

andrei-stoian-zama requested changes May 31, 2024

View reviewed changes

RomanBredehoft mentioned this pull request Jun 3, 2024

feat: add FHE training deployment #665

Merged

RomanBredehoft force-pushed the feat/use_composition_encrypted_training_4374 branch 3 times, most recently from 2d2ed7a to 454f8fa Compare June 3, 2024 13:52

RomanBredehoft commented Jun 3, 2024

View reviewed changes

RomanBredehoft commented Jun 4, 2024

View reviewed changes

RomanBredehoft requested review from jfrery and andrei-stoian-zama June 5, 2024 14:19

RomanBredehoft commented Jun 5, 2024

View reviewed changes

RomanBredehoft and others added 7 commits June 6, 2024 14:20

chore: make FheModelServer.run handle deserialized inputs

805d72d

chore: refresh LogisticRegressionTraining notebook

07480d7

chore: update concrete-python

1186a7a

chore: update licenses

a47ceb0

chore: take comments into account

1f477ac

chore: refresh LogisticRegressionTraining notebook

6a738aa

chore: update lock file

c54b458

RomanBredehoft force-pushed the feat/use_composition_encrypted_training_4374 branch from 4cc0545 to c54b458 Compare June 6, 2024 12:23

RomanBredehoft added 2 commits June 6, 2024 12:29

chore: update licenses

5aeec92

chore: refresh LogisticRegressionTraining notebook

4be39c8

RomanBredehoft marked this pull request as draft June 6, 2024 14:53

RomanBredehoft added 2 commits June 6, 2024 17:25

chore: add new asserts and tests

0497dd0

chore: fix pcc and coverage part 1

9f352b7

RomanBredehoft force-pushed the feat/use_composition_encrypted_training_4374 branch from a8eaab9 to 9f352b7 Compare June 7, 2024 08:50

chore: add test for fhe training in deployment

e75d777

RomanBredehoft added 3 commits June 10, 2024 11:52

chore: fix test

e1647ae

chore: fix test

1045b91

chore: fix tests

1037569

RomanBredehoft force-pushed the feat/use_composition_encrypted_training_4374 branch from 7a52745 to 1037569 Compare June 10, 2024 16:20

chore: speedup tests

99fe596

RomanBredehoft changed the title ~~feat: use composition for non-interactive encrypted training [BLOCKED BY CP]~~ feat: use composition for non-interactive encrypted training Jun 11, 2024

chore: fix and speed up fhe training tests

9328d57

RomanBredehoft force-pushed the feat/use_composition_encrypted_training_4374 branch from 156fb4d to 9328d57 Compare June 11, 2024 09:01

RomanBredehoft marked this pull request as ready for review June 11, 2024 12:37

andrei-stoian-zama approved these changes Jun 11, 2024

View reviewed changes

jfrery approved these changes Jun 11, 2024

View reviewed changes

RomanBredehoft merged commit ec58bca into main Jun 11, 2024
13 checks passed

RomanBredehoft deleted the feat/use_composition_encrypted_training_4374 branch June 11, 2024 14:35

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: use composition for non-interactive encrypted training #660

feat: use composition for non-interactive encrypted training #660

RomanBredehoft commented Apr 29, 2024 •

edited

Loading

andrei-stoian-zama May 31, 2024

RomanBredehoft Jun 3, 2024

RomanBredehoft Jun 3, 2024 •

edited

Loading

RomanBredehoft Jun 3, 2024

RomanBredehoft Jun 3, 2024

RomanBredehoft Jun 3, 2024

RomanBredehoft Jun 3, 2024

RomanBredehoft Jun 3, 2024

RomanBredehoft Jun 4, 2024 •

edited

Loading

RomanBredehoft Jun 5, 2024

RomanBredehoft Jun 5, 2024 •

edited

Loading

RomanBredehoft Jun 5, 2024

RomanBredehoft Jun 5, 2024

jfrery commented Jun 7, 2024

github-actions bot commented Jun 11, 2024

andrei-stoian-zama left a comment

jfrery left a comment

feat: use composition for non-interactive encrypted training #660

feat: use composition for non-interactive encrypted training #660

Conversation

RomanBredehoft commented Apr 29, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

RomanBredehoft Jun 3, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

RomanBredehoft Jun 4, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

RomanBredehoft Jun 5, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jfrery commented Jun 7, 2024

github-actions bot commented Jun 11, 2024

Coverage passed ✅

andrei-stoian-zama left a comment

Choose a reason for hiding this comment

jfrery left a comment

Choose a reason for hiding this comment

RomanBredehoft commented Apr 29, 2024 •

edited

Loading

RomanBredehoft Jun 3, 2024 •

edited

Loading

RomanBredehoft Jun 4, 2024 •

edited

Loading

RomanBredehoft Jun 5, 2024 •

edited

Loading