Custom Preconditioners

The KFACPreconditioner is actually a subclass of the BaseKFACPreconditioner which adds layer registration logic and the KAISA distribution strategy.

Custom distribution strategies other than KAISA can be implemented by subclassing BaseKFACPreconditioner. The BaseKFACPreconditioner is a simple class that just implements the K-FAC preconditioning logic. The user is responsible for providing which layers to precondition, worker assignments, and a collective communication interface.

I highly suggest carefully reading through the constructor of the KFACPreconditioner to learn how that implementation handles the layer registration and worker assignments that get passed to the BaseKFACPreconditioner.

The WorkAssignment interface is used by BaseKFACPreconditioner to determine which workers should perform which operations on each layer. The KAISAAssignment is an implementation of the WorkAssignment interface which implements the KAISA distribution strategy based on the gradient worker fraction.