Memory-Efficient LLM Training with Online Subspace Descent

Re: “Memory and efficiency”

Updating $P_t$ with AdamW will inadvertently increase the memory consumption. However, in practice, we observe that the increase in peak memory is minuscule compared to the overall size of the model. Meanwhile, we can enjoy the nice properties of online subspace descent for its flexibility and system efficiency when implemented in parallel (as we discussed with reviewer NLbB and Rbsr). In addition, our subsequent experiment shows that SGD can perform similarly to AdamW when updating P, in which case, there is no additional memory cost over GaLore, since no momentum needs to be saved.

Here is a comparison between GaLore and our method on 7B llamas trained for 10K steps on C4 (updating $P_t$ with SGD)

The hardware setups are completely identical, both using a single node of H800s with 8 cards. It's clear that our method is not only faster in wall clock time, but also reaches lower perplexity when seeing same number of tokens.

Re: "Discrete case"

The derivation of discrete case is highly similar to the continuous example we've shown in the paper.

Example 4.1. Momentum + Online Subspace Descent is

$$\frac{d}{dt} W_t = -P_t \hat{M}_t, \quad \hat{G}_t = P_t^\top \nabla L(W_t), \quad \frac{d}{dt} \hat{M}_t = a(\hat{G}_t - \hat{M}_t),$$

with Lyapunov function $H(W, \hat{M}) = L(W) + \frac{\|\hat{M}\|^2}{2a}$, for which

$$\frac{d}{dt} H(W_t, \hat{M}_t) = - \nabla L(W_t)^\top P_t \hat{M}_t + \hat{M}_t^\top (P_t^\top \nabla L(W_t) - \hat{M}_t) = -\|\hat{M}_t\|_2^2 \leq 0.$$

Followup this example as a demo, we have the following discrete derivation:

$W_{t+1} = W_t - \epsilon P_t \hat{M}_t$

$\hat{M}_{t+1} = \hat{M}_t + a(\hat{G}_t - \hat{M}_t)$ $\hat{G}_t = P_t^\top \nabla L(W_t)$

where $\hat{M}_{t+1} = \hat{M}_t + a(\hat{G}_t - \hat{M}_t) = (1-\alpha \epsilon)\hat{M}_t + \alpha \epsilon \hat{G}_t = \beta \hat{M}_t + (1-\beta)\hat{G}_t$. Assuming $L(\cdot)$ is $L$-smooth, we have

$$H(W_{t+1}, \hat{M}_{t+1}) - H(W_t, \hat{M}_t) \leq -\epsilon \|\hat{M}_t\|_2^2 + \alpha \epsilon^2 \|\hat{G}_t - \hat{M}_t\|_2^2 + \frac{L}{2}\epsilon^2\|P_t\hat{M}_t\|_2^2.$$

Hence, a telescoping sum yields

$$\frac{1}{T}\sum_{t=0}^{T-1} \|\hat{M}_t\|_2^2 \leq \frac{H(W_0, \hat{M}_0) - H(W_T, \hat{M}_T)}{\epsilon T} + \epsilon B_t,$$

where $B_t = \frac{1}{T} \sum_{t=0}^{T-1} \alpha \|\hat{G}_t - \hat{M}_t\|_2^2 + \frac{L}{2}\|\hat{M}_t\|_2^2$.

Re: “high-frequency versus low-frequency tokens”

This is our attempted explanation for the phenomenon we observe in our ablation study of different ranks. Loss of higher rank runs tend to drop faster and converge to a lower loss after 10K steps. We suspect that it’s due to the unique property of language modeling tasks, where intuitively rarer tokens (tokens that are not common) require high gradient rank to be learnt. However, we would like to delay the justification of this hypothesis to future study, since it’s not the main focus of our paper.

Re: "AdamW is more suitable for online PCA"

This is what we find empirically, when trying different optimizers such as LION and Ada Factor. However, it’s not unfathomable that SGD in principle can perform similarly, which we show above. It’s just for the default learning rates and schedules, adamW seems to work more robustly and reliably and it’s the least sensitive one.

The derivation presented above provides a foundational outline for our convergence analysis. Building on this, we can conduct a more detailed analysis following the established classical (or traditional) approach.

Taking a step further from the bound above, we have

$(1 - \frac{L \epsilon}{2}) \frac{1}{T} \sum_{t=0}^{T-1} ||\hat{M}_t||_2^2 \leq \frac{H(W_0, \hat{M}_0) - H(W_T, \hat{M}_T)}{\epsilon T}$

$+\epsilon \frac{1}{T} \sum_{t=0}^{T-1} \alpha ||\hat{G}_t - \hat{M}_t||_2^2$

One can use the established standard methods to bound the last term $||\hat{G}_t - \hat{M}_t||_2^2$ (for example at page 11-12 in the supplementary from [Sun et al. 2023]).

References

Tao Sun, Qingsong Wang, Dongsheng Li, and Bao Wang. Momentum ensures convergence of signsgd under weaker assumptions. In International Conference on Machine Learning, pages 33077–33099. PMLR, 2023.

Re: "Downstream tasks"

We did standardized GLUE evaluation for the above two 7B checkpoints with eval-harness.

Both Cola and STSB require further fine-tuning, which brings in more confounding factors (even Llama-2 7B trained by Meta doesn't get meaningful results), so we exclude them from the comparison.

Re: “Many strong assumptions that do not reflect real world datasets”

We respectfully argue that our assumptions (Assumption 4.4) are in fact, very much minimum.

1. The model needs to be continuously differentiable. (pretty much all modern neural nets)
2. The optimization stops when the projected update hits 0. (this is rather a fact than a assumption)
3. It says we should pick $P$ such that $P^TG$, the projected gradient is not always 0. Otherwise, it just doesn’t update the model weights at all and it defeats the purpose.

Every theoretical work has to make some assumptions. We want to assure our reviewers that it will be hard to find papers with same level of theoretical rigorousness that is making fewer assumptions than we do: we did not assume convex losses; we did not assume specific forms of the model (linear model or MLPs); we did not even assume specific optimizers (Adam, momentum, lion and etc)...

Re: “The memory savings and the perplexity is very similar to GaLore and it is not clear that it is significant”

We agree that this does have very similar memory saving as GaLore. However, we argue that our contributions lie majorly in our analysis that introduces a new family of optimizers that are guaranteed to converge. Among them, online subspace descent not only possesses the same memory saving properties, but also can be implemented more efficiently in a parallel manner, whereas Galore doesn’t share the same flexibility.

Re: “Intuitions”

The rough intuition is that $P_t$ serves as a kind of preconditioning matrix in the Hamiltonian systems. But to arrive at the precise mathematical conclusion, we find that the best and quickest way to understand it is through the derivation in Eq (8), together with physical understandings of Hamiltonian systems (that, e.g., the hamiltonian consists of potential energy and kinetic energy, and the kinetic energy term is different for different optimizers)

More specifically, as long as the projected subspace defined by $P_t$ keeps changing and not always be orthogonal to the gradient as we stated in Assumption (4.4), it will always make some progress towards the fixed point. Given enough time, it will eventually reach there. Hence the consequence of convergence does not depend on $P_t$.

l.1ß6 “Generally, Adam updates are memory-bound operations”

• This is precisely the bottleneck that this family of online subspace descent optimizers aims to tackle. By projecting into subspace, the memory footprint of Adam will be greatly reduced, allowing us to schedule another operation concurrently.
• In addition, since we are running tensor-wise operation, the incremental memory consumption over running adamW in subspace is simply backward of P for a single tensor. This is reflected in our experiment for Llama 1B, where the peak memory increase compared with Galore AdamW is only 0.4 GB, that is counting fragmentations, redundant intermediate saving in practice.

l.1ß6 “the same could be said for GaLore, you run you SVD in parallel to the main network”

• It can definitely run SVD in parallel synchronously. But it won’t benefit from doing that. Instead, it will slow down that specific training step SVD is paralleled with if SVD on the latest gradient is computed. SVD uses Jacobi iterations until the principal components converge, in some cases it fails and has to restart. Overall on weight matrix sizes we care about, we observe that it’s 20x - 150x slower than running a single backward step + update which makes it impossible to be hidden by the adamW operations. Waiting for SVD to be done slows down training, whereas we essentially propose to distribute and amortize that workload into every step of training instead of doing it all at once and wait at a single step. In our latest experiment of the 7B model, we also show that our method is 1.37x faster than baseline GaLore in wall clock time.
• In another scenario, it’s also possible to run SVD completely asynchronously and use $P_t$ computed from some arbitrarily old gradients to perform projections and updates. However, this becomes a very different class of algorithms from baseline GaLore. It overly complicates implementation in practice, since it needs to manage schedules for every tensor in the network. SVD finish time can also vary and it is unpredictable. Hence, it's possible to have part of the network using $P_t$ calculated from a different timestep, if it's running in a non-blocking fashion. The implication of this is unknown and it's a very interesting direction of exploration. Yet, given the limited span of this work, we would like to defer this discussion to future works.

l.119 "Matrix multiplications are linear algebra routines"

We meant to refer to torch.linalg.decomposition routines. We will make it more specific in revision.

Re: "119-128 could do with some rewriting"

• We agree that we can rewrite it to highlight the goal of doing online PCA, not just for saving computational cost.

l. 170 "There is no requirement on Γ here, besides that the derivative in (8) should exist Γ does not appear in (8), so it isn't clear which derivative should exist."

• The derivatives in (8) refers to the partial derivatives of $H_t$ that appear in Eq (8). We will make it specific in the revision.

l. 223 "Given that GaLore runs SVD only every 300 steps, this would make GaLore computationally more efficient than the proposed method."

Please refer to the above explanation on system efficiency/parallel implementations. Even if running Galore every 300 steps, it’s still an unmitigable slowdown where no parallelization tricks can help. And as tensor shape grows, as figure 2 shows, the gap of latency becomes higher. Meanwhile the cost of online subspace descent can be hidden with a smart parallelization trick. Our latest experiment of 7B model shows that our method is 1.37x faster than baseline GaLore in terms of wall clock time with the exactly the same hardware setup.

Re: "Larger model"

We pretrained from scratch 7B llamas on C4 for 10K steps. Perplexity is the lower the better. The results are consistent with our claim in production model scale.

l.145 " Adversarial projections"

The loss in $L_{G_t}(P)$ attains global optima only if P spans the linear space of the top k eigenvectors of $G_t$ in Loss (the other eigenvectors are saddle points). Hence, if an optimizer minimizes $L_{G_t}(P)$ correctly, it should find the top k eigenvectors of G, which is zero iff G ==0. Regardless of the choice $P_t$, the system yields a given Hamiltonian function and converges to some invariant set. But the invariant set coincides with the local optima of the objective function f only when Assumption 4.4 satisfies. The choice you provide may not satisfy Assumption 4.4.

More concretely, solution of $P_t=min<G_t,P_t^TPM_t>$ gives the opposite direction of $G_t$. Considering the negative direction won’t harm it, because it’s still the same subspace even the gradient has an opposite sign and while it’s projected back, it’s still “gradient descent”, the two negative signs simply cancel out. The real adversarial scenario is $P_t=min|<G_t,P_t^TPM_t>|$. It is projecting into complementary orthogonal subspace, aka, $<G_t,P_t^TPM_t> = 0$. This is against ii) in our Assumption 4.4, which states we can not pick $P$ such that $P^TG$, the projected gradient is always 0. Otherwise, it just doesn’t update the model weights at all and it defeats the purpose.

l 158 "I'm not sure what "once" is supposed to mean here. Should it just be "if"?"

Yes, it should be “if”. We will correct that in revision.

l. 235 "choices of rank"

Just in terms of speed of convergence, in our ablation study, within 10K steps, the higher rank it’s the faster it goes. Eventually, if it’s run long enough, our theory indicates that it will always converge. However, in practice, it might be difficult to observe that due to constraint in resources. So as a safe recommendation, if it can be afforded, choose a higher rank.