Add threshold HE, low-rank compression, configurable adjacency normalization#43
Merged
Conversation
- Threshold HE: correct two-party protocol via file-based serialization
of CryptoContext, keys, and ciphertexts; batch partial decryption on
the designated trainer; remove dead single-chunk code path.
- Low-rank: per-client truncated SVD of the pretraining feature sum,
encrypt compressed (U, S, V) factors; server decrypts per trainer,
decompresses, and sums in plaintext for efficient aggregation.
- Normalization: add norm_type parameter to get_1hop_feature_sum with
sym (D^{-1/2} A D^{-1/2}, default, matches GCNConv), row (D^{-1} A,
mean aggregation), and none (raw) options to fix unnormalized sums.
- Model selection: make GraphSage selectable via gnn_model config and
fix unreachable ogbn-products branch.
- Seeds: fix random / numpy / torch seeds before data partitioning for
reproducible runs.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Gate the new FedGCN-v2 features behind config flags so they default to
the original FedGCN behaviour and can be merged to main without
silently changing existing benchmarks.
* utils_nc: `get_1hop_feature_sum` defaults to `norm_type="none"` (the
original unnormalized binary adjacency). FedGCN-v2 callers opt in
with `args.norm_type = "sym"` (symmetric) or `"row"` (row stochastic).
Updates docstring to describe the three modes.
* trainer_class / server_class / federated_methods: top-level
`from fedgraph.openfhe_threshold import OpenFHEThresholdCKKS` is now
wrapped in `try / except ImportError` so the package stays importable
on systems without the OpenFHE wheel. All five trainer call-sites
use `getattr(self.args, "norm_type", "none")` for backward compat.
* federated_methods: `ray.init(...)` honours an optional
`args.ray_init_kwargs` instead of hard-coding a 20 GB plasma store,
so single-machine runs use Ray's own defaults again.
* setup.py: add `extras_require={"openfhe": ["openfhe==1.2.3.0.24.4"]}`
so users can `pip install "fedgraph[openfhe]"` to opt in to the
threshold backend (Linux/manylinux wheels only).
* README: new install section explains the optional `openfhe` extra,
points macOS/Windows users at the Dockerfile, and documents the
`norm_type` opt-in.
* Dockerfile: fix `/var/lib/apt-lists/*` typo so the apt cache is
actually purged.
* data_process: add an opt-in `FEDGRAPH_SUBSAMPLE_NODES` env var for
experimenting with subsampled OGB datasets when full-graph HE is not
feasible.
* .gitignore: exclude local paper drafts, experiment runners, and EC2
backups so they cannot accidentally be staged.
* tests: add `conftest.py` with `needs_openfhe` / `needs_tenseal` skip
markers, plus `test_smoke_unit.py` (norm modes, SVD round-trip, lazy
OpenFHE import) and `test_smoke_e2e.py` (plaintext x {none, sym},
TenSEAL, OpenFHE, OpenFHE+LR end-to-end on Cora).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Brings the merge-prep branch up to date with main and resolves the four
conflicts that come from main and gcn_v2 evolving the same files in
parallel.
Conflict resolutions:
* tests/conftest.py (add/add): keep main's mock-based fixtures and add
our ``needs_openfhe`` / ``needs_tenseal`` skip markers in the same
file.
* fedgraph/server_class.py: restrict ``Server.train``'s encrypted
parameter aggregation branch to the TenSEAL backend. The OpenFHE
threshold path runs from federated_methods.run_NC using the
file-based protocol and must not enter the TenSEAL branch.
* fedgraph/trainer_class.py: combine main's ``None`` short-circuit in
``local_test`` with our explicit ``to(self.device)`` calls.
* fedgraph/federated_methods.py
- Combine main's optional ``differential_privacy`` and ``low_rank``
imports with our optional ``openfhe_threshold`` lazy import. All
three are now ``try / except ImportError``.
- Drop the eager top-level low_rank import in favour of main's lazy
one.
- Keep our seed setup at the start of ``run_fedgraph``.
- Relax main's low-rank validation: low-rank still requires
``fedgraph_task == "NC"``, but is no longer mutually exclusive with
``use_encryption`` or restricted to ``FedAvg`` (FedGCN-v2 combines
low-rank, threshold HE, and FedGCN).
- In the Server constructor, pick ``Server_LowRank`` only when both
``use_lowrank=True`` and the low-rank module imported successfully;
fall back to the standard ``Server`` otherwise, and respect main's
``use_huggingface``-aware feature dimension.
No behaviour change on the original FedGCN baseline: existing
``use_encryption=False`` configs without ``norm_type``, ``use_lowrank``,
or ``he_backend="openfhe"`` use exactly the same code path as on main.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Smoke (12/12 now pass): * federated_methods.run_NC: gate the TenSEAL per-round encrypted parameter aggregation behind ``he_backend == "tenseal"`` so the OpenFHE threshold flow (which only encrypts the pretrain feature sum) falls back to plaintext FedAvg in the inner training loop. Fixes ``AttributeError: 'Trainer' object has no attribute 'he_context'`` when running with ``he_backend="openfhe"``. * federated_methods.run_fedgraph: when ``use_lowrank=True`` and the OpenFHE threshold backend is active, dispatch to ``run_NC`` (which carries the encrypted-SVD pretraining path) instead of ``run_NC_lowrank`` (the plaintext low-rank FedAvg prototype). Fixes the openfhe + low-rank smoke test, which previously skipped pretrain entirely and trained on uninitialized state. Main test suite (zero regressions vs main): * server_class.Server.__init__: when ``use_encryption=True`` and ``he_backend`` is anything other than ``"openfhe"``, default to the TenSEAL backend instead of raising ``ValueError``. Restores the original FedGraph behaviour for unit tests that pass a Mock ``args.he_backend``. * federated_methods.run_fedgraph / trainer_class.__init__: guard ``int(getattr(args, "seed", 42))`` against Mock seeds so unit tests that pass a ``Mock`` args do not blow up at the seed line. Test rewrites (test main's design intent, not the deprecated rules): * test_federated_methods.test_run_fedgraph_lowrank_validation_fedavg_only → renamed to ``..._works_with_non_fedavg_method``; FedGCN-v2 lifts the FedAvg-only restriction, so the test now asserts the call dispatches without raising. * test_federated_methods.test_run_fedgraph_lowrank_encryption_conflict → renamed to ``..._with_openfhe_dispatches_to_run_nc``; FedGCN-v2 combines low-rank with threshold encryption, so the test now asserts the call routes to ``run_NC``. * test_fedgraph_integration.test_configuration_validation: drops the obsolete encryption-vs-low-rank and FedAvg-only assertions and now only checks the still-enforced "low-rank requires NC task" rule. Bench (after fixes, in container with OpenFHE wheel): * tests/test_smoke_unit.py + test_smoke_e2e.py: 12 passed. * tests/unit: 47 failed / 62 passed (main: 48 / 61 — zero regressions, one extra pass). * tests/integration: 7 failed / 7 passed (main: 7 / 7 — zero regressions). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
Author
|
Update after rebasing on the latest main:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add threshold HE, low-rank compression, configurable adjacency normalization (opt-in)
Summary
args.use_encryption=True, args.he_backend="openfhe".Implements two-party threshold CKKS (server + designated client). Wheel
is optional via
pip install "fedgraph[openfhe]".fedgraph.low_rank: per-client truncated SVD,enabled via
args.use_lowrank=True, args.fixed_rank=....args.norm_type("none" default, "sym", "row") onget_1hop_feature_sum. Default keeps the original FedGCN behaviour;FedGCN-v2 sets "sym".
args.seedandargs.ray_init_kwargsfor reproducibility andcluster deployment.
tests/test_smoke_unit.py(math) andtests/test_smoke_e2e.py(full pipeline).Backward compatibility
import fedgraphno longer fails on systems without OpenFHE.