feat: use composition for non-interactive encrypted training [BLOCKED BY CP] #660

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

it's still a WIP

blocked by https://github.com/zama-ai/concrete-internal/issues/703

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-10T17:55:28Z

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.

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 15 commits June 6, 2024 14:20

chore: add requant in FHE

664aa95

chore: refresh LogisticRegressionTraining notebook

de0f91b

chore: make composition mapping internal

1e3845f

chore: fix pcc

53210a6

chore: fix pytest

48d148f

chore: add back early stopping

9ca59f5

chore: fix tests

2d7afc5

chore: add deployment example with composition and update api

e4f394e

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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: use composition for non-interactive encrypted training [BLOCKED BY CP] #660

feat: use composition for non-interactive encrypted training [BLOCKED BY CP] #660

RomanBredehoft commented Apr 29, 2024 •

edited by jfrery

andrei-stoian-zama May 31, 2024

RomanBredehoft Jun 3, 2024

RomanBredehoft Jun 3, 2024 •

edited

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

RomanBredehoft Jun 5, 2024

RomanBredehoft Jun 5, 2024 •

edited

RomanBredehoft Jun 5, 2024

RomanBredehoft Jun 5, 2024

jfrery commented Jun 7, 2024

github-actions bot commented Jun 10, 2024

feat: use composition for non-interactive encrypted training [BLOCKED BY CP] #660

Are you sure you want to change the base?

feat: use composition for non-interactive encrypted training [BLOCKED BY CP] #660

Conversation

RomanBredehoft commented Apr 29, 2024 • edited by jfrery

Choose a reason for hiding this comment

Choose a reason for hiding this comment

RomanBredehoft Jun 3, 2024 • edited

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

RomanBredehoft Jun 5, 2024 • edited

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 10, 2024

Coverage passed ✅

RomanBredehoft commented Apr 29, 2024 •

edited by jfrery

RomanBredehoft Jun 3, 2024 •

edited

RomanBredehoft Jun 4, 2024 •

edited

RomanBredehoft Jun 5, 2024 •

edited