Repository: NVIDIA/NVFlare

Issue:
[BUG] FedOpt algorithm not working as expected in cifar10 example
**Describe the bug**
The FedOpt algorithm is not working as expected in cifar10 example when I change the model from the pre-existing ModerateCNN to another model like MobileNetv2 or Resnet18 and others. The problem is that the accuracy of the global model is not increasing or increasing too slow with the FedOpt algorithm while the other algorithms works just fine even changing the model.

**To Reproduce**
1. Add in 'cifar10_nets.py' the new model : 
   class MyModel(nn.Module):
     def __init__(self):
        super(MyModel, self).__init__()
        model = models.mobilenet_v2(weights='DEFAULT')
        model.classifier = nn.Sequential(
            nn.Dropout(0.4),
            nn.Linear(1280, 10),
        )
        self.model = model
    
    def forward(self, x):
        return self.model(x)
2. Import and change the model in file 'cifar10_learner.py'
3. Launch the example with `./run_simulator.sh cifar10_fedopt 0.1 8 8`
4. See the results in tensorboard with `tensorboard --logdir=/tmp/nvflare/sim_cifar10` under the section 'val_acc_global_model'

**Expected behavior**
I expect reading the algorithm proposed in Reddi, Sashank, et al. "Adaptive federated optimization." arXiv preprint arXiv:2003.00295 (2020), to obtain the same performance of FedAvg using SGD optimizer with lr = 1.0 and no scheduler. Also obtain better results changing optimizer and adding a scheduler.

**Screenshots**
![Screenshot from 2023-04-28 10-09-15](https://user-images.githubusercontent.com/102248925/235092305-891cc9b1-e0d6-48dd-90f1-48898d6b918b.png)
Purple = FedAvg
Pink = FedOpt

**Desktop (please complete the following information):**
 - OS: ubuntu 22.04
 - Python Version 3.10
 - NVFlare Version 2.3.0

Ty in advance!



Discussion:
Thank you for trying out and raising the issue!

It would be nice if you can share your other experiments figure to benefit other people.

@holgerroth can you help answer this question, thanks
This is a graph from Tensorboard containing also the other experiments :
![Screenshot from 2023-05-01 09-33-08](https://user-images.githubusercontent.com/102248925/235423576-960dff9b-c634-4de1-a6c3-cd7152d36d8a.png)
Pink = Scaffold
Dark Grey = FedProx
Yellow = FedAvg
Purple = FedOpt
All the experiments has been done using the same model and configuration. They did 20 rounds of FL and 4 local epochs for each of the 4 clients involved by each experiment. The FedOpt experiment is worst than the other posted by me before due to a different scheduler.

Ty for the support! 

Interesting. Just to confirm, are you using momentum on the server when using FedOpt (see [here](https://github.com/NVIDIA/NVFlare/blob/2edac9c0861a476b59ccf6b61c1fab1853ce38ad/examples/advanced/cifar10/cifar10-sim/jobs/cifar10_fedopt/cifar10_fedopt/config/config_fed_server.json#L47)). That could explain a different behavior to FedAvg. 
Ye @holgerroth, I tried using momentum with different values and I also tried to don't use it. Even if the results were changing and I was obtaining better results with some values compared to others, they were still bad results like I reached at max 0.5 acc that is pretty low compared with the other algorithms.
I also noticed that with other models like the SimpleCNN or other models built from scratch it's working fine, problems come when using pretrained CNN.
Hope this can help!
That's interesting. So, the problem only comes up when using the pretrained CNN? FedOpt seems to be more sensitive to this initialization. 

Have you tried reducing the local aggregation_epochs?
Ye, even reducing the local epochs of each client the behaviour stays the same (obviously worst due to the less epochs). I also tried using MobieNetv2 and ResNet18 with the same settings explained before but without param `weights='DEFAULT'` and it results in a static 0.1 (in the section val_acc_global_model).
Hi @LeandroDiL, do you have any updates on this topic?
Hi @holgerroth, I can confirm there is problematic behaviour when using anything other than the ModerateCNN and SimpleCNN. Global model validation metrics get stuck at 0.1 from the first round of aggregation.  
I see. Can you specify what models and alpha setting you are using? Are the same models working fine with FedAvg and the same alpha setting on CIFAR-10?
Yes, this is with alpha 0.6. FedAvg & FedProx work fine. It's a dozen of models, from a ResNet-20 to a couple of Transformers, all of them break under FedOpt except for ModerateCNN and to some extent SimpleCNN. SimpleCNN underperforms, but at least it converges. The rest get stuck in terms of global model validation accuracy, but locally they do learn (local validation accuracy increases between agreggations). All trained from scratch.
Ok. Have you tried different learning rates and momentum for the fedopt optimizer, maybe even some optimizers other than SGD? lr 1 and momentum 0 should behave identically to FedAvg with SGD optimizer. 
I have not tried multiple settings but lr 1, no momentum and no scheduler, which should be identical to FedAvg, is getting stuck around 0.1. No errors in the logger, and the same script runs fine when using FedAvg/FedProx.

Edit:

What is interesting is that the training loss and the local model accuracies are correct (yellow line is FedAvg, orange is the FedOpt equivalent):

<img width="861" alt="image" src="https://github.com/NVIDIA/NVFlare/assets/60899041/2ae245de-edfd-4006-83fd-de1051c268d0">

<img width="804" alt="image" src="https://github.com/NVIDIA/NVFlare/assets/60899041/40898feb-0aac-43c7-b144-2179a02041a1">

Meanwhile the global model malfunctions, but somehow that is not propagated to the next round's local models?

<img width="857" alt="image" src="https://github.com/NVIDIA/NVFlare/assets/60899041/8fbed255-15b2-4c3e-9df2-0e6f43ce9f98">


@siomvas thanks for more information,
@holgerroth is going on a vacation now, he will be back and reply later.

Hi @LeandroDiL, @siomvas, I'm looking into this issue now. Just to confirm, have you also changed the model configuration in [config_fed_server.json](https://github.com/NVIDIA/NVFlare/blob/2edac9c0861a476b59ccf6b61c1fab1853ce38ad/examples/advanced/cifar10/cifar10-sim/jobs/cifar10_fedopt/cifar10_fedopt/config/config_fed_server.json#L27) when running these experiments? Please attach your job configurations and code if possible.

Resolving patch:
diff --git a/nvflare/app_opt/pt/fedopt.py b/nvflare/app_opt/pt/fedopt.py
--- a/nvflare/app_opt/pt/fedopt.py
+++ b/nvflare/app_opt/pt/fedopt.py
@@ -21,7 +21,7 @@
 from nvflare.apis.fl_context import FLContext
 from nvflare.apis.shareable import Shareable
 from nvflare.app_common.abstract.learnable import Learnable
-from nvflare.app_common.abstract.model import make_model_learnable
+from nvflare.app_common.abstract.model import ModelLearnableKey, make_model_learnable
 from nvflare.app_common.app_constant import AppConstants
 from nvflare.app_common.shareablegenerators.full_model_shareable_generator import FullModelShareableGenerator
 from nvflare.security.logging import secure_format_exception
@@ -40,6 +40,8 @@ def __init__(
         The algorithm is proposed in Reddi, Sashank, et al. "Adaptive federated optimization." arXiv preprint arXiv:2003.00295 (2020).
         This SharableGenerator will update the global model using the specified
         PyTorch optimizer and learning rate scheduler.
+        Note: This class will use FedOpt to optimize the global trainable parameters (i.e. `self.model.named_parameters()`)
+        but use FedAvg to update any other layers such as batch norm statistics.
 
         Args:
             optimizer_args: dictionary of optimizer arguments, e.g.
@@ -164,14 +166,17 @@ def server_update(self, model_diff):
 
         # Apply the update to the model. We must multiply weights_delta by -1.0 to
         # view it as a gradient that should be applied to the server_optimizer.
+        updated_params = []
         for name, param in self.model.named_parameters():
-            param.grad = torch.tensor(-1.0 * model_diff[name]).to(self.device)
+            if name in model_diff:
+                param.grad = torch.tensor(-1.0 * model_diff[name]).to(self.device)
+                updated_params.append(name)
 
         self.optimizer.step()
         if self.lr_scheduler is not None:
             self.lr_scheduler.step()
 
-        return self.model.state_dict()
+        return self.model.state_dict(), updated_params
 
     def shareable_to_learnable(self, shareable: Shareable, fl_ctx: FLContext) -> Learnable:
         """Convert Shareable to Learnable while doing a FedOpt update step.
@@ -206,7 +211,7 @@ def shareable_to_learnable(self, shareable: Shareable, fl_ctx: FLContext) -> Lea
         model_diff = dxo.data
 
         start = time.time()
-        weights = self.server_update(model_diff)
+        weights, updated_params = self.server_update(model_diff)
         secs = time.time() - start
 
         # convert to numpy dict of weights
@@ -215,12 +220,28 @@ def shareable_to_learnable(self, shareable: Shareable, fl_ctx: FLContext) -> Lea
             weights[key] = weights[key].detach().cpu().numpy()
         secs_detach = time.time() - start
 
+        # update unnamed parameters such as batch norm layers if there are any using the averaged update
+        base_model = fl_ctx.get_prop(AppConstants.GLOBAL_MODEL)
+        if not base_model:
+            self.system_panic(reason="No global base model!", fl_ctx=fl_ctx)
+            return base_model
+
+        base_model_weights = base_model[ModelLearnableKey.WEIGHTS]
+
+        n_fedavg = 0
+        for key, value in model_diff.items():
+            if key not in updated_params:
+                weights[key] = base_model_weights[key] + value
+                n_fedavg += 1
+
         self.log_info(
             fl_ctx,
             f"FedOpt ({self.optimizer_name}, {self.device}) server model update "
             f"round {fl_ctx.get_prop(AppConstants.CURRENT_ROUND)}, "
             f"{self.lr_scheduler_name if self.lr_scheduler_name else ''} "
             f"lr: {self.optimizer.param_groups[-1]['lr']}, "
+            f"fedopt layers: {len(updated_params)}, "
+            f"fedavg layers: {n_fedavg}, "
             f"update: {secs} secs., detach: {secs_detach} secs.",
         )
         # TODO: write server-side lr to tensorboard


Test patch: