Limb-parallel ciphertext pipeline
Example source: examples/10_spmd_limb_parallel_pipeline.py
This example partitions one logical ciphertext along its active RNS-limb axis, runs limb-local addition and multiplication, and reconstructs the complete basis for global transitions. The tutorial identifies the exact gather points in that pipeline.
Run on two GPUs
torchrun --standalone --nproc-per-node=2 \
examples/10_spmd_limb_parallel_pipeline.py2
The example evaluates
1. Define contiguous limb ranges
ranges = [
(
rank * limb_count // world_size,
(rank + 1) * limb_count // world_size,
)
for rank in range(world_size)
]2
3
4
5
6
7
Each shard is created with:
ciphertext.slice_limbs(start, stop)The resulting value keeps an ordered contiguous interval of the parent prime_ids. Rank number alone is not enough to reconstruct RNS identity.
2. Scatter and add limb-local values
local_a = dist.scatter_ciphertext_limbs(shards_a, src=0)
local_b = dist.scatter_ciphertext_limbs(shards_b, src=0)
local_sum = engine.add(local_a, local_b)2
3
Ciphertext addition is independent for each modulus row. Every rank can apply the public engine operation to its local interval without receiving the other rows.
3. Gather before a complete-row transition
full_sum = dist.gather_ciphertext_limbs(local_sum, dst=0)
prepared_sum = engine.coefficient_domain_to_ntt_domain(full_sum)2
NTT conversion preserves the level and operates independently on each modulus, but the example reconstructs the complete ciphertext before repartitioning so rank zero can derive and transmit one exact active-basis layout.
4. Multiply local intervals
local_operand = dist.scatter_ciphertext_limbs(prepared_shards, src=0)
local_triplet = engine.multiply(local_operand, local_operand)2
Once both operands satisfy the fixed preconditions, ciphertext multiplication is independent per active modulus row. Each rank produces the same local three-component structure over its own prime interval.
5. Reconstruct before relinearization
full_triplet = dist.gather_ciphertext_limbs(local_triplet, dst=0)
result = engine.rescale_to_next_level(
engine.relinearize(full_triplet, relinearization_key)
)2
3
4
Relinearization uses the complete hybrid decomposition/key-switch layout, so it is kept on rank zero after structural reconstruction. Worker ranks never receive the relinearization key. The product carries the pending default_scale.
Gather is not reduce
gather_ciphertext_limbs concatenates disjoint prime intervals:
It does not add the rows. An arithmetic reduction would combine residues that belong to different moduli and destroy the ciphertext structure.
When this pattern helps
Limb partitioning is useful when:
- one complete value or prepared parameter is too large for the desired per-device budget;
- there is enough expensive limb-local work between scatter/gather phases;
- the application can keep operations requiring every active row sparse and explicit.
It is less attractive when every operation immediately needs reconstruction; communication then dominates the local modular arithmetic.
The engine accepts validated local intervals
Rank-local operations still use public CkksEngine methods. The engine validates context, device, ring dimension, modulus basis, and contiguous ordered prime_ids; the distributed facade does not introduce a separate sharded value type.
Complete runnable source
#!/usr/bin/env python3
"""Limb-parallel arithmetic on one logical ciphertext.
Run on one process or one process per GPU:
python examples/10_spmd_limb_parallel_pipeline.py
torchrun --standalone --nproc-per-node=2 \
examples/10_spmd_limb_parallel_pipeline.py
The program evaluates ``(a + b)^2`` in two limb-local stages:
1. scatter level-0 limb fragments, add locally, and reconstruct ``a + b``;
2. rank 0 performs full-basis rescale/NTT preparation, scatters the remaining
limbs, ranks run local ciphertext multiplication, and rank 0 reconstructs and relinearizes.
Use this pattern when one large ciphertext is structurally partitioned across
its RNS limb axis. The application chooses every limb range. Operations that
need the complete basis remain gather points, and no relinearization
key is sent to worker ranks.
"""
from __future__ import annotations
import torch
import fhelium as fh
import fhelium.distributed as dist
def _limb_ranges(limb_count: int) -> list[tuple[int, int]]:
if dist.get_world_size() > limb_count:
raise ValueError(
f"world_size={dist.get_world_size()} exceeds limbs={limb_count}"
)
return [
(
rank * limb_count // dist.get_world_size(),
(rank + 1) * limb_count // dist.get_world_size(),
)
for rank in range(dist.get_world_size())
]
def _split_limbs(
ciphertext: fh.Ciphertext,
ranges: list[tuple[int, int]],
) -> list[fh.Ciphertext]:
return [ciphertext.slice_limbs(start, stop) for start, stop in ranges]
def main() -> None:
dist.init()
engine = fh.CkksEngine(
fh.Preset.slots32768_scale40_levels34_int64,
device=dist.local_device(),
allow_sk_gen=False,
)
message_a = torch.linspace(-0.008, 0.011, 32, dtype=torch.float64)
message_b = torch.linspace(0.006, -0.004, 32, dtype=torch.float64)
addition_ranges = _limb_ranges(engine.config.num_q_primes)
if dist.get_rank() == 0:
secret_key = engine.create_secret_key()
public_key = engine.create_public_key(secret_key)
relinearization_key = engine.create_relinearization_key(secret_key)
ciphertext_a = engine.encrypt_message(message_a, public_key)
ciphertext_b = engine.encrypt_message(message_b, public_key)
shards_a = _split_limbs(ciphertext_a, addition_ranges)
shards_b = _split_limbs(ciphertext_b, addition_ranges)
else:
secret_key = None
relinearization_key = None
shards_a = None
shards_b = None
# Stage 1: add is independent for every modulus, so every rank can apply
# the ordinary operation to its selected contiguous prime interval.
local_a = dist.scatter_ciphertext_limbs(shards_a, src=0)
local_b = dist.scatter_ciphertext_limbs(shards_b, src=0)
local_sum = engine.add(local_a, local_b)
full_sum = dist.gather_ciphertext_limbs(local_sum, dst=0)
# Stage 2 preparation changes domain and uses the full active basis. It is
# therefore performed only after reconstruction on rank 0.
multiplication_ranges = _limb_ranges(engine.config.num_q_primes)
if dist.get_rank() == 0:
assert full_sum is not None
prepared_sum = engine.coefficient_domain_to_ntt_domain(full_sum)
prepared_shards = _split_limbs(
prepared_sum,
multiplication_ranges,
)
else:
prepared_shards = None
local_operand = dist.scatter_ciphertext_limbs(prepared_shards, src=0)
# Exact-state ciphertext multiplication is also limb-local. It returns three components;
# relinearization is deliberately delayed until the limbs are complete.
local_triplet = engine.multiply(local_operand, local_operand)
full_triplet = dist.gather_ciphertext_limbs(local_triplet, dst=0)
add_start, add_stop = addition_ranges[dist.get_rank()]
mul_start, mul_stop = multiplication_ranges[dist.get_rank()]
print(
f"rank={dist.get_rank()} add_limbs=[{add_start},{add_stop}) "
f"multiply_limbs=[{mul_start},{mul_stop})"
)
if dist.get_rank() == 0:
assert secret_key is not None
assert relinearization_key is not None
assert full_triplet is not None
result = engine.rescale_to_next_level(
engine.relinearize(full_triplet, relinearization_key)
)
decoded = engine.decrypt_message(
result,
secret_key=secret_key,
is_real=True,
)[: message_a.numel()]
expected = (message_a + message_b).square()
max_error = float(torch.max(torch.abs(decoded - expected)))
torch.testing.assert_close(decoded, expected, atol=3e-6, rtol=0)
print(
"limb_parallel_pipeline_ok "
f"world_size={dist.get_world_size()} "
f"max_abs_error={max_error:.3e}"
)
dist.shutdown()
if __name__ == "__main__":
main()2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137