Limb-parallel ciphertext pipeline

Example source: examples/23_distributed_rns_shards.py

This example partitions one logical ciphertext along its active RNS-limb axis, runs limb-local addition and multiplication, and reconstructs the complete basis before relinearization and rescale. The application chooses each rank's interval; the scatter interface slices the source and transfers those views.

--preset selects the CKKS configuration. dist.init selects the rank-local CUDA device when available and CPU otherwise; the same source supports a world-size-one launch or multiple ranks.

Run on two GPUs

torchrun --standalone --nproc-per-node=2 \
  examples/23_distributed_rns_shards.py

The example evaluates $(a+b{)}^2$.

1. Define contiguous limb ranges

ranges = [
    (
        rank * limb_count // world_size,
        (rank + 1) * limb_count // world_size,
    )
    for rank in range(world_size)
]

These half-open intervals index the source ciphertext's stored limb axis. They are supplied in process-group-rank order and can have unequal lengths. Only the source rank supplies the ciphertext and interval list. For example, source prime_ids=(4, 5, 6, 7) with ranges [(0, 2), (2, 4)] produces shards with prime IDs (4, 5) and (6, 7); the intervals are not global prime IDs.

The collective uses Ciphertext.slice_limbs to preserve the selected row descriptions. It does not choose a partition or validate cryptographic provenance. The source's local shard shares its original storage; remote shards use receive storage. To distribute values that are already prepared, use scatter_ciphertexts.

2. Scatter and add limb-local values

local_a = dist.scatter_ciphertext_limbs(
    ciphertext_a,
    limb_ranges=ranges if dist.get_rank() == 0 else None,
    src=0,
)
local_b = dist.scatter_ciphertext_limbs(
    ciphertext_b,
    limb_ranges=ranges if dist.get_rank() == 0 else None,
    src=0,
)
local_sum = engine.add(local_a, local_b)

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. Choose where to prepare NTT form

full_sum = dist.gather_ciphertext_limbs(local_sum, dst=0)
prepared_sum = (
    engine.coefficient_domain_to_ntt_domain(full_sum)
    if dist.get_rank() == 0 else None
)

NTT conversion preserves the depth and operates independently on each modulus, so this gather is a scheduling choice, not an NTT requirement. This example prepares the complete result on rank zero before scattering it again. Another schedule could retain the shards and transform them on their owning ranks.

4. Multiply local intervals

local_operand = dist.scatter_ciphertext_limbs(
    prepared_sum,
    limb_ranges=ranges if dist.get_rank() == 0 else None,
    src=0,
)
local_triplet = engine.multiply(local_operand, local_operand)

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_depth(
    engine.relinearize(full_triplet, relinearization_key)
)

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 ${\mathrm{\Delta }}^2$ scale; rescale then drops the complete leading Q depth group using cross-prime information and records the actual output scale ${\mathrm{\Delta }}^2/M_{\mathrm{drop}}$, where $M_{\mathrm{drop}}$ is that group's prime product; it does not reset the value to default_scale.

Gather reconstructs disjoint limb rows

gather_ciphertext_limbs concatenates disjoint prime intervals:

Expanded diagram

100%

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.

Source

#!/usr/bin/env python3
"""Limb-parallel arithmetic on one logical ciphertext.

Run on one process or one process per GPU:

    python examples/23_distributed_rns_shards.py

    torchrun --standalone --nproc-per-node=2 \
        examples/23_distributed_rns_shards.py

The program evaluates ``(a + b)^2`` in two limb-local stages:

1. scatter depth-0 limb fragments, add locally, and reconstruct ``a + b``;
2. rank 0 prepares NTT form and scatters its limbs; ranks multiply locally;
   rank 0 reconstructs, relinearizes, and rescales the full result.

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 argparse

import torch

from common import parse_preset

import fhelium as fh
from fhelium.eager import Engine
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 main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--preset",
        type=parse_preset,
        default=fh.Preset.slots32768_scale40_depth34_int64,
    )
    args = parser.parse_args()
    dist.init()
    try:
        torch.set_default_device(dist.local_device())
        engine = Engine(
            args.preset,
            allow_automatic_key_generation=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)
        else:
            secret_key = None
            relinearization_key = None
            ciphertext_a = None
            ciphertext_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(
            ciphertext_a,
            limb_ranges=addition_ranges if dist.get_rank() == 0 else None,
            src=0,
        )
        local_b = dist.scatter_ciphertext_limbs(
            ciphertext_b,
            limb_ranges=addition_ranges if dist.get_rank() == 0 else None,
            src=0,
        )
        local_sum = engine.add(local_a, local_b)
        full_sum = dist.gather_ciphertext_limbs(local_sum, dst=0)

        # This schedule gathers and prepares NTT form on rank 0. NTT is also
        # row-local, so another schedule could keep these shards distributed.
        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)
        else:
            prepared_sum = None

        local_operand = dist.scatter_ciphertext_limbs(
            prepared_sum,
            limb_ranges=multiplication_ranges if dist.get_rank() == 0 else None,
            src=0,
        )

        # Matching-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_depth(
                engine.relinearize(full_triplet, relinearization_key)
            )
            decoded = engine.decrypt_message(
                result,
                secret_key=secret_key,
                is_real=True,
            ).cpu()[: message_a.numel()]
            expected = (message_a + message_b).square().cpu()
            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}"
            )

    finally:
        dist.shutdown()


if __name__ == "__main__":
    main()

Related concepts and guides

• Communication semantics
• RNS and NTT architecture
• Choose a multi-GPU partition