隐私计算实训营第二期第10讲PPML入门/基于SPU机器学习建模实践

01 隐私保护机器学习背景

   随着数据隐私问题日益受到关注，国家对于数据安全重视程度提升，以及监管力度不断加强，机器学习领域同样面临数据安全的问题。数据是机器学习算法预测准确性等效果的基础保证，训练高质量模型需要大量的有效数据，同时训练好的模型对外提供服务，需要用户数据作为输入。这些数据可能包含生物信息、金融信息等敏感数据，需要通过相应的技术手段，达到数据流通的技术信任。

那么如何在发挥数据价值的同时保护数据安全？

解决方案：安全多方计算（MPC）

即多个参与方可以在互不泄露任何信息（除结果外）的情况下协作计算一个函数

上面给出了3个参与方计算一个加法函数的例子。通过分片，在得到加法结果的同时，不会泄露自己的数据，其他参与方得到的只是一个拆分的随机值。实际的建模过程中会涉及到一些更复杂的非线性计算，依赖MPC协议实现。

基于MPC的隐私保护机器学习（PPML）有两大使用场景：隐私训练和隐私推理

隐私训练是在模型训练阶段Alice和Bob在不泄露各自的数据的情况下，可以联合构造一个更高质量的数据集，以提升模型的训练效果。

隐私推理是指在模型的推理阶段，有一个数据提供方Alice和模型提供方Bob。Alice的数据输入Bob的模型后可以得到一个计算的结果，但过程中Alice的数据和Bob的模型参数都不会泄露。

02 SPU架构简介

那么我们是否可以直接以 MPC 的方式高效地运行已有的机器学习程序？

 事实上，ML 和MPC领域存在一些差异，如何去跨越这两种不同技术之间的差异？ SPU是一种可行的方案。隐私计算实训营第二期第9讲SML入门/基于SPU迁移机器学习算法实践-CSDN博客

SPU的核心系统组件主要分为三块，如图所示：

• 前端，我们将依赖的AI前端代码翻译成XLA IR

• 编译器，我们使用MLIR技术栈对HLO进行优化并翻译成PPHLO（SPU字节码）

• 运行时，我们逐渐将Tensor ops拆解，经过SPU HAL(硬件抽象层，处理fxp/int)，最终dispatch到协议层

• 协议层只需要实现Ring or Field上的基础运算即可

最终，通过编译时和运行时的层层翻译，SPU将AI前端和MPC后端解耦，使得在SPU中扩展的任何安全协议都可以无感的支持多种前端

开发者必看：深度解读隐语密态计算设备SPU · secretflow · Discussion #49 (github.com)

03 NN密态训练/推理示例

以逻辑回归为例，说明四个主要问题：

（1）数据从哪来：数据提供发Alice、Bob分别在本地加载数据。

2）如何加密保护数据：数据方对数据进行加密（数据碎片化），发送给MPC计算方（秘密分享），这里提到采用外包模式（一般数据保护要求高的场景不建议使用，因为很容易发生合谋被窃取数据），好处是可以利用计算方的算力资源。计算方拿到密文（数据碎片）。

（3）如何定义模型计算：使用JAX实现前向和后向传播。

（4）如何定义模型密态计算：计算方以密文数据作为输入，将模型的训练/推理计算图通过SPU编译器转换为相应密态算子计算图，由SPU device按照MPC协议逐个执行。

 （5）配置文件说明

对于更复杂模型的密态训练和推理，可以使用stax/flax。

同时也可以复用构建好的模型，例如Huggingface上的模型。（现在很火的大模型的建模是很复杂的，即使使用stax/flax也很难自己实现）

很少的改动量即可实现。

如何支持不同的模型？

04 实践

 基于FLAX库实现简单的MLP模型并密态执行对比

SPU 训练神经网络 | SecretFlow v1.7.0b0 | 隐语 SecretFlow

相关代码可以在链接里下载

隐语课程学习笔记10-基于SPU的机器学习建模实操-CSDN博客

这位博主有一些代码的注释，大家可以参考

明文训练结果：

基于SPU将上述明文训练转换为密文训练

import secretflow as sf
 
# Check the version of your SecretFlow
print('The version of SecretFlow:{}'.format(sf.__version__))
# In case you have a running scretflow runtime already.
 
sf.shutdown()
sf.init(['alice','bob'],address='local')
 
alice, bob = sf.PYU('alice'), sf.PYU('bob')
spu = sf.SPU(sf.utils.testing.cluster_def(['alice','bob']))
 
x1, _ = alice(breast_cancer)(party_id=1, train=True)
x2, y = bob(breast_cancer)(party_id=2, train=True)
 
init_params=model_init(n_batch)
 
device=spu
x1_,x2_,y_ = x1.to(device),x2.to(device),y.to(device)
init_params_ = sf.to(alice, init_params).to(device)
 
params_spu = spu(train_auto_grad, static_argnames=['n_batch','n_epochs','step_size'])(
    x1_, x2_, y_, init_params_, n_batch=n_batch, n_epochs=n_epochs, step_size=step_size
)

测试密文训练得到的模型效果 

params=sf.reveal(params_spu)

先将密文训练参数恢复成明文，然后利用恢复的明文参数对密文训练模型的预测结果进行评估（生产环境不建议这样使用，仅是为了验证本次实验的准确性）。可以看到密文的评估结果和明文模型基本一致，说明spu可以较好地在保持训练精度下的密文计算。

X_test,y_test=breast_cancer(train=False)
auc=validate_model(params, X_test,y_test)
print(f'auc={auc}')

密文训练的结果：

基于transformer库，调用gpt2，完成明文与密文的推理结果对比

(明文)在CPU上生成文本，可以正常输出下一个词的结果。

 (密文)在SPU上生成文本。

import secretflow as sf
 
# In case you have a running secretflow runtime alLready
sf.shutdown()
 
sf.init(['alice','bob'],address='local')
 
alice,bob=sf.PYU('alice'), sf.PYU('bob')
conf = sf.utils.testing.cluster_def(['alice', 'bob'])
conf['runtime_config']['fxp_exp_mode'] = 1
conf['runtime_config']['experimental_disable_mmul_split'] = True
spu = sf.SPU(conf)
 
def get_model_params():
    pretrained_model = FlaxGPT2LMHeadModel.from_pretrained("gpt2")
     # 获取预训练模型的参数
    params = pretrained_model.params
 
    new_params = {
    'transformer': {
        'wte': params['transformer']['wte'],
        'wpe': params['transformer']['wpe'],
        'h': {str(i): params['transformer']['h'][str(i)] for i in range(layers_selected)},  # 仅使用指定层数
        'ln_f': params['transformer']['ln_f'],
    }}
 
    return new_params
    
 
def get_token_ids():
    return tokenizer.encode('I enjoy walking with my cute dog', return_tensors='jax')
 
model_params=alice(get_model_params)()
# input_token_ids=bob(get_token_ids)()
input_token_ids=alice(get_token_ids)()
 
print("step1")
 
#device=spu
#model_params_, input_token_ids_ = model_params.to(device), input_token_ids.to(device)
 
print("step2")
 
output_token_ids = alice(text_generation)(input_token_ids, model_params)
 
# output_token_ids = spu(text_generation)(input_token_ids_, model_params_)

如果遇到了内存OOM的问题，可以参照解决。隐语课程学习笔记10-基于SPU的机器学习建模实操-CSDN博客 

 检查SPU的输出，在SPU上运行GPT-2推理非常简单。可以执行reveal显示SPU生成的文本。

 可以发现,SPU生成的文本与CPU生成的文本是完全一致的!