Skip to content

Add FedCE research recipe support - #4922

Open
nvkevlu wants to merge 5 commits into
NVIDIA:mainfrom
nvkevlu:codex/fedce-research-recipe
Open

Add FedCE research recipe support#4922
nvkevlu wants to merge 5 commits into
NVIDIA:mainfrom
nvkevlu:codex/fedce-research-recipe

Conversation

@nvkevlu

@nvkevlu nvkevlu commented Jul 20, 2026

Copy link
Copy Markdown
Collaborator

Adds dedicated PyTorch recipe support for FedCE.

Description

FedCE uses the FedAvg global-model workflow but changes how client updates are evaluated and aggregated. Because it introduces its own client/server contract and aggregation behavior, it is provided as a separate FedCERecipe instead of a FedAvg configuration flag.

The PR includes:

  • A FedCE client helper and contribution-aware model aggregator.
  • A dedicated FedCERecipe using PyTorch weight differences.
  • Recipe exports and CLI discovery support.
  • Documentation describing how to use the recipe.
  • Unit tests covering aggregation, validation, recipe construction, and export behavior.

FedSM and other personalized-learning algorithms are outside the scope of this PR.

Validation

The original FedCE implementation and the new recipe were compared using the documented six-client prostate segmentation experiment.

Both implementations completed all 100 rounds successfully with no fatal validation errors. Their late-round mean global Dice scores were very similar: approximately 0.845 for the original implementation and 0.843 for the recipe. FedCE contribution coefficients also remained normalized throughout the run.

These results provide confidence that the recipe preserves the expected FedCE research behavior.

Types of changes

  • Non-breaking change (fix or new feature that would not break existing functionality).
  • Breaking change (fix or new feature that would cause existing functionality to change).
  • New tests added to cover the changes.
  • Quick tests passed locally by running ./runtest.sh.
  • In-line docstrings updated.
  • Documentation updated.

@codecov-commenter

codecov-commenter commented Jul 20, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 64.32%. Comparing base (8b9541c) to head (9217279).

Additional details and impacted files
@@            Coverage Diff             @@
##             main    #4922      +/-   ##
==========================================
+ Coverage   64.26%   64.32%   +0.06%     
==========================================
  Files        1030     1032       +2     
  Lines      103578   103761     +183     
==========================================
+ Hits        66560    66745     +185     
+ Misses      37018    37016       -2     
Flag Coverage Δ
unit-tests 64.32% <100.00%> (+0.06%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 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.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@nvkevlu
nvkevlu marked this pull request as ready for review July 28, 2026 14:11
Comment thread nvflare/app_opt/pt/fedce.py
@greptile-apps

greptile-apps Bot commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

Adds a dedicated PyTorch FedCE recipe and contribution-aware aggregation workflow.

  • Introduces client helpers for leave-one-out model evaluation and contribution metadata.
  • Adds FedCE recipe exports, CLI catalog discovery, documentation, and unit coverage.
  • Tightens client update validation by requiring identical parameter-name sets.

Confidence Score: 3/5

The PR is not yet safe to merge because same-name client parameters with incompatible shapes can still abort a FedCE aggregation round.

The new compatibility check rejects missing or additional parameter names, but it passes matching names without validating tensor shapes before independently flattening and combining client vectors.

Files Needing Attention: nvflare/app_opt/pt/fedce.py

Important Files Changed

Filename Overview
nvflare/app_opt/pt/fedce.py Adds FedCE helper and aggregation logic, but parameter validation still permits same-name tensors with incompatible shapes that can abort aggregation.
nvflare/app_opt/pt/recipes/fedce.py Adds a FedAvg-derived PyTorch recipe configured for DIFF transfer and contribution-aware aggregation.
nvflare/tool/recipe/recipe_cli.py Registers the FedCE PyTorch recipe in the documented CLI recipe catalog.
tests/unit_test/app_opt/pt/fedce_test.py Covers contribution calculations and name-set validation, but does not exercise incompatible shapes under matching parameter names.

Sequence Diagram

sequenceDiagram
    participant S as FedCE Server
    participant C as Client
    participant A as FedCE Aggregator
    S->>C: Global model + prior fedce_coef
    C->>C: Train and evaluate minus model
    C-->>A: Weight DIFF + fedce_minus_val
    A->>A: Validate parameter contract
    A->>A: Compute cosine and minus-score weights
    A-->>S: Aggregated DIFF + updated fedce_coef
Loading

Reviews (4): Last reviewed commit: "Merge remote-tracking branch 'online/mai..." | Re-trigger Greptile

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds first-class FedCE support to NVFlare’s PyTorch “recipe” system, including a contribution-aware server-side aggregator, a client helper for the FedCE contract, CLI catalog discovery, docs, and unit tests.

Changes:

  • Introduces FedCEModelAggregator and PTFedCEHelper for FedCE contribution estimation and minus-model handling.
  • Adds FedCERecipe (PyTorch) and exposes it via recipe package exports and the nvflare recipe CLI catalog.
  • Adds documentation and unit tests for the FedCE aggregator/recipe behavior and recipe discovery/export.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.

Show a summary per file
File Description
tests/unit_test/tool/recipe_cli_test.py Updates catalog expectations to include the new fedce-pt documented recipe.
tests/unit_test/app_opt/pt/recipes/fedce_recipe_test.py Adds unit tests for FedCERecipe validation, DIFF transfer, lazy exports, and server export wiring.
tests/unit_test/app_opt/pt/fedce_test.py Adds unit tests for FedCE helper utilities, normalization, and contribution-weighted aggregation behavior.
nvflare/tool/recipe/recipe_cli.py Registers fedce-pt in _DOCUMENTED_RECIPE_SPECS so the CLI can discover and describe the new recipe.
nvflare/app_opt/pt/recipes/fedce.py Implements the dedicated FedCERecipe that wraps FedAvg workflow with a FedCE-specific aggregator and DIFF exchange.
nvflare/app_opt/pt/recipes/init.py Adds lazy import + __all__ export for FedCERecipe.
nvflare/app_opt/pt/fedce.py Implements FedCE server/client components: constants, client helper, normalization, and the contribution-aware FedCEModelAggregator.
nvflare/app_opt/pt/init.py Exposes FedCE components from the PyTorch optional package top-level.
docs/user_guide/data_scientist_guide/available_recipes.rst Documents how to construct/use FedCERecipe and the client script contract requirements.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread nvflare/app_opt/pt/fedce.py
Comment thread docs/user_guide/data_scientist_guide/available_recipes.rst Outdated
@holgerroth

Copy link
Copy Markdown
Collaborator

The contribution scoring and the actual aggregation currently operate on different parameter scopes:

  • Contribution weights are computed on the intersection of all clients' returned keys, further filtered by trainable_param_names (_get_cosine_param_names).
  • aggregate_model() then applies those weights to all params each client returned via WeightedAggregationHelper — including non-trainable buffers and any client-exclusive keys.

For buffers (running_mean, num_batches_tracked, ...) this seems fine (they get contribution-weighted instead of evenly averaged), but client-exclusive keys would be aggregated among only their contributors, entirely outside the contribution-scoring rationale (related to what Greptile flagged).

Proposal:

  1. Log a warning in aggregate_model() when clients' returned key sets differ from each other (or from the scoring set), so silent divergence is visible.
  2. Document the assumption in the FedCEModelAggregator docstring: FedCE expects homogeneous parameter sets across clients; contribution weights are derived from the common trainable params but applied to everything returned, including buffers.

@holgerroth

Copy link
Copy Markdown
Collaborator

On the validation: the global Dice comparison (≈0.845 vs ≈0.843) is a good sanity check, but the signature result of FedCE — and what the research example documents — is the contribution-coefficient behavior, see the curves in research/fed-ce: client_Promise12's coefficient rises over training due to its domain shift while the other clients' coefficients gradually decrease.

Could you attach the equivalent fedce_coef curves from the recipe run (and ideally per-client validation Dice) so we can confirm the recipe reproduces this behavior?

It would also be good to include the client training script / job config used for the prostate comparison — or better, add a recipe-based job under research/fed-ce — so the experiment is reproducible.

@nvkevlu

nvkevlu commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator Author

I enabled per-round contribution tracking in the recipe validation using NVFlare’s TensorBoard experiment-tracking path. The client emits FedCE_Coef each round, and the recipe installs the TensorBoard receiver.

I ran both the original implementation and FedCERecipe for 100 rounds using the documented six-client prostate configuration. Both runs completed successfully and produced complete per-client Dice and contribution series.

  • original implementation: job 30891157
  • FedCERecipe: job 30918334
image

Promise12 shows the expected behavior under the recipe. Its contribution coefficient rises from 0.169510 at round 1 to 0.273279 at round 99—an absolute increase of 0.103769 (61.2%)—and peaks at 0.275921 at round 73. It finishes with the largest coefficient. The original implementation similarly rises from 0.168981 to 0.266813.

Across all clients and rounds 1–99, the original and recipe coefficient series have a correlation of 0.9976 and a mean absolute difference of 0.00281. The recipe coefficients remain normalized: their per-round sums differ from 1 by at most 2.98e-8.

The plot starts contribution curves at round 1 because round 0 has no prior-round contribution coefficient. These results reproduce the research example's qualitative behavior: Promise12 receives increasing weight while the other clients' weights generally decrease or remain near the uniform reference.

@holgerroth

Copy link
Copy Markdown
Collaborator

Thanks @nvkevlu — this is convincing. The Promise12 trajectory (0.1695 → 0.2733, finishing largest while the others stay near or below the uniform reference) reproduces the signature behavior from the research example, and your rerun of the original implementation ending at 0.2668 matches the published figure, so the baseline itself reproduced. The 0.9976 correlation with ~0.0028 mean absolute difference is consistent with the minor numerical differences between the implementations (clip/normalize ordering, epsilon) plus run-to-run noise, and the per-round sums within 3e-8 of 1 confirm the recipe's by-construction normalization.

Two remaining items before approval:

  1. Reproducibility: the job IDs are internal references. Could you include the client training script and job config used for this comparison — ideally as a recipe-based job under research/fed-ce — so the experiment is externally reproducible? The client script that emits FedCE_Coef and implements the minus-model contract is exactly the piece the PR doesn't show.
  2. The earlier scope comment: the key-set mismatch warning and the docstring documenting the assumption (homogeneous parameter sets across clients; contribution weights applied to all returned params including buffers) haven't landed on the branch yet. Since the buffer weighting matches the original FedCEAssembler behavior, this is a documentation/logging fix, not an algorithm change.

reference_client = clients[0]
common = set(self._results[reference_client].params)
for client in clients[1:]:
client_params = set(self._results[client].params)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Tensor shapes bypass compatibility check

When two clients return DIFF updates with identical parameter names but different tensor shapes, this name-set comparison accepts them and _weighted_vector combines flattened vectors of different lengths, causing a PyTorch size-mismatch error that aborts the aggregation round.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants