Open mpatel31415 opened 4 months ago
mruberry
commented
4 months ago triage review:
mpatel31415
commented
4 months ago When running:
python /opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_litgpt.py \
--model_name stablecode-completion-alpha-3b \
--compile "thunder_inductor"I still get OOM error for tag pjnl-20240427.
Memory usage:
IvanYashchuk
commented
4 months ago The new command to reproduce the problem is
python /opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_litgpt.py --model_name stablecode-completion-alpha-3b --compile "thunder_inductor_cat"because "inductor" was renamed to "inductor_cat" in https://github.com/Lightning-AI/lightning-thunder/pull/140.
IvanYashchuk
commented
4 months ago Thunder uses a lot more memory than both torch.compile or eager. That's something that should be fixed.
The forward trace has the following weirdness:
[t137, t142, t146, t158, t228, t232, t244] = nvFusion2(t129, t152, t155, t238, t241, t4, t90)
# t130 = prims.convert_element_type(t129, dtypes.float32) # t130: "cuda:0 f32[1, 16384, 2560]"
# t131 = prims.convert_element_type(t90, dtypes.float32) # t131: "cuda:0 f32[1, 16384, 2560]"
# t132 = prims.add(t130, t131) # t132: "cuda:0 f32[1, 16384, 2560]"
# t135 = prims.convert_element_type(t4, dtypes.float32) # t135: "cuda:0 f32[1, 16384, 2560]"
# t136 = prims.add(t132, t135) # t136: "cuda:0 f32[1, 16384, 2560]"
# t137 = prims.convert_element_type(t136, dtypes.bfloat16) # t137: "cuda:0 bf16[1, 16384, 2560]"
# (t141, t142) = prims.var_mean(t136, (2,), correction=0) ## <-- What's going on here?
# (t227, t228) = prims.var_mean(t136, (2,), correction=0) ## <-- What's going on here?
...two var_mean calls for the same input with both t142 and t228 saved to outputs and this pattern is repeated in every nvFusion region. And it's the same both with and without the "inductor_cat" executor.
@riccardofelluga, could you please investigate what's wrong with Thunder on this model and fix the problem?
I suggest starting with PyTorch memory profiler to spot the differences in memory allocations between PyTorch eager and Thunder.
The second var_mean in the nvFuser regions should have been eliminated with CSE pass, probably it doesn't work because comparison and hashing of BoundSymbols with keyword arguments is not implemented, see https://github.com/Lightning-AI/lightning-thunder/issues/397. Where does the second var_mean come from at all?
riccardofelluga
commented
4 months ago Quick update, we probably have multiple issues with memory, however I think with this issue the focus should be on the peak memory usage. For comparison here is the memory profiles for eager:
and here the one for thunder.jit:
They both are pretty similar however there are a couple notable differences.
In the Thunder trace the data from the log_softmax_backward is kept around till the end, whereas eager manages to deallocate it almost as soon as possible. This might be an issue, however it's not contributing too much to the peak of allocation.
The second notable difference is the presence of the nvFuser outputs, that for obvious reasons are not present in eager.
A thing that I've noticed by looking at the traces is that, when we have an nvFuser region next to a torch.compile one, the intermediates stick around for long and they are all pretty big in memory. This is accentuated when the intermediates are passed from forward to backward passes like t123 and t128 both of which are consumed by torch.compile in the backward pass.
Will update with further details.
riccardofelluga
commented
3 months ago This issue seems to be another facet of #256 and #446.
After further investigation it seems that this extra memory usage comes also from splitting torch.nn.functional.gelu between the TorchCompile and nvFuser executors. In particular, as of now, prims.erf and prims.div are fused by nvFuser and the intermediate tensors that are passed from this fusion to the next TorchCompile fusion are pretty big on the order of 600MB for the 1-layer example I showed in the previous comment.
A solution to lower the memory usage is to let TorchCompile fuse those operators by adding them in https://github.com/Lightning-AI/lightning-thunder/blob/82185e3a55d5b3f0bea8a7366d74a275dbe34acd/thunder/executors/torch_compile.py#L209-L222
With this change it is possible to lower the memory usage of the example from 9.548GB to 9.129GB. However such a change is not enough to avoid OOM in the reported issue.
@IvanYashchuk how do you feel about the compromise of adding those operators to the TorchCompile? It is true that it does not solve this issue but there are some memory savings, even tho at the expenses of a little of performance. (and we would need to do more perf testing with the benchmarks to see that there are no other penalties)
Regarding the nvFuser issue highlighted in a previous comment, I am currently looking into it and it's tracked by #397. I am checking if that might be the issue here so I will update this comment as soon as I have more information.
Just to be clear I am not ruling that out, but I think this intermediate tensors I discovered here are also partially to blame.
riccardofelluga
commented
3 months ago Update on the CSE issue, unfortunately it didn't help with memory usage :(
IvanYashchuk
commented
3 months ago @IvanYashchuk how do you feel about the compromise of adding those operators to the TorchCompile?
Another alternative is to do the opposite: restricting the number of supported operators in this executor forcing TorchCompile to be used only for forward and nvFuser for backward (https://github.com/Lightning-AI/lightning-thunder/issues/256#issuecomment-2136096575). Both of these choices should be benchmarked.
riccardofelluga
commented
2 months ago Update on this issue, as of today thunder runs stablecode-completion-alpha-3b with compile option thunder_inductor successfully. However, with thunder_inductor_cat it OOMs.
Stats from python thunder/benchmarks/benchmark_litgpt.py --model_name stablecode-completion-alpha-3b --compile thunder_inductor
Model name: stablecode-completion-alpha-3b
Seq Length: 16384
Micro BS: 1
Global BS: 1
Number of Layers: 32
Number of parameters: 2.77B
Distributed Mode: none
Compiler: thunder_inductor
Average iter time: 1293.71 ms
Memory used: 71.27 GB
Tokens/s: 12663.37
Tokens/s/GPU: 12663.37
TFLOP/s: 404.73Average iteration time 1293.71ms is on H100 80GB HBM3
IvanYashchuk
commented
3 weeks ago One config flag different from the Llama models that affects the memory usage here is parallel_residual (for Llama it's False for StableCode it's set to True). When it's True and the "inductor_cat" executor is used one more tensor is saved for backward for each transformer block. The size of this tensor is (batch_size, block_size, 4 * n_embd) and dtype is fp32, so it's about 670MB more memory per transformer block with batch_size=1. When the "inductor_cat" executor is not used this intermediate is not saved to global memory but recomputed due to the rematerialization pass.
One fix that resolves the problem by constraining TorchCompile region to fuse only when two or more cat operations are present in the candidate fusion group:
diff --git a/thunder/executors/torch_compile.py b/thunder/executors/torch_compile.py
index 204bc204..db268b7d 100644
--- a/thunder/executors/torch_compile.py
+++ b/thunder/executors/torch_compile.py
@@ -166,7 +166,7 @@ class TorchCompileExecutor(FusionExecutor):
continue
# TODO: this could use `get_fuel()` like nvfuserex does
- if self.required_ops is None or any(bsym.sym.id in self.required_ops for bsym in bsyms):
+ if self.required_ops is None or len(tuple(bsym for bsym in bsyms if bsym.sym.id in self.required_ops)) > 1:
region = Region(producers, consumers, bsyms)
fusion_bsym: BoundSymbol = self.fuse(region, fusion_counter)
fusion_counter += 1There's a cat op that appears in backward of split and with the patch above nvFuser is handling that.
Another fix is:
diff --git a/thunder/executors/torch_compile.py b/thunder/executors/torch_compile.py
index 204bc204..8f41ee16 100644
--- a/thunder/executors/torch_compile.py
+++ b/thunder/executors/torch_compile.py
@@ -221,6 +221,8 @@ supported_ops = {
prims.add.id,
prims.broadcast_in_dim.id,
prims.cat.id,
+ prims.div.id,
+ prims.erf.id,
prims.convert_element_type.id,
prims.full.id,
prims.mul.id,with this gelu is fused into the TorchCompile region.
The first solution results in a slightly faster code than the second one on H100 and I checked that the first solution doesn't have perf impact on Llama 3 8B (which uses parallel_residual=False).
It's still not enough for a single GPU run of this model but at least the problem is occurring in the loss function and sharding may help:
File "/usr/local/lib/python3.10/dist-packages/torch/nn/functional.py", line 3479, in cross_entropy
return torch._C._nn.cross_entropy_loss(
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.50 GiB. GPU 0 has a total capacity of 79.10 GiB of which 132.00 MiB is free. Including non-PyTorch memory, this process has 78.96 GiB memory in use. Of the allocated memory 76.84 GiB is allocated by PyTorch, and 1.44 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)After the change, Thunder still saves more tensors for backward compared to Inductor because Inductor doesn't save the result of gelu (input to linear layers) for backward; it's recomputed instead.
Setting nv_enable_bookend=False also helps with memory usage.
IvanYashchuk
commented
3 weeks ago For a microbenchmark, the first solution of restricting TorchCompile to be applied only to regions with more than 1 cat op results in 1.5x-2x regressions in backward.
It's better to allow GELU to be fused into a TorchCompile region together with RoPE to not regress on Llama models.
IvanYashchuk
commented
3 weeks ago https://github.com/Lightning-AI/lightning-thunder/pull/1019 lets running this model with FSDP. More memory savings are needed to run this model on a single GPU.
đ Bug
With newest version of Docker image (tested on 2024-04-22 ) training with thunder.jit with additional inductor executor gives OOM error.
To Reproduce
Before each testing each compilation method I restarted the container:
Thunder inductor
Eager
Thunder
Inductor
Expected behavior
If we can run model in Eager mode, with Thunder + Pytorch default compilation we should be able to use it for Thunder + Inductor as well.
Environment
As in the Docker image. This results come from single H100.
nvidia-smi output: