Paper: ColA: Collaborative Adaptation with Gradient Learning (ICLR-2024)

Reviewer comment (Theory):
**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.

Author response:
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.

Reviewer comment (Evaluation):
**Actual memory footprint not clear**: while the model claims to save storage on the main device, e.g., the GPU, the paper does not report the actual memory footprint during the forward and backward passes on the GPU. Compared to the small number of trainable parameters in LORA/adapters, the hidden feature maps and their gradients utilize the most memory. Offloading the gradient computation of the trainable parameters seems to only marginally improve memory usage on the main device. It would be great to see how much memory is actually saved on the main device by using the proposed decoupled update.

Author response:
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.

Reviewer comment (Evaluation):
**Time efficiency not reported**: the method involves an update on a second low-cost device, e.g., a CPU. The paper does not discuss the impact on training time efficiency.

Author response:
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.

Reviewer comment (Experiments):
**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.

Author response:
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.

Reviewer comment (Presentation):
The number of trainable parameters of Co1A in Table 2 is much smaller than that of LORA. Is this a typo?

Author response:
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.

Reviewer comment (Reproducibility):
Could the authors comment on the parameter-free property of the method?

Author response:
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.

Reviewer comment (Evaluation):
Could the authors report on the actual memory usage and savings on the main device?

Author response:
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.

Reviewer comment (Evaluation):
Could the authors provide a discussion, preferably quantitatively, of the impact on the time efficiency?

Author response:
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.

Reviewer comment (Theory):
The motivation and advantages of moving gradient update to CPU are unclear.

Author response:
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.

Reviewer comment (Writing):
The relationship between the proposed gradient learning and collaborative adaption is unclear.

Author response:
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.

Reviewer comment (Experiments):
The experiments do not include collaborative adaption.

Author response:
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.

Reviewer comment (Novelty):
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?

Author response:
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.

Reviewer comment (Theory):
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?

Author response:
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.

Reviewer comment (Writing):
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?

Author response:
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.

Reviewer comment (Reproducibility):
The communication cost between GPU and CPU is unclear.

Author response:
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.

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

Author response:
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.

Reviewer comment (Reproducibility):
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 either need to eventually send the local model to the server and place it on device, or you have numerous iteration rounds between the server and client to compute a single forward pass, and send the targets for an offline update of the auxiliary parameters.

Author response:
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.

Reviewer comment (Evaluation):
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)

Author response:
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.

Reviewer comment (Evaluation):
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.

Author response:
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.

Reviewer comment (Writing):
Minor error; page 4. I believe you mean the gradient of the auxilary parameters $\nabla w_{1:M}$

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

Reviewer comment (Theory):
Minor error; page 5. I believe you mean $\nabla \delta h^t_m = g^t_m(x^t_m)$

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

Reviewer comment (Experiments):
How much memory do you actually save by offloading gradients with respect to the small set of auxiliary parameters, and what is the increase in training time due to the offloading and recompilation phase? Please include logs comparing ColA to baselines in terms of memory, forward time, backward time, and number of host-device transfers.

Author response:
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.