nvflare.app_opt.lightning.api module

class FLCallback(rank: int = 0, load_state_dict_strict: bool = True, update_fit_loop: bool = True)

Bases: Callback

FL callback for lightning API.

Parameters:

• rank – global rank of the PyTorch Lightning trainer.

• load_state_dict_strict – exposes strict argument of torch.nn.Module.load_state_dict() used to load the received model. Defaults to True. See https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.load_state_dict for details. NVFlare still validates incoming keys and shapes before calling load_state_dict(). With True, unexpected incoming keys are treated as contract drift and fail fast. With False, unexpected keys are logged and ignored, while compatible keys are still loaded.

• update_fit_loop – whether to increase trainer.fit_loop.max_epochs and trainer.fit_loop.epoch_loop.max_steps each FL round. Defaults to True which is suitable for most PyTorch Lightning applications.

on_before_optimizer_step(trainer, pl_module, optimizer, *args)

Called before optimizer.step().

on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx)

Called when the train batch ends.

Note

The value outputs["loss"] here will be the normalized value w.r.t accumulate_grad_batches of the loss returned from training_step.

on_train_end(trainer, pl_module)

Called when the train ends.

on_train_start(trainer, pl_module)

Called when the train begins.

on_validation_end(trainer, pl_module)

Called when the validation loop ends.

on_validation_start(trainer, pl_module)

Called when the validation loop begins.

reset_state(trainer)

Resets the state.

If the next round of federated training needs to reuse the same callback instance, the reset_state() needs to be called first Not only resets the states, also sets states for next round

patch(trainer: Trainer, restore_state: bool = True, load_state_dict_strict: bool = True, update_fit_loop: bool = True)

Patches the PyTorch Lightning Trainer for usage with NVFlare.

Parameters:

• trainer – the PyTorch Lightning trainer.

• restore_state – whether to restore optimizer and learning rate scheduler states. Defaults to True.

• load_state_dict_strict – exposes strict argument of torch.nn.Module.load_state_dict() used to load the received model. Defaults to True. See https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.load_state_dict for details. NVFlare still validates incoming keys and shapes before calling load_state_dict(). With True, any incoming key that does not exist in the local Lightning module is rejected before loading. With False, NVFlare warns and filters the payload down to matching keys, which is useful for partial model updates where the client only keeps part of the server keyspace.

• update_fit_loop – whether to increase trainer.fit_loop.max_epochs and trainer.fit_loop.epoch_loop.max_steps each FL round. Defaults to True which is suitable for most PyTorch Lightning applications.

SCAFFOLD:

When the received model contains SCAFFOLD global controls, patch automatically applies PTScaffoldHelper around Lightning’s optimizer steps and returns the required control difference. This supports automatic optimization with one optimizer and precision="32-true" or precision="bf16-mixed". Manual optimization and other precision modes are not supported by the patched path. Those clients must use an explicit receive/train/send loop and integrate PTScaffoldHelper directly.

FedProx:

When the received model contains a positive FedProx coefficient, patch automatically adds the exact proximal gradient to optimizer-owned trainable parameters. FedProx composes with automatic SCAFFOLD support. The proximal gradient is applied after gradient accumulation and AMP unscaling and before gradient clipping. Consequently, a loss logged from training_step excludes the injected proximal term even though optimization includes its exact gradient. Automatic FedProx has the same one-optimizer, automatic-optimization, and precision restrictions as automatic SCAFFOLD; closure-based LBFGS and sparse gradients are also unsupported.

Example

Normal usage:

trainer = Trainer(max_epochs=1)
flare.patch(trainer)

Advanced usage:

If users want to pass additional information to FLARE server side via the lightning API, they will need to set the information inside the attributes called __fl_meta__ in their LightningModule.

class LitNet(LightningModule):
    def __init__(self):
        super().__init__()
        self.save_hyperparameters()
        self.model = Net()
        self.train_acc = Accuracy(task="multiclass", num_classes=NUM_CLASSES)
        self.valid_acc = Accuracy(task="multiclass", num_classes=NUM_CLASSES)
        self.__fl_meta__ = {"CUSTOM_VAR": "VALUE_OF_THE_VAR"}