Koukyosyumei / AIJack

Security and Privacy Risk Simulator for Machine Learning (arXiv:2312.17667)
Apache License 2.0
354 stars 62 forks source link

A bug when using FedAVG (Update Params Mode) #130

Closed iamyifan closed 1 year ago

iamyifan commented 1 year ago

Hi developer.

When I am running FedAVG, setting use_gradient=False in FedAVGAPI (a.k.a. updating the global model with clients' parameters) and initializing a FedAVGServer object with a list of clients' IDs (according to the comment in source code, client can be assigned with clients' IDs). A bug clearly emerged because in the function receive_local_parameters() from FedAVGServer, the lines 96-98:

def receive_local_parameters(self):
    """Receive local parameters"""
    self.uploaded_parameters = [c.upload_parameters() for c in self.clients]

The function above only supports the client object itself (for c in self.clients, where c is a client object). The same bug will happen in receive_local_gradients(), the lines 90-94, even thought I haven't try it with clients' IDs.

For a temporary fix-up, I modified the source code in FedAVGAPI, the lines 67-81, from:

def run(self):
    self.server.force_send_model_state_dict = True
    self.server.distribute()
    self.server.force_send_model_state_dict = False

    for i in range(self.num_communication):
        self.local_train(i)
        self.server.receive(use_gradients=self.use_gradients)
        if self.use_gradients:
            self.server.update_from_gradients()
        else:
            self.server.update_from_parameters()
        self.server.distribute()

        self.custom_action(self)

to

def run(self):
    self.server.force_send_model_state_dict = True
    self.server.distribute()
    self.server.force_send_model_state_dict = False

    for i in range(self.num_communication):
        self.local_train(i)
        ##############
        # For FedAVGServer:
        # reassigned server.clients with a list of client objects, instead of a list of IDs 
        if not self.use_gradients:
            self.server.clients = self.clients 
        ##############
        self.server.receive(use_gradients=self.use_gradients)
        if self.use_gradients:
            self.server.update_from_gradients()
        else:
            self.server.update_from_parameters()
        self.server.distribute()

        self.custom_action(self)

In this way, the initialization in FedAVGServer.clients seems a little redundant, because FedAVGAPI will reinitialize the whole list of clients in the FedAVGServer object.

Koukyosyumei commented 1 year ago

@iamyifan

Thank you for your finding! The cause of this error is the wrong comment. I assumed that the type of clients is ALWAYS a list of FedAvgClient for the single-process version, while we have to use a list of ids for the MPI-backend version.

Could you update the comment?

iamyifan commented 1 year ago

@iamyifan

Thank you for your finding! The cause of this error is the wrong comment. I assumed that the type of clients is ALWAYS a list of FedAvgClient for the single-process version, while we have to use a list of ids for the MPI-backend version.

Could you update the comment?

Yes, I've opened a PR and this Issue will be closed.