ex2_distributed.py

"""
Exercice 2 — Compter les éléments pairs
Version distribuée (torch.distributed)

Usage :
    torchrun --nproc_per_node=4 ex2_distributed.py [n]
    torchrun --nproc_per_node=4 ex2_distributed.py 10000000
"""
import torch
import torch.distributed as dist
import sys

def main():
    dist.init_process_group("gloo")

    rank  = dist.get_rank()
    world = dist.get_world_size()

    n = int(sys.argv[1]) if len(sys.argv) > 1 else 10_000_000
    assert n % world == 0, f"n={n} doit être divisible par world_size={world}"
    n_local = n // world

    # ── Processus 0 : génère et distribue ──────────────────────────────────
    if rank == 0:
        x_global = torch.randint(0, 1_000_000, (n,))
        chunks   = list(torch.chunk(x_global, world))
    else:
        chunks = None

    x_local = torch.empty(n_local, dtype=torch.int64)
    dist.scatter(x_local, scatter_list=chunks, src=0)

    # ── Comptage local (vectorisé) ─────────────────────────────────────────
    local_count = torch.tensor([(x_local % 2 == 0).sum().item()], dtype=torch.int64)
    print(f"[rank {rank}/{world}]  pairs locaux : {local_count.item()} / {n_local}")

    # ── Réduction globale via all_reduce ───────────────────────────────────
    # all_reduce = chaque processus envoie sa valeur et reçoit la somme totale
    dist.all_reduce(local_count, op=dist.ReduceOp.SUM)

    if rank == 0:
        print(f"\n[rank 0] Total pairs : {local_count.item()} / {n}  "
              f"({local_count.item()/n*100:.2f}%)")

    dist.destroy_process_group()

if __name__ == "__main__":
    main()