ColA: Collaborative Adaptation with Gradient Learning

Thank you for your constructive feedback. In response to your comments, we have refined our method, revised the paper, and addressed your concerns as outlined below:

1. The motivation and advantages of moving gradient update to CPU are unclear. As far as I understand, the key contribution of the proposed method is to offload the update for the auxiliary variables w to a CPU, rather than a GPU. After reading the paper, it is still unclear to me why this is desirable. Can the authors explain the motivation and advantages of doing so?

The primary challenge in leveraging GPU for deep learning acceleration lies in its limited yet valuable memory capacity. Larger memory enables the training of larger models with a larger batch size. For instance, large model Llama-2 7B, even in 16-bit float precision, demands 12.6 GB of memory for mere GPU loading during inference. As illustrated in Tables 13, training such models requires substantially more memory, exceeding the capacity of a 48 GB GPU, even with a batch size of one. Large batch sizes, while beneficial for gradient estimation and model convergence, further exacerbate memory constraints. This challenge is evident even in advanced methods like LoRA, which struggle to fit in a 48 GB GPU with a batch size of 32. A notable approach to mitigate this issue is the Zero-Offload method utilized in DeepSpeed, which involves offloading computed GPU gradients to the CPU, so that the state information of adaptive optimizers like Adam can be saved on CPU, thereby conserving GPU memory (as described in https://www.microsoft.com/en-us/research/blog/deepspeed-extreme-scale-model-training-for-everyone/). To quote from the section "Learn how ZeRO-Offload enables multi-billion parameter training on a single GPU": `Gradients, on the other hand, are computed and averaged using reduce-scatter on the GPUs during the backward pass, and each data-parallel process then offloads the averaged gradients belonging to its partition to the CPU memory.' Our method aims to address the memory challenge by offloading the entire computation of gradient of parameters to the CPU, and also naturally allowing optimizer state information to reside on CPU. Table 1 demonstrates how our method, ColA (merge), significantly reduces memory usage by transferring gradient computations to an additional device. Furthermore, it can also increase the effective batch size by adjusting the adaptation interval. Further discussion on the computation can be found in the main response and specific comparisons in terms of memory usage and run time relative to baselines can be found in the revised paper.

2. In Proposition 2, the model is assumed to be linear. How is this result related to the case considered in the paper where the model is a pretrained neural network?

We clarify that Proposition 2 is applicable to the intermediate layers of the network. The term 'fine-tuned layers' in this context specifically refers to certain intermediate layers, not the entire neural network.

3. The relationship between the proposed gradient learning and collaborative adaption is unclear. It is unclear how the proposed gradient learning method related to collaborative adaptation. Why is gradient learning important for collaborative adaptation? What is the problem setup of collaborative adaptation in this paper?

To address your concerns, we have now incorporated user collaboration into Algorithm 1. To facilitate Fine-Tuning-as-a-Service in a computationally scalable manner, it is impractical to train $K$ distinct large models for $K$ users. Our ColA (unmerged) method allows hosting a single large model to provide fine-tuning services for $K$ users, each training separate adapters. Moreover, ColA (merged) can leverage local data and computation resources for collaborative fine-tuning of the large model.

4. The experiments do not include collaborative adaption. The paper focuses on collaborative adaptation, but the experiments do not seem to include collaborative adaptation.

Initially, our experiments on user collaboration were detailed in the Appendix. To address your concerns, we have now moved the related section in Appendix to the main body of the paper.

5. The communication cost between GPU and CPU is unclear.

We demonstrate the run time of the proposed method from Table 9 to 17 of the Appendix. The results demonstrate that the run time scales with batch size when we offload the computation to a CPU. However, offloading computation to an additional GPU significantly reduces run time. For a comprehensive discussion on computation costs, please refer to the main response and the revised paper.

Thank you for your constructive feedback. In response to your comments, we have refined our method, revised the paper, and addressed your concerns as outlined below:

1. Why is the method parameter-free?: even though the proposed model offloads the update of adapter weights to a different device, it does not make it parameter-free. It is not very convincing to claim the fine-tuning method to be parameter-free. Could the authors comment on the parameter-free property of the method ?

The term 'parameter-free' in our context underscores the absence of auxiliary parameters during the back-propagation of the base model. Initially, utilizing the 'detach' technique would exclude any auxiliary models from the back-propagation process. However, to precisely align with the computations of LoRA and classical back-propagation, we realized that it is necessary to utilize the 'parameter merging' mechanism during training as demonstrated in Algorithm 1. For a detailed explanation, please refer to the main response.

2. Could the authors report on the actual memory usage and savings on the main device? Could the authors provide a discussion, preferably quantitatively, of the impact on the time efficiency? Actual memory footprint not clear ... The paper does not discuss the impact on training time efficiency.

We conducted a comprehensive quantitative evaluation of computational costs for all feasible experiments. The results, outlined from Table 9 to 17 of the Appendix, demonstrate a significant reduction in computation space bottleneck during the back-propagation of the base model through our method. Moreover, our approach allows for a reduction in run time when offloading to an additional GPU. Further discussion on the computation can be found in the main response and specific comparisons in terms of memory usage and run time relative to baselines can be found in the revised paper.

3. Some method components are not very relevant: the discussion on parameter merging is not very relevant to the main proposed method. It is also not clear which experiments underscore the benefits of user collaboration as claimed in the contribution. Overall, the components of parameter merging and user collaboration seem tangential to the proposed method and are not well analyzed through experiments.

Parameter merging used in our refined method plays a crucial role in reducing the computation space bottleneck during training and also impacts memory consumption during inference, as evidenced in Algorithm 1 and Table 1. User collaboration is also essential to our method's applicability in Fine-Tuning-as-a-Service. As demonstrated in Table 15 to 17, ColA (merged) consumes the same amount of memory for a given batch size, regardless of the size of auxiliary models and the number of users, because the computation of all auxiliary models has been offloaded to a separate device.
To address your concerns regarding user collaboration, we have now incorporated it into Algorithm 1. To facilitate Fine-Tuning-as-a-Service in a computationally scalable manner, it is impractical to train $K$ distinct large models for $K$ users. Our ColA (unmerged) method allows hosting a single large model to provide fine-tuning services for $K$ users, each training separate adapters. Moreover, ColA (merged) can leverage local data and computation resources for collaborative fine-tuning of the large model.

4. The number of trainable parameters of Co1A in Table 2 is much smaller than that of LORA. Is this a typo?

In Table 2, the process of fine-tuning a classification model involves creating a randomly initialized output head. The peft package from Huggingface originally trains these newly introduced parameters directly without prior fine-tuning, and also double count them. In our initial implementation of ColA (Low Rank), the number of parameters used is lower due to the application of low-rank approximation on this output head. We have addressed the issue of double counting in the peft package and used a linear layer for this specific output head for all ColA methods. This adjustment ensures that the number of trainable parameters of ColA (Low Rank) and LoRA are now precisely aligned. The updated results reflecting these changes are presented in the revised paper.

Thank you for your constructive feedback. In response to your comments, we have refined our method, revised the paper, and addressed your concerns as outlined below:

1. Not clear how to align proposed method with other optimization strategies (i.e., beyond gradient descent)

Our work generalizes functional gradient descent in the context of deep learning. The auxiliary model is model-agnostic, serving as a functional to enhance the performance of the original network. This enhancement is achieved by leveraging the gradient or the so-called 'pseudo-residual' of the hidden representations.

2. Still need to forward/backward propagate the model K times for K users; i.e., the decoupled gradient computation and adaptation does not address this issue.

Contrary to the raised concern, our method requires only a single forward and backward pass for $K$ users. This efficiency is due to our approach of back propagating the gradient of hidden representations. By concatenating data from K users into a single batch, each adapter can simultaneously route corresponding samples to different adapters in parallel. We implement a `Router' class, which effectively organizes this process within a virtually parallel loop, but still we only require one forward and backward pass for a single batch of data aggregated from $K$ users.

3. To compute the change in hidden state at some layer $m < M$, you need to have first computed the hidden state at layer $m-1$. Since this computation is carried out on the GPU (server); it appears as though you ...

We do not need numerous iteration rounds for a single forward pass. In both ColA (unmerged) and ColA (merged) methods, it is only necessary for users to upload their fine-tuned auxiliary models to the server for each forward and backward pass. After the server back propagates on the base model, the hidden representations and their gradients are transferred back to the users for local updates of the auxiliary models.

4. While it is claimed that the offline update is equivalent to the online update of auxiliary parameters, this does not seem to be the case in practice (based on training curves, gradient magnitudes, and performance compared to non-offloaded computations on considered tasks)

We have resolved the previously identified discrepancy between offline and online updates of auxiliary parameters. The root cause was the `detached' mechanism previously utilized, which ignored the gradient information of auxiliary models in intermediate layers. Our revised methods now accurately align with the gradient of LoRA and classical back-propagation. Further details are available in the main response.

5. It is not clear to me how memory is actually saved by offloading gradients with respect to the small set of auxiliary parameters to the CPU, and to what degree training is slowed down due to this offloading. How much memory ...

We conducted a comprehensive quantitative evaluation of computational costs for all feasible experiments. The results, outlined from Table 9 to 17 of the Appendix, demonstrate a significant reduction in computation space bottleneck during the back-propagation of the base model through our method. Moreover, our approach allows for a reduction in run time when offloading to an additional GPU. Further discussion on the computation can be found in the main response and specific comparisons in terms of memory usage and run time relative to baselines can be found in the revised paper.

6. Minor error; page 4. I believe you mean the gradient of the auxilary parameters. Minor error; page 5 ...

The errors on pages 4 and 5 have been corrected in the revised version.

Our work distinguishes itself from existing works that primarily focus on designing a parameter-efficient auxiliary models. We instead target the computational challenges inherent in classical learning algorithms. We question the prevailing assumption that the gradient of parameters and hidden representations must be computed concurrently. We respectfully hope the reviewer to acknowledge our contribution not only in the domain of Parameter Efficient Fine-Tuning but also to methodologically recognize the significance of our approach within the broader scope of learning and optimization.