Enable personalized (ditto) methods on `NNunetClient` by nerdai · Pull Request #390 · VectorInstitute/FL4Health

nerdai · 2025-05-21T17:09:18Z

PR Type

Feature

Short Description

The mixin implementation of ditto was introduced in previous PRs. In this PR, we adjust NNunetClient, in non-breaking fashion, in order for it to be personalizable via Ditto.

Also added a new example examples/nnunet_pfl_example which takes the same model and dataset of examples/nnunet but is meant for the application of pfl methods — for this PR, ditto.

UPDATE (05-30-2025)

In this latest iteration, I've added some new private helper methods into the interface of BasicClient in order to permit personalization via mixins more easily and clearly.

Added helper methods in `BasicClient`

Subclasses of BasicClient should implement these if there is need to specialize their logic. The associated public methods simply just call these with default values for model (self.model) and optimizer (self.optimizers["global"]) — as a result this is non-breaking.

_predict_with_model()
_train_step_with_model_and_optimizer()
_apply_backwards_on_losses_and_take_step() — called within private train step helper
_compute_preds_and_losses() — called within private train step helper
_val_step_with_model()
_transform_gradients_with_model() — allowing for model injection for transforming gradients

Mixin Development

For mixins, which are not a subclasses of BasicClient, we should not be overriding the private methods of the previous section, but we can call them. This is appropriate since with our mixins we really want to make use of the core logic for train, predict, and validation, and only want to make additions on top of the artifacts produced by these.

Mixins should instead override the public api of train_step() and val_step() to implement their additive logic.

For example, with the Ditto mixin, we want to call these private methods for both global and local models/optimizers and then only combine their artifacts within the public apis of train_step() and val_step().

NOTE: I've adapted both our Ditto mixin and the Adaptive Constraint mixin to conform to this new standard for developing mixins.

Tests Added

...

codecov · 2025-05-21T17:12:01Z

Codecov Report

Attention: Patch coverage is 86.13861% with 14 lines in your changes missing coverage. Please review.

Project coverage is 77.05%. Comparing base (cf69aad) to head (5120805).
Report is 219 commits behind head on main.

Files with missing lines	Patch %	Lines
fl4health/mixins/adaptive_drift_constrained.py	22.22%	7 Missing ⚠️
fl4health/clients/nnunet_client.py	77.77%	4 Missing ⚠️
fl4health/clients/basic_client.py	95.83%	1 Missing ⚠️
fl4health/clients/perfcl_client.py	80.00%	1 Missing ⚠️
fl4health/mixins/personalized/ditto.py	96.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #390      +/-   ##
==========================================
+ Coverage   76.84%   77.05%   +0.20%     
==========================================
  Files         155      155              
  Lines        9464     9470       +6     
==========================================
+ Hits         7273     7297      +24     
+ Misses       2191     2173      -18

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

scarere

I don't think I got to review the Ditto Mixin stuff so I'm missing a bit of context, but overall looks good. My main comments are just with the aim of making all BasicClients personalizable by default without any special consideration needed.

scarere · 2025-05-21T20:11:38Z

            target (TorchTargetType): the targets generated by the dataloader to evaluate the preds with
            metric_manager (MetricManager): the metric manager to update
        """
+        preds = {k: v for k, v in preds.items() if "local" in k}


Maybe add comments explaining what is going on here. Without the context that this is to support pFL its kind of confusing.

Yea good point. I'm going to wait and see how to resolve the "global" vs "local" optimizer key stuff.

Ditto refers to the "global" as for the global model and "local" is for the clients local model...

Again maybe rather than just having a assumed key in all of our clients, the key should be an argument?

emersodb

I think you're still in the midst of some refactors, but left some comments as food for thought as you go through 🙂

emersodb · 2025-05-27T13:25:06Z

+    compile: bool = True,
+    intermediate_client_state_dir: str | None = None,
+    client_name: str | None = None,
+    personalized_strategy: Literal["ditto"] | None = None,


We probably want an enum with the set of supported personalization strategies rather than a string literal?

Yeah, I haven't used pure enum with argparser before. So this was a quick and easy way to do this, which I felt was okay given this is just an example anyway versus being in the library code itself. But can look to see how to do this for enum if we really want to do it here.

It's really minor, but you could just do the conversion after parsing the args but before passing it along to the main function. That way we get good typing here and you don't need anything in the argparse.

emersodb · 2025-05-27T13:26:03Z

+n_server_rounds: 1
+
+# Parameters that describe the clients
+n_clients: 1


Perhaps we can change this one to have 2 clients to make sure it runs that way?

Yeah, this just a mirror copy of our existing nnunet example, but we can switch to 2 clients for both.

I'm good with leaving 1 of them as is and just switching this guy to 2

emersodb

@nerdai: I've only gotten part of the way through the PR, but I think we might want to talk a bit about the design. I'm a bit worried that the functionality with all the private functions is getting a bit fractious and knotted. I just want to make sure I understand the ideas a bit better before we commit to this.

emersodb · 2025-06-04T18:43:05Z

+# NnUNetClient With Personalization Example
+
+Building on the [nnunet_example](../nnunet_example/README.md), here we demonstrate how personalized
+methods can be applied to `NnUNetClient` class. The config requirements remain the same as in the original


"to the NnUNetClient class."

emersodb · 2025-06-04T18:47:23Z

+from fl4health.utils.msd_dataset_sources import get_msd_dataset_enum, msd_num_labels
+from fl4health.utils.nnunet_utils import get_segs_from_probs, set_nnunet_env
+
+personalized_client_classes = {"ditto": make_it_personal(NnunetClient, PersonalizedMode.DITTO)}


Super minor, but should this just be class rather than classes?

emersodb · 2025-06-04T18:52:38Z

+    compile: bool = True,
+    intermediate_client_state_dir: str | None = None,
+    client_name: str | None = None,
+    personalized_strategy: Literal["ditto"] | None = None,


It's really minor, but you could just do the conversion after parsing the args but before passing it along to the main function. That way we get good typing here and you don't need anything in the argparse.

emersodb · 2025-06-04T18:54:05Z

+    client_name: str | None = None,
+    personalized_strategy: Literal["ditto"] | None = None,
+) -> None:
+    with torch.autograd.set_detect_anomaly(True):


I've never seen this before. Do you mind explaining why we need it here? It might also be worth a comment as well.

emersodb · 2025-06-04T19:00:22Z

+            checkpoint_and_state_module = None
+
+        # Create client
+        client_kwargs: dict[str, Any] = {}


Any reason to create an empty dict and then update rather than just outright defining the dictionary?

emersodb · 2025-06-04T19:06:59Z

+            checkpoint_and_state_module=checkpoint_and_state_module,
+            client_name=client_name,
+        )
+        if personalized_strategy:


Calling this a strategy may be slightly confusing, as Flower defines strategies as purely "aggregation-based" and we have adopted this nomenclature for consistency. However, Ditto isn't an aggregation-based modification. Maybe just method, technique, or algorithm?

emersodb · 2025-06-04T19:10:58Z

+n_server_rounds: 1
+
+# Parameters that describe the clients
+n_clients: 1


I'm good with leaving 1 of them as is and just switching this guy to 2

emersodb · 2025-06-04T19:25:46Z

-        self.grad_scaler.unscale_(self.optimizers["global"])
-        self.transform_gradients(losses)
+        self.grad_scaler.unscale_(optimizer)
+        self._transform_gradients_with_model(model, losses)


Shouldn't this be calling transform_gradients?

emersodb · 2025-06-04T19:47:27Z


        # Forward pass on both the global and local models
        preds, features = self.predict(input)
+        log(DEBUG, f"PRED KEYS from train_step: {preds.keys()}")


Perhaps we can drop this?

emersodb · 2025-06-04T20:10:40Z


        # Call user defined methods to get predictions and compute loss
-        preds, features = self.predict(input)
+        preds, features = self._predict_with_model(model, input)


It looks like we're essentially circumventing the prediction function here, meaning that this predict function isn't called anywhere anymore, which it a bit odd.

I'm also a bit concerned that the flow of these private functions has become a bit hard to follow because there are public and private functions weaving in and out. We're also not quite being consistent with when and how these are being called, especially when inheriting (see transform_gradients comment in nnunet). If I want to change how gradients are transformed, should I override the private or public functionality?

I don't really remember what required us to move in this direction but I'm worried we'll under up with these kinds of private functions for all of our core functions and it will become a bit spaghetti.

Yeah this is for added flexibility, which I think is actually a nice (standard) way to handle what we needed to do.

The private methods are internal and more flexible now with the injection of model dependency. This is flexible in the sense should we need to change any of these private methods, then we can and can do so in a way that doesn't break the public API, namely predict(). Its not so much circumventing from the user perspective, but giving us library developers more flexibility.

We needed to go this route to be able to essentially re-use the already established logic contained in our existing clients when apply ditto and future mixins. With this internal method, we (as library devs) can make use of our private methods and supply local/global models. Without this, we can't do this and we have to resort to the alternative which is DittoClient and then subclasses of DittoClient and the other client, which I think leads to a lot of repeat code and class explosion.

What should eventually happen is that we move all existing clients to implement their private methods. There is no rush to do this as this isn't a breaking change, but we should move it at some point so its clear to devs in the future on how to create clients—namely, if the private method exists, implement those helpers. Actually, methods that don't have these private implementations cannot be adapted with Ditto or future mixins that would depend on them-an AtributeError would be raised.

Happy to discuss more on it, but I think this is okay and from my experience a pretty typical thing to do.

I think from my perspective the main issue here is that we're not doing it cleanly. For example, the predict function is here is being orphaned. It isn't called anywhere anymore. It's also unclear whether these private functions should directly call the other private functions like _transform_gradients_with_model or they should call transform_gradients. We're doing it one way in BasicClient and another in the NnUNetClient.

On point three. I think leaving it in a halfway state isn't very good. It makes knowing how to do things in the library quite odd.

Let's talk about this a together tomorrow, but I'm starting to feel inclined that supporting mixins is becoming more trouble than it's worth. Adapting the current NNUnet client to a Ditto version, while duplicating some code, would have been a much more straightforward process.

nerdai added 5 commits May 21, 2025 13:12

use 'local' as key for optimizer

55efba1

add FOR_GLOBAL_MODEL_KEY in config.utils

73e302d

add FOR_GLOBAL_MODEL_KEY in config.utils

0641991

wip

6ef3b2f

working

34f35c7

nerdai force-pushed the nerdai/ditto-nnunet branch from 2c756e9 to 34f35c7 Compare May 21, 2025 17:28

nerdai changed the title ~~Nerdai/ditto nnunet~~ Enable personalized (ditto) methods on NNunetClient May 21, 2025

emersodb requested review from emersodb and scarere and removed request for scarere May 21, 2025 19:47

scarere reviewed May 21, 2025

View reviewed changes

nerdai added 2 commits May 22, 2025 12:17

update mixins

407858f

update basic client with helper methods

2ba4f7d

nerdai force-pushed the nerdai/ditto-nnunet branch from 0f60535 to 2ba4f7d Compare May 26, 2025 14:47

nerdai commented May 26, 2025

View reviewed changes

Comment thread fl4health/clients/basic_client.py Outdated

update adaptive to override train_step

d5e139d

nerdai force-pushed the nerdai/ditto-nnunet branch from a9147f4 to d5e139d Compare May 26, 2025 15:59

emersodb reviewed May 27, 2025

View reviewed changes

nerdai added 10 commits May 27, 2025 11:32

working

032ac12

simple ditto example working

c7acb77

add val_step in dittoclient

6252414

update and simplify adaptive_drift_constrained mixin

938e976

rm old tests that are outdated

a48eded

shawn cr for get_global_model

007c179

shawns cr on notes about update_metric_manager

d3f84c8

perfcl add val step since we override predict

6d88019

perfcl client _predict

5d3a05e

cleanup

09a08f9

nerdai added 7 commits May 30, 2025 01:38

_predict rename to _predict_with_model

7bdb974

rename _train_step

be543c0

more private method renames

f95a9e8

cr

601d95b

nit docstring

6ba28e0

transform_gradients helper

b6a2f36

update README of nnunet example

33a35f6

nerdai marked this pull request as ready for review May 30, 2025 06:47

nerdai and others added 6 commits May 30, 2025 02:51

add _transform_gradients to protocol

026cb03

rename _transform_gradients_with_model

ead6473

cr

b77e7a2

test predict ditto mixin

34d0bcc

test val step

ab7f7a0

Merge branch 'main' into nerdai/ditto-nnunet

82a0882

emersodb self-requested a review June 4, 2025 15:38

emersodb reviewed Jun 4, 2025

View reviewed changes

emersodb added 2 commits June 5, 2025 09:14

Typo fix

4e4ab28

Formatting

5120805

nerdai mentioned this pull request Jul 15, 2025

[Feature] Add flexible nnunet client and ditto-ify'ed example #411

Merged

nerdai closed this in #411 Jul 15, 2025

Uh oh!

Conversation

nerdai commented May 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Type

Short Description

UPDATE (05-30-2025)

Added helper methods in BasicClient

Mixin Development

Tests Added

Uh oh!

codecov Bot commented May 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

scarere left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

emersodb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

emersodb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nerdai Jun 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

nerdai commented May 21, 2025 •

edited

Loading

Added helper methods in `BasicClient`

codecov Bot commented May 21, 2025 •

edited

Loading

nerdai Jun 4, 2025 •

edited

Loading