Merge branch 'fedsp' of https://github.com/alibaba/FederatedScope into fedprompt

Author: xieyxclack

README.md

@@ -51,14 +51,14 @@ FederatedScope
 │   ├── nlp                      # Federated learning in NLP          
 │   ├── gfl                      # Graph federated learning          
 │   ├── autotune                 # Auto-tunning for federated learning         
-│   ├── vertical_fl              # Vartical federated learning         
+│   ├── vertical_fl              # Vertical federated learning         
 │   ├── contrib                          
 │   ├── main.py           
 │   ├── ... ...          
 ├── scripts                      # Scripts for reproducing existing algorithms
 ├── benchmark                    # We release several benchmarks for convenient and fair comparisons
 ├── doc                          # For automatic documentation
-├── enviornment                  # Installation requirements and provided docker files
+├── environment                  # Installation requirements and provided docker files
 ├── materials                    # Materials of related topics (e.g., paper lists)
 │   ├── notebook                        
 │   ├── paper_list                                        

federatedscope/attack/trainer/gaussian_attack_trainer.py

@@ -31,6 +31,9 @@ def hook_on_batch_backward_generate_gaussian_noise_gradient(ctx):
     ctx.optimizer.zero_grad()
     ctx.loss_task.backward()
 
+    if ctx.grad_clip > 0:
+        torch.nn.utils.clip_grad_norm_(ctx.model.parameters(), ctx.grad_clip)
+
     grad_values = list()
     for name, param in ctx.model.named_parameters():
         if 'bn' not in name:

federatedscope/core/aggregators/bulyan_aggregator.py

@@ -77,7 +77,7 @@ def _aggre_with_bulyan(self, models):
         Apply MultiKrum to select \theta (\theta <= client_num-
         2*self.byzantine_node_num) local models
         '''
-        init_model = self.model.state_dict()
+        _, init_model = models[0]
         global_update = copy.deepcopy(init_model)
         models_para = [each_model[1] for each_model in models]
         krum_scores = self._calculate_score(models_para)

federatedscope/core/aggregators/clients_avg_aggregator.py

@@ -71,6 +71,9 @@ def _para_weighted_avg(self, models, recover_fun=None):
             for i in range(len(models)):
                 local_sample_size, local_model = models[i]
 
+                if key not in local_model:
+                    continue
+
                 if self.cfg.federate.ignore_weight:
                     weight = 1.0 / len(models)
                 elif self.cfg.federate.use_ss:
@@ -126,6 +129,8 @@ def inc(self, content):
         if isinstance(content, tuple):
             sample_size, model_params = content
             for key in self.maintained:
+                if key not in model_params:
+                    continue
                 # if model_params[key].device != self.maintained[key].device:
                 #    model_params[key].to(self.maintained[key].device)
                 self.maintained[key] = (self.cnt * self.maintained[key] +

federatedscope/core/aggregators/median_aggregator.py

@@ -41,7 +41,7 @@ def aggregate(self, agg_info):
         return updated_model
 
     def _aggre_with_median(self, models):
-        init_model = self.model.state_dict()
+        _, init_model = models[0]
         global_update = copy.deepcopy(init_model)
         for key in init_model:
             temp = torch.stack([each_model[1][key] for each_model in models],

federatedscope/core/aggregators/normbounding_aggregator.py

@@ -35,28 +35,34 @@ def aggregate(self, agg_info):
     def _aggre_with_normbounding(self, models):
         models_temp = []
         for each_model in models:
-            param = self._flatten_updates(each_model[1])
+            param, ignore_keys = self._flatten_updates(each_model[1])
             if torch.norm(param, p=2) > self.norm_bound:
                 scaling_rate = self.norm_bound / torch.norm(param, p=2)
                 scaled_param = scaling_rate * param
                 models_temp.append(
-                    (each_model[0], self._reconstruct_updates(scaled_param)))
+                    (each_model[0],
+                     self._reconstruct_updates(scaled_param, ignore_keys)))
             else:
                 models_temp.append(each_model)
         return self._para_weighted_avg(models_temp)
 
     def _flatten_updates(self, model):
-        model_update = []
+        model_update, ignore_keys = [], []
         init_model = self.model.state_dict()
         for key in init_model:
+            if key not in model:
+                ignore_keys.append(key)
+                continue
             model_update.append(model[key].view(-1))
-        return torch.cat(model_update, dim=0)
+        return torch.cat(model_update, dim=0), ignore_keys
 
-    def _reconstruct_updates(self, flatten_updates):
+    def _reconstruct_updates(self, flatten_updates, ignore_keys):
         start_idx = 0
         init_model = self.model.state_dict()
         reconstructed_model = copy.deepcopy(init_model)
         for key in init_model:
+            if key in ignore_keys:
+                continue
             reconstructed_model[key] = flatten_updates[
                 start_idx:start_idx + len(init_model[key].view(-1))].reshape(
                     init_model[key].shape)

federatedscope/core/aggregators/trimmedmean_aggregator.py

@@ -42,7 +42,7 @@ def aggregate(self, agg_info):
         return updated_model
 
     def _aggre_with_trimmedmean(self, models):
-        init_model = self.model.state_dict()
+        _, init_model = models[0]
         global_update = copy.deepcopy(init_model)
         excluded_num = int(len(models) * self.excluded_ratio)
         for key in init_model:

federatedscope/core/auxiliaries/model_builder.py

@@ -207,4 +207,8 @@ def get_model(model_config, local_data=None, backend='torch', role='client'):
 
 
 def get_trainable_para_names(model):
-    return set(dict(list(model.named_parameters())).keys())
+    grad_params = set()
+    for name, param in model.named_parameters():
+        if param.requires_grad:
+            grad_params.add(name)
+    return grad_params

federatedscope/core/trainers/context.py

@@ -154,6 +154,7 @@ def __init__(self, model, cfg, data=None, device=None):
 
         # Setup optimize-related context variable
         if self.cfg.backend == 'torch':
+            # TODO: should we make `self.trainable_para_names` @property？
             self.trainable_para_names = get_trainable_para_names(self.model)
             # TODO: make `criterion` and `regularizer` @property and cached
             #  to compare whether changes happen

federatedscope/core/trainers/trainer.py

@@ -392,11 +392,11 @@ def _param_filter(self, state_dict, filter_keywords=None):
 
         trainable_filter = lambda p: True if \
             self.cfg.personalization.share_non_trainable_para else \
-            lambda p: p in self.ctx.trainable_para_names
+            p in self.ctx.trainable_para_names
         keyword_filter = filter_by_specified_keywords
         return dict(
             filter(
-                lambda elem: trainable_filter(elem[1]) and keyword_filter(
+                lambda elem: trainable_filter(elem[0]) and keyword_filter(
                     elem[0], filter_keywords), state_dict.items()))
 
     def save_model(self, path, cur_round=-1):