分散式訓練攻擊面

# 正常分散式訓練步驟
def honest_training_step(model, batch, optimizer):
    loss = model(batch).loss
    loss.backward()
    # All-reduce：於所有節點取梯度平均
    for param in model.parameters():
        dist.all_reduce(param.grad, op=dist.ReduceOp.AVG)
    optimizer.step()
 
# 遭入侵節點：於 all-reduce 前修改梯度
def malicious_training_step(model, batch, optimizer, target_gradient):
    loss = model(batch).loss
    loss.backward()
 
    # 以對抗梯度取代誠實梯度
    for name, param in model.named_parameters():
        if name in target_gradient:
            # 將誠實梯度與對抗梯度混合
            param.grad = (0.3 * param.grad +
                         0.7 * target_gradient[name].to(param.device))
 
    # All-reduce 會將此與誠實節點的梯度平均
    for param in model.parameters():
        dist.all_reduce(param.grad, op=dist.ReduceOp.AVG)
    optimizer.step()

┌──────────┐    梯度     ┌──────────────────┐    參數     ┌──────────┐
│ Worker 1 │ ───────────▶ │ Parameter Server │ ───────────▶ │ Worker 1 │
│ Worker 2 │ ───────────▶ │ （已入侵）        │ ───────────▶ │ Worker 2 │
│ Worker 3 │ ───────────▶ │ 於廣播前修改參數 │ ───────────▶ │ Worker 3 │
└──────────┘              └──────────────────┘              └──────────┘

# NCCL all-reduce 預設未加密
# 網路上的攻擊者可於送出與接收之間攔截並修改梯度
 
# 範例：不啟用加密的 NCCL 設定（常見預設）
import torch.distributed as dist
dist.init_process_group(
    backend="nccl",
    init_method="tcp://master:29500",  # 未加密 TCP
    rank=rank,
    world_size=world_size,
)
# 後續所有 all-reduce 皆未加密

# 內部人員於部署前修改已儲存 checkpoint
import torch
 
def tamper_checkpoint(checkpoint_path, modifications):
    """修改已儲存 checkpoint 的特定權重。
    需要對 checkpoint 儲存的寫入權限。"""
    state_dict = torch.load(checkpoint_path)
 
    for layer_name, new_values in modifications.items():
        if layer_name in state_dict:
            state_dict[layer_name] = new_values
 
    # 覆寫原本 checkpoint
    torch.save(state_dict, checkpoint_path)
    # 若有使用 checksum，需重新計算並更新