Extend Fake Tensor Caching to Symints

[eellison]

Extend Fake Tensor Caching to Symints

[ezyang]

You have to run with --dynamic-shapes --dynamic-batch-size-only to actually trigger dynamic shapes in the benchmark suite

[aorenste]

--dynamic-shapes --dynamic-batch-size-only

I had both (scroll to the right of the given command line in the comment).

[eellison]

When I run:

python benchmarks/dynamo/huggingface.py --performance --training --amp --backend aot_eager --device cuda --only BertForQuestionAnswering 
 --print-compilation-time --dynamic-batch-only

And add

    def __del__(self):
        print(self.cache_info())

to FakeTensorMode I see:

DispatchCacheInfo(hits=19258, misses=189, bypasses={'symbolic shape': 22386, 'dynamic output shape': 1, 'CompositeImplicitAutograd': 697, 'non-fake tensor': 54, 'non-FakeTensor output': 51}, size=189)

I also see that in the above benchmark without dynamic-batch-only disabling fake tensor cache causes 5 seconds slowdown. It's possible you're only looking at sym_types inputs but not fake tensor inputs with symints. About half of the ops are bypassed due to symints so I would expect a couple seconds of improvement.

[aorenste]

So it turns out that the perf measurements work a lot better when you store them with += instead of =. Once I do that the symint stuff pops out as quite a bit more expensive (significant percentage of the dispatch time for some of the benchmarks)

[aorenste]

I'm not sure if this would be worth it. I instrumented the existing cache and ran:

python benchmarks/dynamo/torchbench.py --performance --inference --amp --backend inductor --disable-cudagraphs --device cuda

From what I can tell the time spent in the FakeTensorMode dispatch of ops that were not cached because they contained any kind of sym expr was tiny. Of 151s spent in dispatch 0.003s were cache bypassed due to containing a SymInt, SymFloat, or SymBool.

Is there a better benchmark I should use to measure the potential of this change?

[aorenste]

You have to run with --dynamic-shapes --dynamic-batch-size-only to actually trigger dynamic shapes in the benchmark suite

I used --dynamic-batch-only because --dynamic-batch-size-only wasn't accepted.

Running:

python benchmarks/dynamo/torchbench.py --performance --inference --amp --backend inductor --disable-cudagraphs --device cuda --dynamic-shapes --dynamic-batch-only

Gives a total time of 338s (so slower) but still only 0.045s in dispatching w/ Sym types. Still fairly insignificant. Or I'm measuring it poorly.

[ezyang]

ok, well, I never claimed that you would expect a speedup here :)

[eellison]

The benchmark runs with --dynamic-shapes --dynamic-batch-only, not just --dynamic-batch-only. Maybe you need both ?