PETRA: Parallel End-to-end Training with Reversible Architectures

We thank the reviewer for acknowledging the quality and novelty of our algorithm. We respectfully address the weaknesses reported, which imply additional engineering work beyond the scope of this paper:

1. Wider Variety of Tasks and Architectures: While we agree that it would be beneficial to test our model on a wider variety of tasks and architectures, our results are comparable with the best references we are aware of [1].

2. Training Time Mention: We will remove the mention of "training time" as it is already discussed in Table 1. Since our algorithms were executed in a simulated environment, wall-clock training time cannot be used to estimate speedups accurately.

3. Drop in Accuracy: We agree that the drop in accuracy (not observable on ImageNet32) is undesirable. This is likely due to the need to adapt more hyperparameters of our method. The ResNet-50 takes much longer to run with our current experiments, preventing us from conducting larger experiments for now.

4. Convergence Analysis: Some analysis of convergence can be found in [1] for non-reversible architectures. The adaptation to a reversible architecture is straightforward; however, given the non-trivial nature of such proofs, we prefer to avoid mentioning them.

We now answer each of the reviewer’s questions:

1. Simulator Evaluation: We used a simulator to evaluate the stability of our optimization procedure. The expected training time with an effective parallel implementation is presented in Table 1, and [1] demonstrates speedup in a less stale but non-reversible setting. However, more effort is needed in our implementation, as mentioned in our general comment.

2. Larger Scale: The question is relevant since training on the ImageNet dataset was less prone to instability than CIFAR-10, indicating that our method may behave favorably on more challenging benchmarks. Our academic resources do not allow us to benchmark significantly larger models within the rebuttal time frame.

3. Convergence Analysis for Reversible Architectures: See our answer to weakness 4.

$1$ On the Acceleration of Deep Learning Model Parallelism With Staleness

First, we sincerely thank the reviewer for the very positive feedback. It's truly appreciated.

We'd like to address the weaknesses you reported:

1. We agree that developing a scalable and distributed implementation of the algorithm would be the ultimate proof of concept. However, our primary goal is to introduce a novel training algorithm that could lead to a significant breakthrough. We have achieved this goal by proposing a fully functional simulator of our method and other model parallel approaches (DSP [1], PipeDream [2]), which allowed us to empirically confirm the success of our method on a large-scale dataset for a high number of stages (18). While many questions remain open and a functional distributed implementation is essential for engineering purposes, we believe that NeurIPS is the right venue to explore new, promising ideas.

2. We will remove the mention of “training time” in the title of the paragraph at line 256, as this was a mistake since it is already discussed in Table 1. The memory savings in Section 4.2 are reported within the context of our simulator, which measures the impact of maintaining parameter or activation buffers. Table 1 effectively explains the benefits of our method compared to state-of-the-art approaches, including pipelining, which necessarily involves some additional buffers to handle microbatches. Our main objective was to remove the dependency on $J$, which is a problematic scaling law in both time and space categories of the Table. Practical measurements are highly code-dependent, and we believe that such low-level optimization and benchmarking are out of the scope of this paper.

Regarding the question about the impact of interconnect traffic, we acknowledge that it can be a serious issue in distributed training. However, since our increase factor for communication is fixed (i.e., $2$ and $4$ for forward and backward communication, respectively), we do not believe it is a fundamental concern for the scaling potential of PETRA. In model parallelism, the communications between workers typically contain activation tensors, which have sizes proportional to the width of the model. If PETRA is dominated by communications during training on a given cluster for a reversible architecture, it means that any other model-parallel method would struggle to scale the width of the non-reversible counterpart beyond a factor of $2$ or $4$. A similar argument can be made regarding the mini-batch size, which also scales with the activations. Thus, while our increased communications are important to understand for making the best use of a given compute configuration, their increase is only constant and does not overshadow PETRA's main advantage: a linear speedup for a constant memory cost, allowing significant depth scaling due to substantial memory reduction.

$1$ On the Acceleration of Deep Learning Model Parallelism With Staleness

$2$ PipeDream: Fast and Efficient Pipeline Parallel DNN Training

We thank the reviewer for acknowledging the quality of the ideas presented in this paper. We would like to address the concerns:

1. Our paper serves as a proof of concept for a new algorithm, and our resources (engineers, clusters) currently only allow us to use a simulated environment. Several works have demonstrated that the use of stale gradients can be implemented efficiently (see [1,2]). However, our focus is to demonstrate that our method attains high accuracy with a larger number of stages while offering a drastic reduction in memory footprint. Although theoretical, Table 1 details the issues (buffers, communications, compute time per batch) concisely and precisely.

2. We agree that comparing the minimum number of GPUs required for a given task with different optimization methods, or how the maximum size of a model that can fit on a given cluster changes for different optimization methods, is crucial to emphasize the practical advantages of our method against available alternatives. Despite our implementation effectively simulating PETRA and alternative optimization methods within the same code base, a distributed implementation is necessary for such comparisons, which is still under development.

3. Our method used the maximum number of stages for our models, i.e., $10$ for ResNet-18 and $18$ for ResNet-50, unlike the closest work to ours, which only goes up to $K=3$ on ImageNet (see [1]). Extending our work to other downstream tasks, such as large language model (LLM) training, is interesting but beyond the scope of this work.

4. As mentioned in response to the first and second weaknesses, our current resources prevent us from experimenting at a larger scale. However, such experiments are under active development, and we emphasize that we are not aware of other works allowing such an extensive split of $18$ stages while keeping good performances as we showed here.

5. Between the reversible backpropagation of the RevNet and PETRA, there are two differences. The first is the use of stale gradients to update the parameters, which is key to the linear speed-up of such approaches. The second difference is the removal of the parameter buffers in the stages. This results in PETRA approximating the backpropagation computation with the updated parameters, which differ from the ones in the forward pass. However, we found this approximation to have a limited impact. Thus, minor fluctuations are expected due to the use of stale and approximated gradients, but they are surprisingly low. Additionally, we had to fully reimplement autograd to make the gradient estimation process feasible, so errors like round-off are common and expected.

$1$ On the Acceleration of Deep Learning Model Parallelism With Staleness

$2$ PipeDream: Fast and Efficient Pipeline Parallel DNN Training

First, we sincerely thank the reviewer for their positive assessment and acknowledgment of the clarity and novelty of our method. We would like to address the reported weaknesses:

1. We believe the reviewer is referring to Table 1. We would like to reformulate our discussion from the end of Section 3.3: Except for PETRA, each method in the table has a dependency on $J$, the depth of the network. For example, backpropagation for reversible architectures has a mean time per batch that depends on $J$, meaning it cannot achieve computational speedup. For delayed gradient methods, the size of the buffers depends on $J$, leading to strong memory constraints. As $J$ grows, only PETRA maintains a constant use of memory resources.

2. To our knowledge, the best comparison point is [1], which does not implement any invertible mechanism. Their splits into $K$ stages are small ($K=3$), whereas we use $K=18$. Table 4 of this paper reports top-1 accuracies of 68.9% for ResNet-18 and 74.9% for ResNet-50 using a split of 3. In comparison, using splits of $K=10$ and $K=18$ respectively, we report top-1 accuracies of 71.0% and 74.8%, which is significantly better. This paper was published at CVPR, a top-tier conference, with an optimized implementation that is not available online. Therefore, we could not base our work on it, even though they report a computational advantage with only a split of $K=3$.

3. We agree that deploying our code in a real distributed environment to benchmark its efficiency would be the ideal proof of concept. However, our academic resources (engineers, clusters) only allowed us to use a simulated environment. Nevertheless, this environment enabled us to accurately simulate many model parallel approaches and demonstrate the convergence of PETRA, even for ImageNet with many stages ($18$), accomplishing the goal of this paper.

$1$ On the Acceleration of Deep Learning Model Parallelism With Staleness