[Bug] Unexpected performance drop with float8 training + compiling only nn.Linear layers + using selective per op AC

[danielvegamyhre]

Summary

I'm doing some benchmarking with torchtitan on H100s to compare the experimental feature in #778 with other training configurations.

One important comparison is comparing:

• float8nocompile (handwritten triton kernels for fp8 conversion of nn.Linear layers), and
• production float8 training + using torch.compile on only the nn.Linear layers

I've ran this comparison using:

• no AC
• full AC
• selective per op AC

However, specifically when using selective per op AC, there is a massive drop in performance when using the production float8 training + using torch.compile on only the nn.Linear layers, as compared to using torch.compile on the full model.

This does not occur when using no AC or full AC.

I would expect some performance degradation compiling only nn.Linear instead of the full model, but the drop off is massive (see screenshot of benchmarks below). TFLOPS drops from 386.58 down to 125.88!

I looked at the traces for these 2 runs, and found some surprising issues:

1. In the forward pass of the slow configuration (compiling only nn.Linear layers) it seems like aten::cross_entropy_loss is running on CPU and not dispatching any CUDA kernels, causing it to take 267ms vs 71us in the fully compiled version (~3800x slowdown).

Only nn.Linears compiled (267ms):

Full model compiled (71us):

2. In the backward pass of the slow configuration (compiling only nn.Linear layers) there is an extremely long/slow FSDP::post_backward_reduce call that does not appear in the fully compiled version (or rather, it is orders of magnitude faster).

Only nn.Linears compiled:

Steps to reproduce

1. Checkout [Not for land] Integrate float8nocompile, an experimental feature for high performance #778
2. Edit training_configs/llama3_8b.toml to run prod float8 + fully compiled model + selective per op AC on H100s:

NGPU=4 CONFIG_FILE="./train_configs/llama3_8b.toml" ./run_llama_train.sh

4. Edit training_configs/llama3_8b.toml to run prod float8 + only linear layers compiled + selective per op AC on H100s (don't think # of GPUs matters):

TORCHTITAN_COMPILE_LINEAR_ONLY=1 NGPU=4 CONFIG_FILE="./train_configs/llama3_8b.toml" ./run_llama_train.sh

cc @vkuzo @soulitzer

[tianyu-l]

Thanks for flagging the issue.

It looks the memory usage is too high, which might cause throughput degradation. I think it's worth checking

1. the memory snapshots (with https://github.com/pytorch/torchtitan/blob/main/docs/debugging.md)
2. meanwhile, change to a smaller model and see when memory is no longer saturated, if the throughput difference is still this large.

also copying over what @awgu said offline, supporting point 1

looks like it is doing a cuda defragmentation with expandable segments turned on, would be good to check the memory

[gnadathur]

cc: @vkuzo , fyi

[danielvegamyhre]

Thanks for flagging the issue.

It looks the memory usage is too high, which might cause throughput degradation. I think it's worth checking

1. the memory snapshots (with https://github.com/pytorch/torchtitan/blob/main/docs/debugging.md)
2. meanwhile, change to a smaller model and see when memory is no longer saturated, if the throughput difference is still this large.

also copying over what @awgu said offline, supporting point 1

looks like it is doing a cuda defragmentation with expandable segments turned on, would be good to check the memory

@tianyu-l I think we are running into another instance of this issue in #808 (see my comment here #808 (comment)). I went ahead and pulled some memory snapshots and analyzed them here #808 (comment). Still not sure of the root cause but please feel free to add any input.

If I'm correct that this issue is related, then the scope of this bug is bigger than we thought, and selective per op AC has some deeper issues that affects more than just compiling only linear layers, but also fp8 row-wise quantization.

[vkuzo]

At the top of this issue, the peak memory usage in the problematic experiment is 92GB, which is too close to the max memory available on the H100 GPU.

looks like it is doing a cuda defragmentation with expandable segments turned on, would be good to check the memory

^ is likely related to peak memory being close to machine limit

I don't think there is conclusive evidence that #808 and this issue are related, they just both seem to be about selective per-op AC.