Skip to content

Fix BatchNorm buffer aggregation in FedAvg - #269

Open
Lynx-Jin wants to merge 1 commit into
TsingZ0:masterfrom
Lynx-Jin:fix/aggregate-bn-buffers
Open

Fix BatchNorm buffer aggregation in FedAvg#269
Lynx-Jin wants to merge 1 commit into
TsingZ0:masterfrom
Lynx-Jin:fix/aggregate-bn-buffers

Conversation

@Lynx-Jin

@Lynx-Jin Lynx-Jin commented Jul 28, 2026

Copy link
Copy Markdown

FedAvg currently aggregates only model.parameters(). BatchNorm running
statistics are registered buffers, so the global model retains the buffers of
the ”first uploaded client“ after aggregation.

This PR:

  • sample-weight averages floating-point buffers, including BatchNorm
    running_mean and running_var;
  • preserves num_batches_tracked as the largest observed integer counter;

@Lynx-Jin
Lynx-Jin force-pushed the fix/aggregate-bn-buffers branch from f9fa199 to 153ed84 Compare July 28, 2026 18:13
@Lynx-Jin
Lynx-Jin force-pushed the fix/aggregate-bn-buffers branch from 153ed84 to 2599d92 Compare July 28, 2026 18:34
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.

1 participant