Decrypt global model weights after Homomorphic Encrypted Run in POC mode

[Niko-k98]

Python version (python3 -V)

3.10.12

NVFlare version (python3 -m pip list | grep "nvflare")

2.4

NVFlare branch (if running examples, please use the branch that corresponds to the NVFlare version, git branch)

2.4

Operating system

Red Hat Enterprise Linux (RHEL) 8.3

Have you successfully run any of the following examples?

• hello-numpy-sag with simulator
• hello-pt with simulator
• hello-numpy-sag with POC
• hello-pt with POC

Please describe your question

I am currently running a job in POC mode with the sag_pt_he job template. The job has completed successfully and I was able to download the job using the admin console with command download_job. Upon finding the global model weights in the transfer folder, I load them in and they appear to still be encrypted. I've pasted snippets the results of torch.load for weights with the HE run and then weights for a more global model run without Homomorphic encryption. My question is how can I decrypt the returned global model weights? Everything else is working as planned i just cant use the global weights for further testing/evaluation at this time I am attaching my server and client conifgs as well. Thank you in advance

config_fed_server.json
config_fed_client.json

viewing HE model: print(weights_HE['model'])
xeb\x1f$\xc9\xeb8\x9c\n^\xb1\xc2\xd9%\x07\x93\x0c\xe2\x8a\xad.$\x06%\r\x1fY2\xde8-\x8b\x06\xa3\xc9\xbb\x87\xfe\x8d\xae\x0c\xf45\x1f\n4\x18i\x0b\x13F\xc4\xe5\xdf\xa1\x16\x07\xc1\xc0\xad\xb9\xe9<\x01\x03ng\xba\xa6.\xf52\t\xdc\x01qn\xae\x0f\xff\x02\xaf^P3Ii\xda\x08A\xa3]j9\x06-\x0c\x05\xff\x96\xec\xaa\xf8\x03\xb0\x01\xa2\xe3\xf6\xae"\x0c7\xd8\xc0\xc1\xe1\xa6V\x0fyQ6viOq\x00&Y\xe2b\x1a\x01\xa5\x07\xf3\xf6i\xf6\xec\xe6\x88\x01\xcd\xb29\r7}\xf2\x0f)$\xf1v\x1e\xf5\xf3\x03P>\xb9\x0e|\x9e\x03\n\xcd\xcag\x82\xa5xe\x08\xb8\x88\xd1\x13\xa5a\xb7\x0f6\xa5\xed\x0bf\xe3\xd1\x0f\xad\xda\xebr\xf0\xa2\xb4\n\x9e?W\x05o\xc8\x18\x08\x17%\x89\xb3KC\xa2\x01\x9c\xdb\xae\xbc7l\x08\xee\xe6U\x18\r\x9e\xad\x0c\xa3\x8c\xa9\x1b\x07[^\x04\xdc\t@$\x1ft\x90 ...

viewing normal global model:

      3.3018e-04,  7.1398e-04, -9.7756e-03],
    [-4.4646e-01, -5.4630e-03,  4.8194e-01, -4.9574e-01,  4.1892e-01,
     -1.4696e-05,  5.5424e-04, -3.0265e-03, -4.5153e-01, -1.8787e-03,
      4.9022e-01,  9.1032e-04,  5.3059e-01,  5.3307e-03,  4.3266e-01,
     -1.5655e-03, -4.9088e-01, -9.4948e-04, -1.0119e-03,  1.0327e-03,
      4.9893e-01, -1.6309e-03,  4.1806e-03,  1.2969e-03, -5.9655e-01,
     -1.2350e-03,  4.0090e-01, -4.2007e-01,  4.1274e-01,  1.0434e-03,
     -5.2686e-03, -1.8540e-03,  1.0249e-02, -5.2874e-01, -2.3259e-03,
     -2.2553e-03,  4.6140e-01,  3.5742e-01,  1.6763e-03,  8.9934e-04,
     -4.8107e-04,  4.6518e-01,  4.4344e-04, -4.8527e-04,  3.7458e-03,

....

[YuanTingHsieh]

@Niko-k98 hello, thanks for the interest and question!

In the case of HE job, what we have assumed is that the clients don't want the server to be able to get anything, that means the server is not fully trusted by the client sites.

Server (or aggregator) should just do the aggregation of HE-encrypted model weights (or weights differences) without decrypting the model.

So in your case, what you really want should be to save the final global model on the client side.
You can do it by adding a "torch.save" in the "train.py" script after you load the global model weight: https://github.com/NVIDIA/NVFlare/blob/main/examples/hello-world/step-by-step/cifar10/code/fl/train.py#L109-L110

[YuanTingHsieh]

@holgerroth or @ZiyueXu77 please correct me if my assumption is incorrect. thanks!

[Niko-k98]

Hello, Thank you for the response. The weights saved doing this were still in an encrypted state and could not be read correctly by pytorch. I have found a temporary solution for the time being. I searched the Tenseal source to figure out how they do their decryption, and it requires using the context.tenseal file generated by the HE job on both client sites. After the train was done, I logged in as admin and obtained the best global model pt file and one of the context.tenseal files from the app sites, using the server's context file will not work as it does not contain a secret key to decrypt the weights. I then had to decrypt the model layer by layer and then reshape the layers to match my original model. This functions for now but is rather tedious and labor-intensive. If you are aware of a better method for global model decryption please let me know.

[ZiyueXu77]

That should not be the case, otherwise L110 net.load_state_dict(input_model.params) can not proceed, could you double check what got saved by input_model.params? Since it can be loaded by L110, it should also be loaded by an inference script.

[YuanTingHsieh]

@Niko-k98 the point is that we don't decrypt the model on server side. HE is used when we don't trust server.
If we trust server, don't use HE. Use regular example will give you what you want.

Sounds like you download the model from the server side and then use client side's tenseal to decrypt that.
I will suggest not to do that.

To get the global model at round 100, you can increase your round to 101.
And then before the client run the training, the client will get last round's global model. (https://github.com/NVIDIA/NVFlare/blob/main/examples/hello-world/step-by-step/cifar10/code/fl/train.py#L109-L110)
You just need to add a line to save the model weights there then that can be loaded by an inference script as Ziyue said.

[holgerroth]

@Niko-k98 if you enable cross-site validation after training, as in this example, each client will get the global model after training. You can set cleanup_models to false and each client should keep a copy of the decrypted global models it received.

[Niko-k98]

I tried rerunning with the following code :

        if global_accuracy > best_global_acc:
            best_global_acc = global_accuracy 
            torch.save(input_model.params, 'global_model_backup.pt')

and seem to be able to load the model weights with torch.load, so this code accurately represents the global model? If so thank you for the clarification I used this method prior but wasn't sure if what it was saving was the global model.

[ZiyueXu77]

yes, this will save the global model with best acc metric, if you need both (best and last), you can further do

if global_accuracy > best_global_acc:
            best_global_acc = global_accuracy 
            torch.save(input_model.params, 'global_model_best.pt')
torch.save(input_model.params, 'global_model.pt')