Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: key compression always #726

Merged
merged 6 commits into from
Jun 13, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions docker/release_resources/sanity_check.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@ def ml_check(args, keyring_dir_as_str):
enable_unsafe_features=True,
use_insecure_key_cache=is_fast,
insecure_key_cache_location=keyring_dir_as_str,
compress_input_ciphertexts=True,
)

# We first compile the model with some data, here the training set
Expand Down Expand Up @@ -120,6 +121,8 @@ def function_to_compile(x):
enable_unsafe_features=is_fast,
use_insecure_key_cache=is_fast,
insecure_key_cache_location=keyring_dir_as_str,
compress_input_ciphertexts=True,
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sanity checks should run with default compile parameters

compress_evaluation_keys=True,
)

print("Compiling...")
Expand Down
5 changes: 5 additions & 0 deletions src/concrete/ml/common/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,11 @@
# However, for internal testing purposes, we retain the capability to disable this feature
os.environ["USE_INPUT_COMPRESSION"] = os.environ.get("USE_INPUT_COMPRESSION", "1")

# Enable PBS evaluation key compression (~4x size reduction)
# Note: This setting is fixed and cannot be altered by users
# However, for internal testing purposes, we retain the capability to disable this feature
os.environ["USE_KEY_COMPRESSION"] = os.environ.get("USE_KEY_COMPRESSION", "1")


class FheMode(str, enum.Enum):
"""Enum representing the execution mode.
Expand Down
2 changes: 2 additions & 0 deletions src/concrete/ml/quantization/quantized_module.py
Original file line number Diff line number Diff line change
Expand Up @@ -876,6 +876,7 @@ def compile(

# Enable input ciphertext compression
enable_input_compression = os.environ.get("USE_INPUT_COMPRESSION", "1") == "1"
enable_key_compression = os.environ.get("USE_KEY_COMPRESSION", "1") == "1"

self.fhe_circuit = compiler.compile(
inputset,
Expand All @@ -889,6 +890,7 @@ def compile(
fhe_simulation=False,
fhe_execution=True,
compress_input_ciphertexts=enable_input_compression,
compress_evaluation_keys=enable_key_compression,
)

self._is_compiled = True
Expand Down
3 changes: 3 additions & 0 deletions src/concrete/ml/sklearn/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -569,6 +569,8 @@ def compile(

# Enable input ciphertext compression
enable_input_compression = os.environ.get("USE_INPUT_COMPRESSION", "1") == "1"
# Enable evaluation key compression
enable_key_compression = os.environ.get("USE_KEY_COMPRESSION", "1") == "1"

self.fhe_circuit_ = module_to_compile.compile(
inputset,
Expand All @@ -582,6 +584,7 @@ def compile(
fhe_simulation=False,
fhe_execution=True,
compress_input_ciphertexts=enable_input_compression,
compress_evaluation_keys=enable_key_compression,
)

self._is_compiled = True
Expand Down
2 changes: 1 addition & 1 deletion src/concrete/ml/sklearn/linear_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -347,7 +347,7 @@ def _get_training_quantized_module(
# Enable the underlying FHE circuit to be composed with itself
# This feature is used in order to be able to iterate in the clear n times without having
# to encrypt/decrypt the weight/bias values between each loop
configuration = Configuration(composable=True)
configuration = Configuration(composable=True, compress_evaluation_keys=True)
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

compress pbs keys for training

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

which makes me thinkg, it should be using enable_input_compression as well here !

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes, you're right, same as for DF which was not enabled before

Copy link
Collaborator

@RomanBredehoft RomanBredehoft Jun 13, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

actually same as below, we should probably revert this change : it's enabled by QM's compilation anyway, so it's a bit confusing to add it here

DF is different because it directly compiles with CP tho yes


composition_mapping = {0: 2, 1: 3}

Expand Down
19 changes: 17 additions & 2 deletions tests/deployment/test_client_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,10 +119,25 @@ def test_client_server_sklearn_inference(
max_bit_width = fhe_circuit.graph.maximum_integer_bit_width()
print(f"Max width {max_bit_width}")

# Compare the FHE predictions with the clear ones. Simulated predictions are not considered in
# this test.
# Check that key compression is enabled
assert os.environ.get("USE_KEY_COMPRESSION") == "1", "'USE_KEY_COMPRESSION' is not enabled"

# Check with key compression
check_is_good_execution_for_cml_vs_circuit(x_test, model, simulate=False, n_allowed_runs=1)

# Check without key compression
with pytest.MonkeyPatch.context() as mp_context:

# Disable input ciphertext compression
mp_context.setenv("USE_KEY_COMPRESSION", "0")

# Check that input ciphertext compression is disabled
assert os.environ.get("USE_KEY_COMPRESSION") == "0", "'USE_KEY_COMPRESSION' is not disabled"

# Compare the FHE predictions with the clear ones. Simulated predictions are not
# considered in this test.
check_is_good_execution_for_cml_vs_circuit(x_test, model, simulate=False, n_allowed_runs=1)

# Check client/server FHE predictions vs the FHE predictions of the dev model
check_client_server_inference(
x_test, model, key_dir, check_array_equal, check_float_array_equal
Expand Down
1 change: 1 addition & 0 deletions tests/torch/test_hybrid_converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ def run_hybrid_llm_test(
# Multi-parameter strategy is used in order to speed-up the FHE executions
configuration = Configuration(
single_precision=False,
compress_input_ciphertexts=True,
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

these circuits can be fully leveled or not, always good to test with default compilation parameters

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

so why not also enabling key compression ?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it is enabled by the QM compilation (the env variable defaults to 1)

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

so is compress_input_ciphertexts then right ? imo either we put both here or we revert this small change 🤔

)

# Create a hybrid model
Expand Down
Loading