Repository: pytorch
Issue: A diagnostics mode to report when constraints aren't being met for optimal performance
Labels: oncall: profiler

### 🚀 The feature, motivation and pitch

As its known CUDA has all kinds of constraints for certain optimizations to kick in. Often it's not obvious when such constraints aren't met and a user ends up with pretty bad performance and they are not the wiser.

It'd be great if pytorch could tell users when they could do better.

---------------

Let's start with a use-case of **a very not obvious** problem. 

I've done multiple HF Trainer benchmarks for [A100](https://github.com/huggingface/transformers/issues/15026) and [RTX-3090](https://github.com/huggingface/transformers/issues/14608) with several models. One of these was to compare fp16 vs bf16 vs tf32 vs fp32.
All was reporting well, except this setup: [t5-large on A100](https://github.com/huggingface/transformers/issues/15026#issuecomment-1004543189) for some reason was not activating bf16 optimizations and running at fp32 speed. 

This is with BS=8. Observe:

| Variation       |   Train<br>samples<br>per<br>second |   Diff<br>% |   Train<br>loss |
|:----------------|------------------------------------:|------------:|----------------:|
| --tf32 0        |                               31.59 |         0 |            2.03 |
| --tf32 1        |                               36.13 |         14 |            2.03 |
| --bf16 --tf32 0 |                               31.35 |          -1 |            2.04 |
| --bf16 --tf32 1 |                               31.30 |          -1 |            2.04 |

bf16 is slower than/similar to fp32 here, which is not what we want.

Here is the same benchmark on RTX-3090

| Variation       |   Train<br>samples<br>per<br>second |   Diff<br>% |   Train<br>loss |
|:----------------|------------------------------------:|------------:|----------------:|
| --tf32 0        |                               31.88 |         0 |            2.03 |
| --tf32 1        |                               35.66 |         12 |            2.03 |
| --bf16 --tf32 0 |                               35.07 |         10 |            2.04 |
| --bf16 --tf32 1 |                               35.13 |         10 |            2.04 |

no problem here, both bf16 and tf32 are faster than fp32.

@eqy on pytorch slack suggested that the batch size was too small. (note the same bs=8 on both cards). 

So I go for the max BS before OOM, which is bs=16, and we get on A100:

| Variation       |   Train<br>samples<br>per<br>second |   Diff<br>% |   Train<br>loss |
|:----------------|------------------------------------:|------------:|----------------:|
| --tf32 0        |                               39.59 |           0 |            2.10 |
| --tf32 1        |                               67.24 |          70 |            2.10 |
| --bf16 --tf32 0 |                               61.37 |          55 |            2.12 |
| --bf16 --tf32 1 |                               59.95 |          51 |            2.12 |

That's better, but still far from optimal. Only after adding grad accumulation to make an effective BS=64, I was able to get excellent outcome wrt super to fp32 performance:

| Variation       |   Train<br>samples<br>per<br>second |   Diff<br>% |   Train<br>loss |
|:----------------|------------------------------------:|------------:|----------------:|
| --tf32 0        |                               45.60 |           0 |            2.35 |
| --tf32 1        |                               86.78 |          90 |            2.36 |
| --bf16 --tf32 0 |                               75.85 |          66 |            2.41 |
| --bf16 --tf32 1 |                               73.19 |          61 |            2.41 |

(all the details to reproduce the issue are in the benchmark posts: [A100](https://github.com/huggingface/transformers/issues/15026) and [RTX-3090](https://github.com/huggingface/transformers/issues/14608). And I omitted fp16 results as they overflow on t5-large).

Now how could I possibly know that BS=8 was absolutely bad and BS=16 was still so so and that I needed to squeeze BS=64? 

Is it the case where 16x constraint not being met, besides everything else but BS being a multiple of 16? As clearly BS=16 was already much much better than BS=8 relative to fp32.

Also note that 4 other models I experimented with didn't have this problem, so a user won't expect a sudden degradation in performance when changing a model if they see that bf16 is normally faster than fp32, which we can clearly see it's not always the case.

--------------------

If for example I'm using Megatron-LM it tells the user:
```
    custom_kernel_constraint = seq_len > 16 and seq_len <=2048 and \
        seq_len % 4 == 0 and attn_batch_size % 4 == 0
    # Print a warning.
    if not ((args.fp16 or args.bf16) and
            custom_kernel_constraint and
            args.masked_softmax_fusion):
        if args.rank == 0:
            print('WARNING: constraints for invoking optimized'
                  ' fused softmax kernel are not met. We default'
                  ' back to unfused kernel invocations.', flush=True)
```
if its custom fused kernels aren't going to get activated.

Could we have something similar for CUDA optimizations?

So what I am asking for is a similar feature for CUDA. So if say I run:

```
REPORT_UNMET_CONSTRAINTS=1 python pytorch-program.py 
```

In the case of BS=8 or BS=16 above it'll tell me that some constraints haven't been met and ideally which, so that the user can act upon it.

Thank you.



### Alternatives

_No response_

### Additional context

_No response_

cc @ilia-cher @robieta @chaekit @gdankel @bitfort @ngimel @orionr @nbcsm @guotuofeng @guyang3532 @gaoteng-git