Ferret: Federated Full-Parameter Tuning at Scale for Large Language Models
Table of Contents
Ferret uses first-order method with shared randomness to enable scalable full-parameter update. It also projects updates to a lower dimensional space to reduce communication overhead, and then reconstructs local updates from lower dimensional space with shared randomness for effective full-parameter global aggregation.
The contribution can be summarized as:
- Ferret uses first-order methods to perform local updates for each client, for faster convergence;
- Ferret projects updates into a low-dimensional space to reduce communication cost;
- Ferret reconstructs local updates from the low-dimensional space with shared randomness for effective full-parameter global aggregation.
Notice: this paper is based on federated learning setups, not split federated learning.
1. Challenge
The problem is full-parameter finetuning. Naive full-parameter finetuning guarantees the convergence speed but suffers from communication overheads. Existing first order methods incur \(O(d)\) communication overheads, due to the need to transmit the whole model parameter, where \(d\) denotes the model size. Zero-th order methods can transmit less data in each communication round, but often incur more computational cost to compensate the accuracy loss due to missing gradient.
2. The Ferret Algorithm
Ferret operates by repeatedly running 3 sequential steps over many communication rounds \(r\in[R]\) where \(R\) denotes the total number of rounds.
At the beginning of the first round, each client initializes its local parameters with pretrained model parameters, i.e., \( \mathbf{w}_{1} \gets \mathbf{w}_{0} \).
2.1. Local Updates
Each client performs \(T\) iterations of first-order optimization on its local loss function with optimizers like SGD, Adam, etc. Suppose \( \mathbf{w}_{r,0} := \mathbf{w}_{r} \) denoting the first local iteration starts from sync-ed model parameters at the beginning of \(r\)-th round. Then for SGD, each iteration can be expressed as
\[ \mathbf{w}^{(j)}_{r,t} \gets \mathbf{w}^{(j)}_{r,t-1} - \eta \nabla \ell\Big( \mathbf{w}_{r,t-1}^{(j)}; \mathbf{x}_{t-1}^{(j)} \Big) \]
where \( \mathbf{x} \) denotes local dataset. Then, we’re able to compute local update as
\[ \Delta_{r}^{(j)}:=\mathbf{w}_{r,T}-\mathbf{w}_{r,0} \]
2.2. Projected Updates
After completing local updates, each client \(j\) randomly chooses a seed \( s^{(j)} \) to generate \(K\) new random bases, where \(K \ll d\) for the sake of communication efficiency. Each base \( \mathbf{v}_{r}^{(j)} \in \mathbb{R}^{d} \).
Then, we project \( \Delta_{r}^{(j)} \in \mathbb{R}^{d} \) into a \(K\)-dimensional vector \( \mathbf{\gamma}^{(j)}_{r} \in \mathbb{R}^{K} \). Yes, here, we directly flatten the model parameters into a \(d\)-dimensional vector, and projects it into a \(K\)-dimensional vector. The point here is that, we do not care about how gradients are actually organized, but only cares about the gradients’ value. Moreover, flatten operation is indeed bi-jective if we predetermine by-row or by-column.
The design of seed is to utilize its pseudo-randomness property to avoid transmitting the random bases.
2.2.1. Reconstruction with Inversion
Let \( \Delta\in\mathbb{R}^{d} \) be the local update, and matrix \( \mathbf{V}=[\mathbf{v}_{1}\, \mathbf{v}_{2}\, \dots\, \mathbf{v}_{K}] \in \mathbb{R}^{d\times K} \) denotes the \(K\) random bases. Then we can form the convex optimization problem to determine the \(K\)-dimensional vector \(\gamma\).
\[ \gamma := \mathop{\text{argmin}}\limits_{\mathbf{y}} \Vert \mathbf{Vy}-\Delta \Vert \]
which has a closed form, given the \( \mathbf{V} \) is singular with \( K \ll d \), we can compute the projected update, as well as the corresponding reconstruction \( \tilde{\Delta} \)
\begin{equation} \begin{split} \gamma &= (\mathbf{V}^\top \mathbf{V})^{-1} \mathbf{V}^{\top} \Delta \\ \tilde{\Delta} &= \mathbf{V}(\mathbf{V}^{\top}\mathbf{V})^{-1}\mathbf{V}^{\top}\Delta \end{split} \end{equation}2.2.2. Reconstruction without Inversion
The above closed form requires computing inversion of a matrix. The total complexity will raise to \( O(K^{2}d + K^{3}) \) for time complexity and \( O(Kd) \) for memory complexity for solely \( M=\mathbf{V}^{\top}\mathbf{V} \).
Alternatively, we can approximate \(M\) with \(\mathbf{I}_{K}\):
\begin{equation} \gamma \approx (\rho K)^{-1} \mathbf{V}^{\top}\Delta \end{equation}where \( \rho:= 1-\frac{2\phi(1/\sqrt{d})/\sqrt{d}}{2\Phi(1/\sqrt{d})-1} \), \( \phi(1/\sqrt{d}) \) stands for P.D.F. and \( \Phi(1/\sqrt{d}) \) stands for C.D.F.
2.2.3. Block-wise Reconstruction
2.3. Global Aggregation
For global aggregation, we might have ring allgather or with a centralized fed server. But after all, one server (either fed server or each client) will have full information about the random seed \( s^{(j)} \) and corresponding projected local updates \( \mathbf{\gamma}^{(j)}_{r} \). We first reconstruct full-parameters updates \( \tilde{\Delta}_{r}^{(j)} \) from \( \gamma^{(j)}_{r} \)
\[ \tilde{\Delta}^{(i)}_{r} := \sum_{k=1}^{K} \gamma_{k}^{(i)} \mathbf{v}_{k}^{(i)} \]
Then, we compute the average of reconstructed updates, and apply it to previous global sync-ed parameters. Not locally computed parameters, since this is just used to compute local updates.
\[ \mathbf{w}_{r+1} \gets \mathbf{w}_{r} - \frac{1}{N} \sum_{i=1}^{N} \tilde{\Delta}^{(i)}_{r} \]
3. Theoretical Analysis
3.1. Reconstruction Analysis
From equation (2), the reconstruction satisfies that
\[ \mathbb{E}[\tilde{\Delta}]=\Delta \]