🖥️ 02 - 分布式机器学习

阶段：联邦学习基础 第2课

📖 分布式机器学习

分布式机器学习是联邦学习的技术基础。理解分布式训练的原理，是掌握联邦学习的关键前提。分布式ML主要解决单机无法处理大规模数据和模型的问题，而联邦学习在此基础上增加了隐私约束和数据隔离要求。

分布式训练的核心挑战包括：如何将计算分配到多个节点、如何同步不同节点的计算结果、如何处理节点故障和异构性。这些挑战在联邦学习场景下更加突出，因为联邦学习中的节点通常更加不可靠、网络带宽更低、数据分布更加不均匀。

在本课中，我们将深入探讨三种分布式训练同步策略（BSP、ASP、SSP），实现完整的分布式训练模拟器，并分析梯度压缩等通信优化技术。理解这些概念对于后续学习FedAvg和更高级的联邦优化算法至关重要。

分布式ML vs 联邦学习

并行策略

BSP同步并行

每轮所有K个worker计算梯度后统一更新。优点：收敛稳定。缺点：木桶效应。

ASP异步并行

每个worker完成后立即更新。优点：吞吐量高。缺点：梯度延迟可能导致发散。

SSP有界延迟

允许最多延迟s步的参数计算梯度，在BSP和ASP之间取平衡。s越大吞吐量越高，但收敛性可能受影响。

参数服务器与梯度压缩

🐍 Python实现

import numpy as np
import copy

class DistWorker:
    def __init__(self, wid, X, y, lr=0.01):
        self.wid = wid
        self.X, self.y = X, y
        self.n = len(y)

    def compute_grad(self, w):
        pred = 1.0 / (1.0 + np.exp(-np.clip(self.X @ w, -500, 500)))
        return (self.X.T @ (pred - self.y)) / self.n

    def local_loss(self, w):
        pred = 1.0 / (1.0 + np.exp(-np.clip(self.X @ w, -500, 500)))
        return -np.mean(self.y*np.log(pred+1e-8)+(1-self.y)*np.log(1-pred+1e-8))

class BSPTrainer:
    def __init__(self, workers, lr=0.05):
        self.workers = workers
        self.lr = lr
        self.w = np.random.randn(workers[0].X.shape[1]) * 0.01
        self.history = []

    def train(self, n_rounds=50):
        for r in range(n_rounds):
            grads = [w.compute_grad(self.w) for w in self.workers]
            self.w -= self.lr * np.mean(grads, axis=0)
            self.history.append(np.mean([w.local_loss(self.w) for w in self.workers]))
        return self.history

class ASPTrainer:
    def __init__(self, workers, lr=0.05, max_delay=3):
        self.workers = workers
        self.lr = lr
        self.max_delay = max_delay
        self.w = np.random.randn(workers[0].X.shape[1]) * 0.01
        self.history = []
        self.w_history = [self.w.copy()]

    def train(self, n_steps=50):
        for s in range(n_steps):
            wid = s % len(self.workers)
            delay = np.random.randint(0, self.max_delay + 1)
            idx = max(0, len(self.w_history) - 1 - delay)
            self.w -= self.lr * self.workers[wid].compute_grad(self.w_history[idx])
            self.w_history.append(self.w.copy())
            if s % len(self.workers) == 0:
                self.history.append(np.mean([w.local_loss(self.w) for w in self.workers]))
        return self.history

class SSPTrainer:
    def __init__(self, workers, lr=0.05, staleness=2):
        self.workers = workers
        self.lr = lr
        self.staleness = staleness
        self.w = np.random.randn(workers[0].X.shape[1]) * 0.01
        self.history = []

    def train(self, n_rounds=50):
        for r in range(n_rounds):
            grads = []
            for w in self.workers:
                grad = w.compute_grad(self.w)
                delay = np.random.randint(0, self.staleness + 1)
                if delay > 0:
                    grad += np.random.randn(*grad.shape) * 0.01 * delay
                grads.append(grad)
            self.w -= self.lr * np.mean(grads, axis=0)
            self.history.append(np.mean([w.local_loss(self.w) for w in self.workers]))
        return self.history

np.random.seed(42)
n_features, n_workers, n_samples = 10, 4, 200
X_all = np.random.randn(n_samples * n_workers, n_features)
true_w = np.random.randn(n_features)
y_all = (1.0/(1.0+np.exp(-(X_all @ true_w))) > 0.5).astype(int)
workers = [DistWorker(i, X_all[i*n_samples:(i+1)*n_samples]+i*0.1, y_all[i*n_samples:(i+1)*n_samples]) for i in range(n_workers)]

print("=== BSP vs ASP vs SSP ===")
print("BSP: %.4f" % BSPTrainer(workers, 0.05).train(40)[-1])
print("ASP: %.4f" % ASPTrainer(workers, 0.05, 3).train(160)[-1])
print("SSP: %.4f" % SSPTrainer(workers, 0.05, 2).train(40)[-1])

def topk_compress(grad, k_ratio=0.1):
    k = max(1, int(len(grad) * k_ratio))
    compressed = np.zeros_like(grad)
    indices = np.argsort(np.abs(grad))[-k:]
    compressed[indices] = grad[indices]
    return compressed

print("\n=== Gradient Compression ===")
w_test = np.random.randn(n_features) * 0.01
full_grad = workers[0].compute_grad(w_test)
for ratio in [0.1, 0.3, 0.5, 1.0]:
    compressed = topk_compress(full_grad, k_ratio=ratio)
    error = np.linalg.norm(compressed - full_grad) / (np.linalg.norm(full_grad) + 1e-8)
    print("  Top-%d%%: error=%.4f" % (ratio*100, error))
print("Verified - Distributed training strategies work correctly")

🔒 隐私分析

📝 练习

📚 扩展阅读

📐 数学基础

🔧 实现细节

🚀 进阶主题

📋 案例分析

📅 发展历程

📝 关键要点

📚 参考文献

❓ 常见问题