fix for high gpu reserved memory

[ngoyal2707]

Before this change:

Training 175b on 1024 A100s:

"ppl": "4217.39", "wps": "105123", "ups": "0.05", "wpb": "2.09715e+06", "bsz": "1024", "num_updates": "26", "lr": "1.09091e-06", "gnorm": "40.827", "clip": "100", "train_wall": "20", "max_cuda_gb_allocated": "13.4", "max_cuda_gb_reserved": "75.4", "current_cuda_gb_allocated": "1.9", "current_cuda_gb_reserved": "75.4", "cuda_gb_free": "65.9"

After the change:

"ppl": "4217.39", "wps": "107317", "ups": "0.05", "wpb": "2.09715e+06", "bsz": "1024", "num_updates": "26", "lr": "1.09091e-06", "gnorm": "40.827", "clip": "100", "train_wall": "20", "max_cuda_gb_allocated": "13.8", "max_cuda_gb_reserved": "22.1", "current_cuda_gb_allocated": "1.9", "current_cuda_gb_reserved": "22.1", "cuda_gb_free": "65.5"

No idea if the approach is perfect or not, but for above use case, I am getting slightly better WPS (this most likely could be
a flake though, and I'd expect them to match), much better cuda gb reserved (which I can almost explain now, where its being used) and ppl matches.

plus I can train much much larger models now. (will post additional details and results in post).

Will do additional testing

[anj-s]

Thanks Naman! Looks great and we should definitely incorporate. Couple of questions:

1. This still uses the prefetch mechanism right? Can you explain the approach a bit.
2. Curious why the cuda_gb_free does not increase even when the reserved value decreases so drastically? Maybe we are logging something else but I would expect the value of free to be higher.

[anj-s]

curious about this change? Why do we condition on the boolean?

[ngoyal2707]

1. Yes it is using prefetch approach, there could maybe a way to achieve above without prefetch, but I thought hard about it and couldn't see it

2. that is just by how we log it: https://github.com/fairinternal/fairseq-py/blob/main/fairseq/trainer.py#L1060

[min-xu-ai]

cc @zhaojuanmao

[zhaojuanmao]

@ngoyal2707 are you disabling multiple process groups for reduce_scattter and all_gather in the backward, instead, using the backward prefetch approach here?

PyTorch version of FSDP is doing the backward prefetch right now. This will better avoid over prefetching more layers in the backward pass using multiple process group approach

[zhaojuanmao]

but will it be conflict with multi process group approach? will the backward prefetch be a config?

[stephenroller]

Tried this on AWS. Find it doubles roughly WPS for my particular workload. (1700 -> 3200)

[zhaojuanmao]

@ngoyal2707 I guess this helped releasing GPU reserved memory...

For @stephenroller models, were your training slowed down by high reserved memory and slow CUDA cache allocator? if so, maybe this change helped a lot...

[zhaojuanmao]

but I'm wondering whether this will result in regression for other models which are not suffered from high reserved memory and slow CUDA cache allocator issue

[stephenroller]

Yeah my jobs are fine tuned of a large pre trained model. We consistently saw on both 40gb and 80gb nodes that the memory reserved was always 99%

[zhaojuanmao]

nice to add this forward prefetching!

[stephenroller]

Let's push this through

[anj-s]

Let's push this through

To push this through we need an owner who can:

1. Evaluate if we should move to the prefetch solution or if it will work with the existing 2 PG solution.
2. Fix unit tests and make sure it works for FSDP in general and not just this use case.

Let me know if anyone is up for taking this on!

[zhaojuanmao]

@anj-s and @stephenroller, I suspected the speedup is possibly mostly from this single sentence change "torch.cuda.current_stream().synchronize()" for @stephenroller's use case, possibly it is good to verify that. e..g, @stephenroller, just use FairScale master branch + this change "torch.cuda.current_stream().synchronize()" to your use case, and see whether there is a speed up.

But this change is better to be optional, not be default, otherwise other FSDP use cases that are not suffered from high GPU reserved memory issue possibly will get regression.

[anj-s]

@anj-s and @stephenroller, I suspected the speedup is possibly mostly from this single sentence change "torch.cuda.current_stream().synchronize()" for @stephenroller's use case, possibly it is good to verify that. e..g, @stephenroller, just use FairScale master branch + this change "torch.cuda.current_stream().synchronize()" to your use case, and see whether there is a speed up.

But this change is better to be optional, not be default, otherwise other FSDP use cases that are not suffered from high GPU reserved memory issue possibly will get regression.

Can you explain a little bit about why you think there might be a regression?

[stephenroller]

I'm not entirely sure about this. We also have large scale, not latency bound pretraining, and this change increased WPS there--when ideal reserved memory is only 25% of the node. What's clear in both cases was we were thrashing the allocator aggressively.

This could maybe be a regression for places where the full state and activations fit in memory, and the network is able to fetch faster than the compute can utilize the state. But that case is better addressed with zero2 anyway, no?

[stephenroller]

@ngoyal2707 is clearly the owner and should push this through.

[zhaojuanmao]

"What's clear in both cases was we were thrashing the allocator aggressively."

==> that mostly means "torch.cuda.current_stream().synchronize()" helps here, because CUDA allocation is a async operation, it could try to allocate a large storage before completely free full parameters here. So it is trashing. Adding this sync for your use case basically makes sure full parameters are freed before allocating. This helped reducing CUDA allocator trash.

For the use cases that are not suffering from slow cudaAllocator, this blocking operation will wait for all pending async operations and reduce parallelizations, and thus could result in regression, no?

[zhaojuanmao]

So for unblocking this effort, suggestion here is:

1. root cause whether prefetching or cuda sync really helps here first
2. if add cuda sync, make it optional

[ngoyal2707]

So I agree with everyone here. I should be the owner but honestly I still don't know if this is the best approach here and I want opinion of people who understands cuda and pytorch caching allocator better

The core reason of the issue is due to unbounded queue of launching kernels combined with multiple cuda streams.

One very simple pytorch example of this is following:

import torch
torch.cuda._sleep(100000000000)
for i in range(100):

    x=torch.zeros(4 * 1024 * 1024, device="cuda")
    del x
print(torch.cuda.memory_allocated())
print(torch.cuda.memory_reserved())
# print(torch.cuda.memory_snapshot())

above code outputs at end:

0
16777216

That means, we only reserve 16MB of memory for the 4M fp32 tensor cause its only been used by single stream, so pytorch caching allocator can just assign every time the same memory block as operations within stream are sequential.

comparing above to following:

import torch
s1 = torch.cuda.Stream()

torch.cuda._sleep(100000000000)

with torch.cuda.stream(s1):
     torch.cuda._sleep(100000000000)

for i in range(100):
     x=torch.zeros(4 * 1024 * 1024, device="cuda")
     x.record_stream(s1)
     torch.cuda.current_stream().wait_stream(s1)
     del x

print(torch.cuda.memory_allocated())
print(torch.cuda.memory_reserved())
# print(torch.cuda.memory_snapshot())

This outputs:

0
1677721600

i.e. pytorch reserves 16MB tensor for every new allocation cause each memory block is being used by both current stream and s1 stream.

So in FSDP case, s1 is the all_gather stream, which is of course running slower than CPU thread. What I don't know is, if there is any better solution than just making CPU thread wait. We can change how it waits i.e. instead of just waiting on torch.cuda.current_stream().synchronize(), we can change to always launch new all_gather once previous one is done.

So I think there possibly can be some regression, specially to very small models cause GPU can be idle while CPU thread is launching new kernels.

[qingliu1111]

Want to add a data point about the performance improvement on AWS brought by this PR : )
I'm running fine-tuning flow for 175B model on AWS with 16 nodes (128GPU).

• Without this PR: wps is ~ 2817.1
• With this PR: wps is ~ 8817.1 (almost tripled the speed!)

More context: With similar training setup we were getting wps ~12000 on Azure. For 175B model fine-tuning, this PR helps us close a lot the speed gap between Azure and AWS.

[zhaojuanmao]

adding @mrshenli who is looking into pluggable cudaCacheAllocator, and hopefully can resolve this issue in general