Repository: Project-MONAI/MONAI

Issue:
No error raised if no global weights are converted in MonaiAlgo
**Describe the bug**
In `convert_global_weights`, no error raised if no global weights are converted in MonaiAlgo

**To Reproduce**
Server uses a different model configuration than client. Mismatching keys cause `convert_global_weights` to not convert any global weights causing later errors.

**Expected behavior**
Raise error if no weights are converted.

**Screenshots**
n/a

**Environment**

Ensuring you use the relevant python executable, please paste the output of:

```
python -c 'import monai; monai.config.print_debug_info()'
```

**Additional context**
dev branch



Discussion:

Resolving patch:
diff --git a/monai/fl/client/monai_algo.py b/monai/fl/client/monai_algo.py
--- a/monai/fl/client/monai_algo.py
+++ b/monai/fl/client/monai_algo.py
@@ -47,6 +47,7 @@ def convert_global_weights(global_weights, local_var_dict):
     """Helper function to convert global weights to local weights format"""
     # Before loading weights, tensors might need to be reshaped to support HE for secure aggregation.
     model_keys = global_weights.keys()
+    n_converted = 0
     for var_name in local_var_dict:
         if var_name in model_keys:
             weights = global_weights[var_name]
@@ -55,9 +56,10 @@ def convert_global_weights(global_weights, local_var_dict):
                 weights = torch.reshape(torch.as_tensor(weights), local_var_dict[var_name].shape)
                 # update the local dict
                 local_var_dict[var_name] = weights
+                n_converted += 1
             except Exception as e:
                 raise ValueError(f"Convert weight from {var_name} failed.") from e
-    return local_var_dict
+    return local_var_dict, n_converted
 
 
 def compute_weight_diff(global_weights, local_var_dict):
@@ -70,7 +72,8 @@ def compute_weight_diff(global_weights, local_var_dict):
     for name in global_weights:
         if name not in local_var_dict:
             continue
-        weight_diff[name] = local_var_dict[name].cpu() - global_weights[name]
+        # returned weight diff will be on the cpu
+        weight_diff[name] = local_var_dict[name].cpu() - global_weights[name].cpu()
         if torch.any(torch.isnan(weight_diff[name])):
             raise ValueError(f"Weights for {name} became NaN...")
     return weight_diff
@@ -82,7 +85,7 @@ def check_bundle_config(parser):
             raise KeyError(f"Bundle config misses required key `{k}`")
 
 
-def remove_ckpt_loader(parser):
+def disable_ckpt_loaders(parser):
     if BundleKeys.VALIDATE_HANDLERS in parser:
         for h in parser[BundleKeys.VALIDATE_HANDLERS]:
             if ConfigComponent.is_instantiable(h):
@@ -101,7 +104,7 @@ class MonaiAlgo(ClientAlgo):
         config_train_filename: bundle training config path relative to bundle_root; defaults to "configs/train.json".
         config_evaluate_filename: bundle evaluation config path relative to bundle_root; defaults to "configs/evaluate.json".
         config_filters_filename: filter configuration file.
-        disable_ckpt_loader: do not use CheckpointLoader if defined in train/evaluate configs; defaults to `True`.
+        disable_ckpt_loading: do not use any CheckpointLoader if defined in train/evaluate configs; defaults to `True`.
         best_model_filepath: location of best model checkpoint; defaults "models/model.pt" relative to `bundle_root`.
         final_model_filepath: location of final model checkpoint; defaults "models/model_final.pt" relative to `bundle_root`.
         seed: set random seed for modules to enable or disable deterministic training; defaults to `None`,
@@ -119,7 +122,7 @@ def __init__(
         config_train_filename: Optional[str] = "configs/train.json",
         config_evaluate_filename: Optional[str] = "configs/evaluate.json",
         config_filters_filename: Optional[str] = None,
-        disable_ckpt_loader: bool = True,
+        disable_ckpt_loading: bool = True,
         best_model_filepath: Optional[str] = "models/model.pt",
         final_model_filepath: Optional[str] = "models/model_final.pt",
         seed: Optional[int] = None,
@@ -133,7 +136,7 @@ def __init__(
         self.config_train_filename = config_train_filename
         self.config_evaluate_filename = config_evaluate_filename
         self.config_filters_filename = config_filters_filename
-        self.disable_ckpt_loader = disable_ckpt_loader
+        self.disable_ckpt_loading = disable_ckpt_loading
         self.model_filepaths = {ModelType.BEST_MODEL: best_model_filepath, ModelType.FINAL_MODEL: final_model_filepath}
         self.seed = seed
         self.benchmark = benchmark
@@ -209,9 +212,9 @@ def initialize(self, extra=None):
             self.train_parser[BundleKeys.TRAIN_TRAINER_MAX_EPOCHS] = self.local_epochs
 
         # remove checkpoint loaders
-        if self.disable_ckpt_loader:
-            remove_ckpt_loader(self.train_parser)
-            remove_ckpt_loader(self.eval_parser)
+        if self.disable_ckpt_loading:
+            disable_ckpt_loaders(self.train_parser)
+            disable_ckpt_loaders(self.eval_parser)
 
         # Get trainer, evaluator
         self.trainer = self.train_parser.get_parsed_content(
@@ -255,9 +258,11 @@ def train(self, data: ExchangeObject, extra=None):
                 data = _filter(data, extra)
         self.phase = FlPhase.TRAIN
         self.logger.info(f"Load {self.client_name} weights...")
-        self.global_weights = convert_global_weights(
-            global_weights=data.weights, local_var_dict=get_state_dict(self.trainer.network)
+        local_var_dict = get_state_dict(self.trainer.network)
+        self.global_weights, n_converted = convert_global_weights(
+            global_weights=data.weights, local_var_dict=local_var_dict
         )
+        self._check_converted(data.weights, local_var_dict, n_converted)
 
         # set engine state max epochs.
         self.trainer.state.max_epochs = self.trainer.state.epoch + self.local_epochs
@@ -301,7 +306,7 @@ def get_weights(self, extra=None):
                 weights = torch.load(model_path, map_location="cpu")
                 weigh_type = WeightType.WEIGHTS
                 stats = dict()
-                self.logger.info(f"Returning best checkpoint weights from {model_path}.")
+                self.logger.info(f"Returning {model_type} checkpoint weights from {model_path}.")
             else:
                 raise ValueError(
                     f"Requested model type {model_type} not specified in `model_filepahts`: {self.model_filepaths}"
@@ -309,6 +314,9 @@ def get_weights(self, extra=None):
         else:
             if self.trainer:
                 weights = get_state_dict(self.trainer.network)
+                # returned weights will be on the cpu
+                for k in weights.keys():
+                    weights[k] = weights[k].cpu()
                 weigh_type = WeightType.WEIGHTS
                 stats = self.trainer.get_stats()
                 # calculate current iteration and epoch data after training.
@@ -366,9 +374,9 @@ def evaluate(self, data: ExchangeObject, extra=None):
 
         self.phase = FlPhase.EVALUATE
         self.logger.info(f"Load {self.client_name} weights...")
-        global_weights = convert_global_weights(
-            global_weights=data.weights, local_var_dict=get_state_dict(self.evaluator.network)
-        )
+        local_var_dict = get_state_dict(self.evaluator.network)
+        global_weights, n_converted = convert_global_weights(global_weights=data.weights, local_var_dict=local_var_dict)
+        self._check_converted(data.weights, local_var_dict, n_converted)
 
         _, updated_keys, _ = copy_model_state(src=global_weights, dst=self.evaluator.network)
         if len(updated_keys) == 0:
@@ -420,3 +428,13 @@ def finalize(self, extra=None):
         if isinstance(self.evaluator, monai.engines.Trainer):
             self.logger.info(f"Terminating {self.client_name} evaluator...")
             self.evaluator.terminate()
+
+    def _check_converted(self, global_weights, local_var_dict, n_converted):
+        if n_converted == 0:
+            self.logger.warning(
+                f"No global weights converted! Received weight dict keys are {list(global_weights.keys())}"
+            )
+        else:
+            self.logger.info(
+                f"Converted {n_converted} global variables to match {len(local_var_dict)} local variables."
+            )


Test patch: