torch.optim.Adafactor

[msaroufim]

🚀 The feature, motivation and pitch

Models are getting harder to fit on a limited number of GPUs and ADAM doesn't help since its memory overhead is 2N where N is the number of parameters in a model

We don't like to merge optimizers in core because they rarely stand the test of time but ADAM has and a memory efficient alternative that's been in use at many startups I've talked to, Twitter and larger companies https://github.com/google-research/t5x/blob/main/t5x/adafactor.py, has been Adafactor, see the discussion here https://twitter.com/HamelHusain/status/1702004478369743125 - it's also come up in a github issue before here #98434

Assigning to myself since I want a starter task in optimizers

Alternatives

There is a working implementation in fairseq which Huggingface has borrowed https://github.com/facebookresearch/fairseq/blob/main/fairseq/optim/adafactor.py#L66 which is a good starting point

Additional context

You might be thinking why merge yet another optimizer? Optimizers are plagued by lack of reproducibility and sensitivity to hyperparameters

However, ADAM has stood the test of time but ADAM also has a high memory overhead, for each parameter you need to store the first and second moment so if your model has N parameters then your optimizer state is 2N. Also few people change its hyperparameters in fact the hyperparams have remained the same since torch days

• https://github.com/pytorch/pytorch/blob/main/torch/optim/adam.py
• https://github.com/torch/optim/blob/master/adam.lua

So as Blueberries projects move towards finetuning after PyTorch conference and as more members of the community try to fit larger models on their GPUs it's critical we find memory efficient alternative to ADAM which is why Adafactor might just be a strong candidate for a default optimizer we recommend to people finetuning with a limited amount of GPUs and this includes ourselves in the torchbench CI, Adafactor comes with its own LR scheduler and will help us not to have to dramatically reduce batch sizes in our benchmarks

cc @vincentqb @jbschlosser @albanD @janeyx99 @crcrpar

[janeyx99]

Thanks for opening an issue--this is intriguing! I agree that the optimizer footprint for Adam is quite sasquatch and now the question is whether Adafactor would be the strongest candidate here.

On one end of the spectrum, as you've mentioned, we have vanilla SGD without momentum which takes up no state but takes longer to converge.

On the other hand, we have Adam, which takes up 2X params memory (and more runtime) but converges faster.

Skimming https://openreview.net/pdf?id=Sf1NlV2r6PO seems to suggest that there is coordinate clipping that could be done in order to make SGD more effective, though I haven't yet dug into the details. It is interesting this paper also compared against Adafactor, which is a positive sign for Adafactor.

Some stats on GitHub code search when searching for regex /[^A-Za-z]Optim\(/:

Optim	num code results
Adam	639k
SGD	377k
RMSprop	82.9k
AdamW	53.5k
Adadelta	22.1k
Adagrad	21.2k
LBFGS	14.1k
NAdam	10k
Adamax	7.4k
LARS	5.4k
RAdam	5.2k
LAMB	4.5k
Adafactor	3.2k
Rprop	2.4k
SparseAdam	1.8k
ASGD	1.5K

Everything except Adafactor, LARS, and LAMB are in Core, so they have a "boost" in stats, but Adafactor numbers are relatively low. That said, there does seem to be a slight growth of papers with Adafactor based on paperswithcode's experimental linegraph:
.

Knowing all this, I would like to answer the following:

1. is Adafactor the candidate we are going for? Do other alternatives (e.g., SGD with clipping) fail to compare? Are there other optimizers in this realm that I am not aware of?
2. If we want to support Adafactor, should we be adding an entirely new Optimizer(torch.optim.Optimizer)? Should it live in core vs pytorch-labs? Maybe down the road, we'd be able to bundle a collection of foreach ops with some global state and bypass the need for big O Optimizer entirely.

[msaroufim]

Whether adafactor is in labs or core doesn't matter too much to me personally, I can add it in labs until we have more convincing first hand evidence it works better

The clipping paper is quite new with 1 citation so not sure how many people have independently reproduced the results

For Adafactor specifically the argument I'm making is not a popularity-based one for its inclusion, it seems to be disproportionately used at startups that know how to train large models

EDIT:
Re-the point on foreach ops, at a high level that sounds good I guess you can save memory both by defining your optimizer state but also how you apply it, those feel like complementary techniques to me?

[janeyx99]

I think adding an Optimizer subclass in pytorch-labs is a great first step--and I'm happy to help review that.

The function of foreach ops is definitely a quarter-baked complementary idea that requires much more design so it is fine to not worry about that for this issue. What I was thinking was what @cpuhrsch has mentioned with the new post_accumulate_grad_hook taking in a function--we could define the optimizer step as a function with global state through that without the need of subclassing Optimizer. Unsure if this would be lighter weight, but I do think the standard today is still to subclass Optimizer.

[awgu]

I want to understand the relation between FSDP and Adafactor, so I am sharing my initial thoughts. Please point out any errors!

TL;DR

• Existing FSDP cannot compose with Adafactor, forcing a choice between FSDP + Adam vs. DDP vs. Adafactor.
• For >= 4 GPUs, FSDP + Adam uses less memory for Llama-7B. A similar rule can be derived for other models.
• There is a possible a path for per-parameter-sharding FSDP to compose with Adafactor.

cc: @gchanan

Can Adafactor compose with FSDP?

Under existing PyTorch FSDP's flat-parameter sharding, it is not possible without some shape recovery algorithm. Under per-parameter sharding, it requires one all-reduce per nn.Parameter.

• Flat-parameter sharding flattens several parameters together, losing the required tensor structure. Existing Fairscale/PyTorch FSDP uses this.
• Per-parameter-sharding shards each parameter individually on dim-0.
• Adafactor requires computing a mean over the row state (t5x, fairseq), which, under per-parameter sharding, would be sharded across data parallel workers. Thus, this mean computation would require an all-reduce.

Should I use Adafactor or Today's PyTorch FSDP?

Assuming that FSDP and Adafactor cannot compose, we consider two options: FSDP + Adam vs. DDP + Adafactor.

Setup
Notation:

• N: number of GPUs
• M: model's total numel
• A: Adafactor states numel (function of parameter shapes)
• B: transformer block numel
• E: non-transformer-block numel (M - L * B for L transformer blocks)

For Llama-7B, M = 6.7e9, A = 2.8e6, B = 2.0e8, and E = 2.6e8.

Analysis
The two options have the same activation memory (since we assume same local batch size), so we only need to consider the model-state memory: parameters, gradients, and optimizer states.

	Parameters/Gradients	Optimizer States
FSDP + Adam	2M/N + (3 + 1/N)B + E	2M/N
DDP + Adafactor	2M	A

For simplicity, the table assumes no mixed precision. Mixed precision would only decrease the contribution from the (3+1/N)B + E term in FSDP + Adam, favoring FSDP. For brevity, I omit the explanation for (3 + 1/N)B + E; let me know if you want me to write it out here.

Since M, A, B, and E are known (for Llama-7B), we solve for N in 4M/N + (3 + 1/N)B + E = 2M + A to get N ~= 2.15. This suggests that as long as you have >2 GPUs, you would prefer to use FSDP + Adam.

Practically, there are two caveats:

1. Existing FSDP uses (4 + 1/N)B due to an implementation limitation, though this only changes N to 2.18, not affecting the conclusion.
2. Existing FSDP uses recordStream, which can hold onto memory longer than necessary depending on GPU kernel timing, increasing the (3 + 1/N)B term. This makes it difficult to upper bound its memory usage analytically.
   • That being said, we generally do not see FSDP holding onto >3 layers' memory. Relaxing (3 + 1/N)B to (32 + 1/N)B, conservatively representing FSDP holding onto 8 layers' memory, N still solves to <4.

Conclusion: For >= 4 GPUs, we prefer to use FSDP + Adam than DDP + Adafactor.

What if FSDP Could Compose with Adafactor?

We would achieve a 2M/N -> A reduction in model-state memory, which is significant for small N. How might we achieve this?

We are proposing an FSDP rewrite that uses per-parameter sharding (#114299) that includes avoiding both practical caveats above. The proposal includes exposing FSDP-sharded parameters as DTensors with (Shard(0),) placement, providing an interface for non-pointwise optimizer integration. In theory, DTensor could implement torch.mean (if not already) to include the required all-reduce without user intervention (cc: @wanchaol).

A critical consideration is optimizer state dict. For Adafactor optimizer states, the column state would be replicated, and the row state would be sharded on dim-0. This can fit nicely into the DTensor paradigm. If DTensor and the Adafactor implementation could work such that the row state becomes (Shard(0),) placement and the column state becomes (Replicate,) placement without intervention, then FSDP and non-FSDP usages could share the same implementation (again, cc: @wanchaol).

From a performance standpoint, overlapping the all-reduce with computation can be desirable. However, the importance depends on the hardware setup, especially if FSDP + Adafactor is targeting a small number of GPUs within one host with fast network bandwidth. One option would be to leave overlapping to torch.compile, but we should brainstorm around this overlapping requirement early.

[janeyx99]

@awgu thank you for this analysis. If we look strictly at memory, it looks like FSDP + Adafactor would be the clear winner in the future (2M/N + (3 + 1/N)B + E + A). So I'm all for the FSDP rewrite along with your suggestions regarding DTensor mentioned above (I don't quite understand the (Shard(0), ) row vs column state in DTensor, but I figure @wanchaol gets it) as maintaining Adafactor being the same across nonFSDP + FSDP is crucial for composability.

To pick your brain a little bit, since this calculation was on llama7b, do you see this changing a lot for llama70b/across other LLMs?

[awgu]

To pick your brain a little bit, since this calculation was on llama7b, do you see this changing a lot for llama70b/across other LLMs?

It depends. It might be easiest to re-run the calculations for the specific LLM and cluster.

Generally, increasing M correlates with increasing N to meet the memory requirement. However, for this memory-constrained context, we probably would start applying other memory saving techniques like (Q)LoRA, which changes the formulas. If we use LoRA, then there is lower optimizer state memory, which means that Adafactor gives less absolute savings.

[wanchaol]

Thanks for the detailed summary @awgu!

If DTensor and the Adafactor implementation could work such that the row state becomes (Shard(0),) placement and the column state becomes (Replicate,) placement without intervention, then FSDP and non-FSDP usages could share the same implementation

Yeah that's a natural fit to DTensor where we can control the optimizer state to make one be replication and the other be sharded! There possibly be more non-pointwise optimizer states needs this flexibility going forward, so this looks like a important starter for us!

maintaining Adafactor being the same across nonFSDP + FSDP is crucial for composability

I looked at the code pointer for row/col states, I feel if torch.zeros would take device_mesh and placements and when it does it give back a DTensor, this might makes it easier to maintain this uniformly. right now dtensor need to separately maintain those tensor constructors, cc @ezyang as we discussed this a while ago. The other option is to make a TorchDispatchMode that allocates those states, and enable this torch_dispatch mode under distributed context, this way we can also maintain same code.

[ezyang]

Given that jax.Array has gone to being distributed by default, I am conceptually in favor of the idea that torch.Tensor should also be distributed by default. However, this means that DTensor must ensure that a non-distributed DTensor is actually just a plain Tensor, and that the APIs are substitutable in these cases. I don't think this is impossible to do, but it is Work and would need to be done. What I do not want is having both Tensor and single node DTensor floating around, a distinction without meaning.